rocm_mma_bare_dispatcher.h

Go to the documentation of this file.

#pragma once
 
#include "sdfg/targets/gpu/gpu_mma_dispatcher.h"
 
namespace sdfg::gpu::rocm {
 
inline data_flow::ImplementationType ImplementationType_ROCM_MMA_BARE("ROCM_MMA_BARE");

class RocmMmaBareStrategy : public GpuMmaDispatcherStrategy {
public:
    void emit_fill(codegen::CodegenOutput& out, const std::string& acc_frag, types::PrimitiveType fill_type) const override;
 
    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 override;
 
    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 override;
 
    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 override;
 
    void emit_eltwise(
        codegen::CodegenOutput& out,
        const std::string& main_frag,
        const std::vector<std::string>& additional_frags,
        const EltwiseCompute& compute
    ) const override;
 
    std::string fragment_element(const std::string& frag, const std::string& idx) const override {
        return frag + "[" + idx + "]";
    }
 
    std::string fragment_element_count(const std::string& frag) const override {
        return "static_cast<int>(sizeof(" + frag + ") / sizeof(" + frag + "[0]))";
    }

    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 override {
        const std::string lane = "(threadIdx.x % 64)";
        const std::string lo = "(" + lane + " % 32)";
        const std::string kq = "(4 * (" + lane + " / 32) + " + idx + ")";
        condition.clear();
        switch (frag_type) {
            case MmaFragmentType::A:
                return std::make_pair(lo, kq);
            case MmaFragmentType::B:
                return std::make_pair(kq, lo);
            default:
                return std::make_pair("(8 * (" + idx + " / 4) + 4 * (" + lane + " / 32) + " + idx + " % 4)", lo);
        }
    }
 
private:
    void emit_fragment_transfer(
        codegen::CodegenOutput& out,
        bool load,
        MmaFragmentType type,
        const GpuMmaFromMemoryLayout& layout,
        types::PrimitiveType element_type,
        const std::string& frag,
        const std::string& base
    ) const;
};
 
} // namespace sdfg::gpu::rocm