gpu_mma_dispatcher.h

Go to the documentation of this file.

#pragma once
 
#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
 
#include "sdfg/codegen/dispatchers/block_dispatcher.h"
#include "sdfg/targets/gpu/gpu_mma_fragment.h"
#include "sdfg/types/type.h"
 
namespace sdfg::gpu {

class GpuMmaDispatcherStrategy {
public:
    using EltwiseCompute = std::function<void(
        codegen::CodegenOutput& out,
        const std::string& main_frag_elem,
        const std::string& idx,
        const std::vector<std::string>& other_frag_elems
    )>;
 
    virtual ~GpuMmaDispatcherStrategy() = default;

    virtual void emit_needed_declarations(codegen::CodegenOutput& out) const {
    }

    virtual void
    emit_fill(codegen::CodegenOutput& out, const std::string& acc_frag, types::PrimitiveType fill_type) const = 0;

    virtual void emit_mma(
        codegen::CodegenOutput& out,
        const std::string& frag_d,
        const std::string& frag_a,
        const std::string& frag_b,
        const std::string& frag_c
    ) const = 0;

    virtual void emit_fragment_load(
        codegen::CodegenOutput& out,
        const std::string& frag,
        const std::string& base,
        const GpuMmaFromMemoryLayout& layout,
        MmaFragmentType frag_type,
        types::PrimitiveType element_type,
        MmaFragmentLayout fragment_layout
    ) const = 0;

    virtual void emit_fragment_store(
        codegen::CodegenOutput& out,
        const std::string& frag,
        const std::string& base,
        const GpuMmaFromMemoryLayout& layout,
        MmaFragmentType frag_type,
        types::PrimitiveType element_type
    ) const = 0;

    virtual void emit_eltwise(
        codegen::CodegenOutput& out,
        const std::string& main_frag,
        const std::vector<std::string>& additional_frags,
        const EltwiseCompute& compute
    ) const = 0;

    virtual std::string fragment_element(const std::string& frag, const std::string& idx) const {
        throw std::runtime_error("GpuMmaDispatcherStrategy: clipped fragment transfers are not supported");
    }

    virtual std::string fragment_element_count(const std::string& frag) const {
        throw std::runtime_error("GpuMmaDispatcherStrategy: clipped fragment transfers are not supported");
    }

    virtual std::optional<std::pair<std::string, std::string>> fragment_element_coords(
        MmaFragmentType frag_type, const MmaBlockSize& block_size, const std::string& idx, std::string& condition
    ) const {
        return std::nullopt;
    }
};

void emit_clipped_fragment_transfer(
    codegen::CodegenOutput& out,
    const GpuMmaDispatcherStrategy& strategy,
    bool load,
    const std::string& frag,
    const std::string& base,
    const GpuMmaFromMemoryLayout& layout,
    const MmaBlockSize& block_size,
    MmaFragmentType frag_type,
    types::PrimitiveType element_type,
    MmaFragmentLayout fragment_layout
);

class GpuMmaNodeDispatcher : public codegen::LibraryNodeDispatcher {
protected:
    std::unique_ptr<GpuMmaDispatcherStrategy> strategy_;
 
public:
    GpuMmaNodeDispatcher(
        codegen::LanguageExtension& language_extension,
        const Function& function,
        const data_flow::DataFlowGraph& data_flow_graph,
        const data_flow::LibraryNode& node,
        std::unique_ptr<GpuMmaDispatcherStrategy> strategy
    )
        : LibraryNodeDispatcher(language_extension, function, data_flow_graph, node), strategy_(std::move(strategy)) {
    }
};

class GpuMmaMatmulDispatcher : public GpuMmaNodeDispatcher {
public:
    using GpuMmaNodeDispatcher::GpuMmaNodeDispatcher;
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};

class GpuMmaFillDispatcher : public GpuMmaNodeDispatcher {
public:
    using GpuMmaNodeDispatcher::GpuMmaNodeDispatcher;
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};

class GpuMmaFragmentLoadDispatcher : public GpuMmaNodeDispatcher {
public:
    using GpuMmaNodeDispatcher::GpuMmaNodeDispatcher;
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};

class GpuMmaFragmentStoreDispatcher : public GpuMmaNodeDispatcher {
public:
    using GpuMmaNodeDispatcher::GpuMmaNodeDispatcher;
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};

class GpuMmaEltwiseAddDispatcher : public GpuMmaNodeDispatcher {
public:
    using GpuMmaNodeDispatcher::GpuMmaNodeDispatcher;
 
    void dispatch_code_with_edges(
        codegen::CodegenOutput& out,
        std::vector<codegen::DispatchInput>& inputs,
        std::vector<codegen::DispatchOutput>& outputs
    ) override;
};
 
} // namespace sdfg::gpu