sdfg::gpu::GpuMmaDispatcherStrategy Class Reference

Hardware-dependent emit primitives for the composable MMA library nodes. More...

#include <sdfg/targets/gpu/gpu_mma_dispatcher.h>

Inheritance diagram for sdfg::gpu::GpuMmaDispatcherStrategy:

Public Types
using	EltwiseCompute
	Per-element callback used by emit_eltwise: writes one output element from its inputs.

Public Member Functions
virtual	~GpuMmaDispatcherStrategy ()=default
virtual void	emit_needed_declarations (codegen::CodegenOutput &out) const
	Target library includes / helper declarations needed before the first emit.
virtual void	emit_fill (codegen::CodegenOutput &out, const std::string &acc_frag, types::PrimitiveType fill_type) const =0
	Zero-initialize the accumulator fragment behind acc_frag.
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
	frag_d = frag_a * frag_b + frag_c (single MMA step; all operands are fragments).
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
	Load a tile from memory (base + layout) into the fragment behind frag.
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
	Store the fragment behind frag to memory (base + layout).
virtual void	emit_eltwise (codegen::CodegenOutput &out, const std::string &main_frag, const std::vector< std::string > &additional_frags, const EltwiseCompute &compute) const =0
	Element-wise loop over main_frag, invoking compute for each element.
virtual std::string	fragment_element (const std::string &frag, const std::string &idx) const
	Expression for element idx of this lane's share of frag (needed for clipped transfers).
virtual std::string	fragment_element_count (const std::string &frag) const
	Number of elements of this lane's share of frag (needed for clipped transfers).
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
	Analytic (row, col) block coordinates of this lane's element idx, if the fragment's register layout is known; clipped transfers otherwise recover them by loading an iota table through the fragment load.

Detailed Description

Hardware-dependent emit primitives for the composable MMA library nodes.

A GpuMma*Dispatcher owns one strategy and delegates only the parts that differ between targets (CUDA vs ROCm) and implementations (fragment/WMMA vs MFMA) to it. Everything reusable - reading the nodes, wiring up operands - stays in the dispatchers.

Member Typedef Documentation

EltwiseCompute

using sdfg::gpu::GpuMmaDispatcherStrategy::EltwiseCompute

Initial value:

 std::function<void(
        codegen::CodegenOutput& out,
        const std::string& main_frag_elem,
        const std::string& idx,
        const std::vector<std::string>& other_frag_elems
    )>

Per-element callback used by emit_eltwise: writes one output element from its inputs.

Constructor & Destructor Documentation

~GpuMmaDispatcherStrategy()

virtual sdfg::gpu::GpuMmaDispatcherStrategy::~GpuMmaDispatcherStrategy ( )

virtualdefault

Member Function Documentation

emit_eltwise()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::emit_eltwise ( codegen::CodegenOutput & out, const std::string & main_frag, const std::vector< std::string > & additional_frags, const EltwiseCompute & compute ) const

pure virtual

Element-wise loop over main_frag, invoking compute for each element.

Implemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

emit_fill()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::emit_fill ( codegen::CodegenOutput & out, const std::string & acc_frag, types::PrimitiveType fill_type ) const

pure virtual

Zero-initialize the accumulator fragment behind acc_frag.

Implemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

emit_fragment_load()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::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

pure virtual

Load a tile from memory (base + layout) into the fragment behind frag.

Implemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

emit_fragment_store()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::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

pure virtual

Store the fragment behind frag to memory (base + layout).

Implemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

emit_mma()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::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

pure virtual

frag_d = frag_a * frag_b + frag_c (single MMA step; all operands are fragments).

Implemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

emit_needed_declarations()

virtual void sdfg::gpu::GpuMmaDispatcherStrategy::emit_needed_declarations ( codegen::CodegenOutput & out ) const

inlinevirtual

Target library includes / helper declarations needed before the first emit.

Reimplemented in sdfg::gpu::cuda::CudaMmaFragmentStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

fragment_element()

virtual std::string sdfg::gpu::GpuMmaDispatcherStrategy::fragment_element ( const std::string & frag, const std::string & idx ) const

inlinevirtual

Expression for element idx of this lane's share of frag (needed for clipped transfers).

Reimplemented in sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

fragment_element_coords()

virtual std::optional< std::pair< std::string, std::string > > sdfg::gpu::GpuMmaDispatcherStrategy::fragment_element_coords ( MmaFragmentType frag_type, const MmaBlockSize & block_size, const std::string & idx, std::string & condition ) const

inlinevirtual

Analytic (row, col) block coordinates of this lane's element idx, if the fragment's register layout is known; clipped transfers otherwise recover them by loading an iota table through the fragment load.

condition receives a preprocessor condition under which the layout holds (empty = always).

Reimplemented in sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

fragment_element_count()

virtual std::string sdfg::gpu::GpuMmaDispatcherStrategy::fragment_element_count ( const std::string & frag ) const

inlinevirtual

Number of elements of this lane's share of frag (needed for clipped transfers).

Reimplemented in sdfg::gpu::rocm::RocmMmaBareStrategy, and sdfg::gpu::rocm::RocmMmaFragmentStrategy.

---

The documentation for this class was generated from the following file:

• opt/include/sdfg/targets/gpu/gpu_mma_dispatcher.h