gpu_offload_base_dispatcher.h

Go to the documentation of this file.

#pragma once
 
#include <string>
#include <vector>
 
#include "sdfg/analysis/analysis.h"
#include "sdfg/codegen/code_snippet_factory.h"
#include "sdfg/codegen/dispatchers/node_dispatcher.h"
#include "sdfg/codegen/instrumentation/instrumentation_info.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/targets/gpu/gpu_offload_schedule_type.h"
#include "sdfg/types/type.h"
 
namespace sdfg {
namespace gpu {
 
// Interface to handle everything cuda/rocm specific
class GPUOffloadDispatcherStrategy {
public:
    virtual ~GPUOffloadDispatcherStrategy() = default;
 
    virtual bool is_matching_kernel_language_extension(
        const StructuredSDFG& sdfg, const StructuredLoop& loop, codegen::LanguageExtension& language_extension
    ) = 0;
    virtual std::unique_ptr<codegen::LanguageExtension>
    create_kernel_language_extension(StructuredSDFG& sdfg, const StructuredLoop& loop) = 0;
    void require_matching_kernel_language_extension(
        const StructuredSDFG& sdfg, const StructuredLoop& loop, codegen::LanguageExtension& language_extension
    ) {
        if (!is_matching_kernel_language_extension(sdfg, loop, language_extension)) {
            throw std::runtime_error("Outermost loop should have already provided kernel language extension");
        }
    }

    virtual void emit_target_header_declarations(codegen::PrettyPrinter& kernel_header_stream) = 0;
 
    virtual void emit_target_global_includes(codegen::CodeSnippetFactory& library_snippet_factory) = 0;
 
    virtual void dispatch_kernel_call(
        codegen::PrettyPrinter& main_stream,
        const std::string& kernel_name,
        codegen::LanguageExtension& host_language_ext,
        symbolic::Expression& num_blocks_x,
        symbolic::Expression& num_blocks_y,
        symbolic::Expression& num_blocks_z,
        symbolic::Expression& block_size_x,
        symbolic::Expression& block_size_y,
        symbolic::Expression& block_size_z,
        std::vector<std::string>& arguments_device
    ) = 0;
 
    virtual codegen::TargetType get_instrumentation_kernel_target_type() const = 0;
 
    virtual int get_warp_size() const = 0;
 
    virtual bool is_device_pointer_storage(const types::StorageType& storage) const = 0;

    virtual std::string kernel_file_extension() const = 0;
    virtual std::string kernel_header_file_extension() const = 0;

    virtual std::string warp_shuffle_xor(const std::string& value, const std::string& lane_mask) const = 0;
};

class GPUOffloadBaseDispatcher : public codegen::NodeDispatcher {
protected:
    structured_control_flow::StructuredLoop& node_;
    std::unique_ptr<GPUOffloadDispatcherStrategy> strategy_;
    // Threads per block when statically known (0 otherwise); emitted as __launch_bounds__.
    long launch_bounds_ = 0;
 
    GPUOffloadBaseDispatcher(
        codegen::LanguageExtension& language_extension,
        StructuredSDFG& sdfg,
        analysis::AnalysisManager& analysis_manager,
        structured_control_flow::StructuredLoop& node,
        codegen::InstrumentationPlan& instrumentation_plan,
        codegen::ArgCapturePlan& arg_capture_plan,
        std::unique_ptr<GPUOffloadDispatcherStrategy> strategy
    );
 
    bool is_outermost_map(analysis::AnalysisManager& analysis_manager);
 
    void dispatch_header(
        codegen::PrettyPrinter& globals_stream,
        const std::string& kernel_name,
        std::vector<std::string>& arguments_declaration
    );
 
    void dispatch_kernel_preamble(
        codegen::PrettyPrinter& library_stream,
        const std::string& kernel_name,
        std::vector<std::string>& arguments_declaration
    );
 
    // Emit the includes for every kernel-local library dependency (e.g. rocwmma)
    // tracked by the nested snippet factory, into the kernel's include header.
    void emit_lib_dependency_includes(
        codegen::PrettyPrinter& kernel_header_stream, codegen::NestedCodeSnippetFactory& nested_snippet_factory
    );

    virtual void validate_before_dispatch(analysis::AnalysisManager& analysis_manager) {
    }
 
    virtual void dispatch_kernel_body(
        codegen::NestedCodeSnippetFactory& kernel_snippet_factory,
        codegen::PrettyPrinter& kernel_source_stream,
        codegen::PrettyPrinter& kernel_header_stream,
        codegen::LanguageExtension& kernel_lang_ext,
        symbolic::Symbol indvar,
        std::vector<std::string>& scope_variables,
        symbolic::Expression& num_iterations
    ) = 0;
 
public:
    void dispatch_node(
        codegen::PrettyPrinter& main_stream,
        codegen::PrettyPrinter& globals_stream,
        codegen::CodeSnippetFactory& library_snippet_factory
    ) override;
 
    codegen::InstrumentationInfo instrumentation_info() const override;
};
 
} // namespace gpu
} // namespace sdfg