sdfg::gpu::rocm::RocmMmaFragmentStrategy Class Reference

Classic rocwmma fragment-based emit strategy (load/store/mma_sync/fill_fragment). More...

#include <sdfg/targets/rocm/rocm_mma_fragment_dispatcher.h>

Inheritance diagram for sdfg::gpu::rocm::RocmMmaFragmentStrategy:

Public Member Functions
void	emit_needed_declarations (codegen::CodegenOutput &out) const override
	Target library includes / helper declarations needed before the first emit.
void	emit_fill (codegen::CodegenOutput &out, const std::string &acc_frag, types::PrimitiveType fill_type) const override
	Zero-initialize the accumulator fragment behind acc_frag.
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
	frag_d = frag_a * frag_b + frag_c (single MMA step; all operands are fragments).
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
	Load a tile from memory (base + layout) into the fragment behind frag.
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
	Store the fragment behind frag to memory (base + layout).
void	emit_eltwise (codegen::CodegenOutput &out, const std::string &main_frag, const std::vector< std::string > &additional_frags, const EltwiseCompute &compute) const override
	Element-wise loop over main_frag, invoking compute for each element.
std::string	fragment_element (const std::string &frag, const std::string &idx) const override
	Expression for element idx of this lane's share of frag (needed for clipped transfers).
std::string	fragment_element_count (const std::string &frag) const override
	Number of elements of this lane's share of frag (needed for clipped transfers).
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
	On wave64 targets (per the queried wavefront size), 16x16 accumulators keep the MFMA output layout: lane l holds column l % 16, rows 4 * (l / 16) + e.
Public Member Functions inherited from sdfg::gpu::GpuMmaDispatcherStrategy
virtual	~GpuMmaDispatcherStrategy ()=default

Additional Inherited Members
Public Types inherited from sdfg::gpu::GpuMmaDispatcherStrategy
using	EltwiseCompute
	Per-element callback used by emit_eltwise: writes one output element from its inputs.

Detailed Description

Classic rocwmma fragment-based emit strategy (load/store/mma_sync/fill_fragment).

Member Function Documentation

emit_eltwise()

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

overridevirtual

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

Implements sdfg::gpu::GpuMmaDispatcherStrategy.

emit_fill()

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

overridevirtual

Zero-initialize the accumulator fragment behind acc_frag.

Implements sdfg::gpu::GpuMmaDispatcherStrategy.

emit_fragment_load()

void sdfg::gpu::rocm::RocmMmaFragmentStrategy::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

overridevirtual

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

Implements sdfg::gpu::GpuMmaDispatcherStrategy.

emit_fragment_store()

void sdfg::gpu::rocm::RocmMmaFragmentStrategy::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

overridevirtual

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

Implements sdfg::gpu::GpuMmaDispatcherStrategy.

emit_mma()

void sdfg::gpu::rocm::RocmMmaFragmentStrategy::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

overridevirtual

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

Implements sdfg::gpu::GpuMmaDispatcherStrategy.

emit_needed_declarations()

void sdfg::gpu::rocm::RocmMmaFragmentStrategy::emit_needed_declarations ( codegen::CodegenOutput & out ) const

overridevirtual

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

Reimplemented from sdfg::gpu::GpuMmaDispatcherStrategy.

fragment_element()

std::string sdfg::gpu::rocm::RocmMmaFragmentStrategy::fragment_element ( const std::string & frag, const std::string & idx ) const

inlineoverridevirtual

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

Reimplemented from sdfg::gpu::GpuMmaDispatcherStrategy.

fragment_element_coords()

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

overridevirtual

On wave64 targets (per the queried wavefront size), 16x16 accumulators keep the MFMA output layout: lane l holds column l % 16, rows 4 * (l / 16) + e.

Other wavefront sizes fall back to the iota path.

Reimplemented from sdfg::gpu::GpuMmaDispatcherStrategy.

fragment_element_count()

std::string sdfg::gpu::rocm::RocmMmaFragmentStrategy::fragment_element_count ( const std::string & frag ) const

inlineoverridevirtual

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

Reimplemented from sdfg::gpu::GpuMmaDispatcherStrategy.

---

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

• opt/include/sdfg/targets/rocm/rocm_mma_fragment_dispatcher.h