Tensor matrix multiplication following ONNX MatMul operator specification. More...

#include <sdfg/data_flow/library_nodes/math/tensor/matmul_node.h>

Inheritance diagram for sdfg::math::tensor::MatMulNode:

Public Member Functions
	MatMulNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, data_flow::DataFlowGraph &parent, const TensorLayout &layout_a, const TensorLayout &layout_b, QuantizationType quantization=QUANTIZATION_MATCH_INPUTS, const TensorLayout *layout_y=nullptr, const data_flow::ImplementationType &impl_type=data_flow::ImplementationType_NONE)
	Construct a matmul node.
	MatMulNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, data_flow::DataFlowGraph &parent, const TensorLayout &layout_a, const TensorLayout &layout_b, const ContractionTypes &types, const TensorLayout *layout_y=nullptr, const data_flow::ImplementationType &impl_type=data_flow::ImplementationType_NONE)
	Construct a matmul node with explicit element types.
symbolic::Expression	m () const
	Get the M dimension (rows of A, rows of output).
symbolic::Expression	n () const
	Get the N dimension (columns of B, columns of output).
const ContractionTypes &	types () const
	The configured element types; Void entries are inferred, see input_type(), accumulator_type(), output_type().
void	set_types (const ContractionTypes &types)
types::PrimitiveType	input_type () const
	Element type of A and B: configured, or inferred from their memlets (Void if they disagree).
types::PrimitiveType	accumulator_type () const
	Element type the products are summed in: configured, or the output type.
types::PrimitiveType	output_type () const
	Element type of Y: configured, or inferred from its memlet.
std::optional< types::PrimitiveType >	uniform_quantization (const data_flow::DataFlowGraph &dataflow) const
	The single type of inputs, accumulator and output, if they all agree.
QuantizationType	quantization () const
QuantizationType	quantization (const data_flow::DataFlowGraph &dataflow) const
types::PrimitiveType	input_quantization () const
types::PrimitiveType	output_quantization () const
QuantizationType	fixed_quantization () const
	The configured type if inputs and output share one (Void otherwise).
void	set_fixed_quantization (const QuantizationType quant)
const TensorLayout &	layout_a () const
const TensorLayout &	layout_b () const
const TensorLayout &	layout_y () const
TensorLayout &	layout_a ()
TensorLayout &	layout_b ()
TensorLayout &	layout_y ()
symbolic::Expression	k () const
	Get the K dimension (columns of A, rows of B - contraction dimension).
void	validate (const Function &function) const override
	Validate the code node's configuration.
void	verify_data_types (const data_flow::DataFlowGraph &graph) const override
	Called as part of validate.
virtual passes::LibNodeExpander::ExpandOutcome	expand (passes::LibNodeExpander::ExpandContext &, structured_control_flow::Block &block) override
	Expand matmul into nested maps.
symbolic::SymbolSet	symbols () const override
	Get all symbols used in this node.
void	replace (const symbolic::Expression old_expression, const symbolic::Expression new_expression) override
void	replace (const symbolic::ExpressionMapping &replacements) override
std::unique_ptr< data_flow::DataFlowNode >	clone (size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph &parent) const override
	Clone this node for graph transformations.
std::string	toStr () const override
	Convert node to string representation.
symbolic::Expression	flop () const override
	Calculate floating point operations for this node.
data_flow::PointerAccessType	pointer_access_type (int input_idx) const override
	Describes what a pointer is used for.
bool	relocalize_operand (int input_idx, const TensorLayout &packed) override
	Repoint the operand at `input_idx` to a locally staged buffer.
bool	can_relocalize_operand (int input_idx, const TensorLayout &packed) const override
	Side-effect-free predicate mirroring relocalize_operand: whether the operand at `input_idx` can be repointed at the buffer view `packed`.
bool	supports_integer_types () const override
	Check if this operation supports integer types.
Public Member Functions inherited from sdfg::math::tensor::TensorNode
	TensorNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, data_flow::DataFlowGraph &parent, const data_flow::LibraryNodeCode &code, const std::vector< std::string > &outputs, const std::vector< std::string > &inputs, const data_flow::ImplementationType &impl_type)
	Construct a tensor node.
void	validate (const Function &function) const override
	Validate tensor node constraints.
types::PrimitiveType	primitive_type (const data_flow::DataFlowGraph &graph) const
	Get the primitive type of this tensor operation.
Public Member Functions inherited from sdfg::math::MathNode
	MathNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, data_flow::DataFlowGraph &parent, const data_flow::LibraryNodeCode &code, const std::vector< std::string > &outputs, const std::vector< std::string > &inputs, const data_flow::ImplementationType &implementation_type, bool side_effects=false)
	Construct a mathematical library node.
Public Member Functions inherited from sdfg::data_flow::LibraryNode
	LibraryNode (const LibraryNode &data_node)=delete
LibraryNode &	operator= (const LibraryNode &)=delete
virtual	~LibraryNode ()=default
ElementType	type_id () const override
	Returns the type of the element.
const LibraryNodeCode &	code () const
	Get the operation code.
const ImplementationType &	implementation_type () const
	Get the implementation type (const).
ImplementationType &	implementation_type ()
	Get the implementation type (mutable).
void	set_implementation_type (const ImplementationType &impl_type)
bool	side_effect () const override
	Check if this node has side effects.
PointerAccessType	pointer_access_type (const Memlet &edge) const
virtual std::optional< SubgroupAccess >	operand_access (int input_idx, unsigned subgroup_size) const
	The elements one subgroup of `subgroup_size` lanes reads from the operand at `input_idx`, if the node declares them.
EdgeRemoveOption	can_remove_out_edge (const data_flow::DataFlowGraph &graph, const Memlet *memlet) const override
	Is the edge in question removable, or will this make the node its an output on invalid?
EdgeRemoveOption	can_remove_in_edge (const data_flow::DataFlowGraph &graph, const Memlet *memlet) const override
bool	pointer_use_creates_side_effects (const DataFlowGraph &dataflow, const Function &func)
Public Member Functions inherited from sdfg::data_flow::CodeNode
	CodeNode (const CodeNode &data_node)=delete
CodeNode &	operator= (const CodeNode &)=delete
const std::vector< std::string > &	outputs () const
	Get output connector names (const).
const std::vector< std::string > &	inputs () const
	Get input connector names (const).
std::vector< std::string > &	outputs ()
	Get output connector names (mutable).
std::vector< std::string > &	inputs ()
	Get input connector names (mutable).
const std::string &	output (size_t index) const
	Get output connector name by index.
const std::string &	input (size_t index) const
	Get input connector name by index.
bool	has_constant_input (size_t index) const
	Check if an input is connected to a constant node.
Public Member Functions inherited from sdfg::data_flow::DataFlowNode
	DataFlowNode (const DataFlowNode &data_node)=delete
DataFlowNode &	operator= (const DataFlowNode &)=delete
graph::Vertex	vertex () const
	Get the graph vertex for this node.
const DataFlowGraph &	get_parent () const
	Get the parent dataflow graph (const).
DataFlowGraph &	get_parent ()
	Get the parent dataflow graph (mutable).
virtual bool	update_edge_removed (const std::string &out_conn)
Public Member Functions inherited from sdfg::Element
	Element (size_t element_id, const DebugInfo &debug_info)
virtual	~Element ()=default
size_t	element_id () const
const DebugInfo &	debug_info () const
void	set_debug_info (const DebugInfo &debug_info)
void	add_metadata (const std::string &key, const std::string &value)
void	remove_metadata (const std::string &key)
const std::string &	metadata (const std::string &key) const
const std::string *	metadata_if_exists (const std::string &key) const
const std::unordered_map< std::string, std::string > &	metadata () const
std::string_view	element_type () const
	Returns the type of the element as a human-readable string.

Static Public Member Functions
static TensorLayout	get_linear_result_layout (const TensorLayout &layout_a, const TensorLayout &layout_b)
Static Public Member Functions inherited from sdfg::data_flow::LibraryNode
static bool	classof (const Element &element)
Static Public Member Functions inherited from sdfg::data_flow::CodeNode
static bool	classof (const Element &element)
	LLVM-style RTTI predicate: true if `element` is a CodeNode.
Static Public Member Functions inherited from sdfg::data_flow::DataFlowNode
static bool	classof (const Element &element)
	LLVM-style RTTI predicate: true if `element` is a DataFlowNode.

Static Public Attributes
static auto constexpr	Y_INPUT_IDX = 0
static auto constexpr	A_INPUT_IDX = 1
static auto constexpr	B_INPUT_IDX = 2
static constexpr const char *	OPERAND_METADATA_KEY = "ml.matmul.operand"
	Metadata key tagging a memlet with the matmul operand it carries ("A", "B", or "Y"), attached during expansion so the originating operand can be matched on the expanded dataflow.
Static Public Attributes inherited from sdfg::data_flow::CodeNode
static constexpr ElementType	TypeGroup = ElementType::Tasklet \| ElementType::LibraryNode
	Bitmask of all concrete ElementTypes that are CodeNodes.
Static Public Attributes inherited from sdfg::data_flow::DataFlowNode
static constexpr ElementType	TypeGroup
	Bitmask of all concrete ElementTypes that are DataFlowNodes.

Additional Inherited Members
Protected Member Functions inherited from sdfg::math::tensor::TensorNode
void	validate_shape_matches (const std::vector< symbolic::Expression > &required_shape, const TensorLayout &layout, const std::string &name) const
Protected Member Functions inherited from sdfg::data_flow::LibraryNode
	LibraryNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, DataFlowGraph &parent, const LibraryNodeCode &code, const std::vector< std::string > &outputs, const std::vector< std::string > &inputs, const bool side_effect, const ImplementationType &implementation_type)
	Protected constructor for library nodes.
Protected Member Functions inherited from sdfg::data_flow::CodeNode
	CodeNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, DataFlowGraph &parent, const std::vector< std::string > &outputs, const std::vector< std::string > &inputs)
	Protected constructor for code nodes.
Protected Member Functions inherited from sdfg::data_flow::DataFlowNode
	DataFlowNode (size_t element_id, const DebugInfo &debug_info, const graph::Vertex vertex, DataFlowGraph &parent)
	Protected constructor for dataflow nodes.
Static Protected Member Functions inherited from sdfg::math::tensor::TensorNode
static data_flow::TaskletCode	get_integer_minmax_tasklet (types::PrimitiveType prim_type, bool is_max)
	Get the appropriate tasklet code for min/max operations on integers.
Protected Attributes inherited from sdfg::data_flow::LibraryNode
LibraryNodeCode	code_
	Operation identifier (e.g., "GEMM", "Add").
bool	side_effect_
	Whether this node has side effects.
ImplementationType	implementation_type_
	How to implement this node.
Protected Attributes inherited from sdfg::data_flow::CodeNode
std::vector< std::string >	outputs_
	Names of output connectors.
std::vector< std::string >	inputs_
	Names of input connectors.
Protected Attributes inherited from sdfg::Element
size_t	element_id_
DebugInfo	debug_info_
std::unordered_map< std::string, std::string >	metadata_

Detailed Description

Tensor matrix multiplication following ONNX MatMul operator specification.

MatMulNode represents a tensor matrix multiplication operation that is compatible with the ONNX MatMul operator. The operation performs matrix multiplication on the last two dimensions and supports broadcasting over batch dimensions.

Input/Output Requirements

Input connector "A": Input tensor [..., M, K]
Input connector "B": Input tensor [..., K, N]
Output connector "Y": Output tensor [..., M, N]

Broadcasting

The batch dimensions are broadcast following numpy broadcasting rules:

Dimensions are compared from right to left (excluding last two matrix dims)
Dimensions must be equal, or one of them must be 1
The output shape takes the maximum of each dimension

Example

For inputs A[B, M, K] and B[B, K, N]:

Y = A @ B has shape [B, M, N]
Each (b, m, n) element is computed as: sum_k(A[b, m, k] * B[b, k, n])

Element types

A and B share the input type; the products are summed in the accumulator type and stored in the output type (see ContractionTypes), e.g. Int8 x Int8 accumulated and stored as Int32.

Constructor & Destructor Documentation

◆ MatMulNode() [1/2]

sdfg::math::tensor::MatMulNode::MatMulNode	(	size_t	element_id,
		const DebugInfo &	debug_info,
		const graph::Vertex	vertex,
		data_flow::DataFlowGraph &	parent,
		const TensorLayout &	layout_a,
		const TensorLayout &	layout_b,
		QuantizationType	quantization = QUANTIZATION_MATCH_INPUTS,
		const TensorLayout *	layout_y = nullptr,
		const data_flow::ImplementationType &	impl_type = data_flow::ImplementationType_NONE )

Construct a matmul node.

Parameters

element_id	Unique element identifier
debug_info	Debug information
vertex	Graph vertex
parent	Parent dataflow graph

◆ MatMulNode() [2/2]

sdfg::math::tensor::MatMulNode::MatMulNode	(	size_t	element_id,
		const DebugInfo &	debug_info,
		const graph::Vertex	vertex,
		data_flow::DataFlowGraph &	parent,
		const TensorLayout &	layout_a,
		const TensorLayout &	layout_b,
		const ContractionTypes &	types,
		const TensorLayout *	layout_y = nullptr,
		const data_flow::ImplementationType &	impl_type = data_flow::ImplementationType_NONE )

Construct a matmul node with explicit element types.

Parameters

types input, accumulator and output types; Void entries are inferred

Member Function Documentation

◆ accumulator_type()

types::PrimitiveType sdfg::math::tensor::MatMulNode::accumulator_type ( ) const

Element type the products are summed in: configured, or the output type.

◆ can_relocalize_operand()

bool sdfg::math::tensor::MatMulNode::can_relocalize_operand	(	int	input_idx,
		const TensorLayout &	packed ) const

overridevirtual

Side-effect-free predicate mirroring relocalize_operand: whether the operand at input_idx can be repointed at the buffer view packed.

Reimplemented from sdfg::data_flow::LibraryNode.

◆ clone()

std::unique_ptr< data_flow::DataFlowNode > sdfg::math::tensor::MatMulNode::clone	(	size_t	element_id,
		const graph::Vertex	vertex,
		data_flow::DataFlowGraph &	parent ) const

overridevirtual

Clone this node for graph transformations.

Parameters

element_id	New element identifier for the clone
vertex	New graph vertex for the clone
parent	Parent graph for the clone

Returns: Unique pointer to the cloned node

Pure virtual function that must be implemented by derived classes to support graph transformations and optimizations.

Implements sdfg::data_flow::DataFlowNode.

◆ expand()

virtual passes::LibNodeExpander::ExpandOutcome sdfg::math::tensor::MatMulNode::expand	(	passes::LibNodeExpander::ExpandContext &	,
		structured_control_flow::Block &	block )

overridevirtual

Expand matmul into nested maps.

Expands the matmul operation by:

Creating outer maps for parallel iteration over batch and M, N dimensions
Creating inner for loop for sequential accumulation over K dimension
Computing matrix multiplication using FMA (fused multiply-add) tasklets
Writing results to output tensor

Parameters

context	way to request isolation of node or add new blocks, loops etc.
block	the block in which the node is contained

Returns: outcome of expand, created by context

Implements sdfg::math::MathNode.

◆ fixed_quantization()

QuantizationType sdfg::math::tensor::MatMulNode::fixed_quantization ( ) const

The configured type if inputs and output share one (Void otherwise).

◆ flop()

symbolic::Expression sdfg::math::tensor::MatMulNode::flop ( ) const

overridevirtual

Calculate floating point operations for this node.

Returns: Symbolic expression for FLOP count (or null)

Reimplemented from sdfg::data_flow::LibraryNode.

◆ get_linear_result_layout()

TensorLayout sdfg::math::tensor::MatMulNode::get_linear_result_layout	(	const TensorLayout &	layout_a,
		const TensorLayout &	layout_b )

static

◆ input_quantization()

types::PrimitiveType sdfg::math::tensor::MatMulNode::input_quantization ( ) const

◆ input_type()

types::PrimitiveType sdfg::math::tensor::MatMulNode::input_type ( ) const

Element type of A and B: configured, or inferred from their memlets (Void if they disagree).

◆ k()

symbolic::Expression sdfg::math::tensor::MatMulNode::k ( ) const

Get the K dimension (columns of A, rows of B - contraction dimension).

Returns: K dimension expression

◆ layout_a() [1/2]

TensorLayout & sdfg::math::tensor::MatMulNode::layout_a ( )

inline

◆ layout_a() [2/2]

const TensorLayout & sdfg::math::tensor::MatMulNode::layout_a ( ) const

◆ layout_b() [1/2]

TensorLayout & sdfg::math::tensor::MatMulNode::layout_b ( )

inline

◆ layout_b() [2/2]

const TensorLayout & sdfg::math::tensor::MatMulNode::layout_b ( ) const

◆ layout_y() [1/2]

TensorLayout & sdfg::math::tensor::MatMulNode::layout_y ( )

inline

◆ layout_y() [2/2]

const TensorLayout & sdfg::math::tensor::MatMulNode::layout_y ( ) const

◆ m()

symbolic::Expression sdfg::math::tensor::MatMulNode::m ( ) const

Get the M dimension (rows of A, rows of output).

Returns: M dimension expression

◆ n()

symbolic::Expression sdfg::math::tensor::MatMulNode::n ( ) const

Get the N dimension (columns of B, columns of output).

Returns: N dimension expression

◆ output_quantization()

types::PrimitiveType sdfg::math::tensor::MatMulNode::output_quantization ( ) const

◆ output_type()

types::PrimitiveType sdfg::math::tensor::MatMulNode::output_type ( ) const

Element type of Y: configured, or inferred from its memlet.

◆ pointer_access_type()

data_flow::PointerAccessType sdfg::math::tensor::MatMulNode::pointer_access_type ( int input_idx ) const

overridevirtual

Describes what a pointer is used for.

Parameters

input_idx index of input that is a pointer.

Returns: Invalid if not asked about a pointer input

Reimplemented from sdfg::data_flow::LibraryNode.

◆ quantization() [1/2]

QuantizationType sdfg::math::tensor::MatMulNode::quantization ( ) const

inline

◆ quantization() [2/2]

QuantizationType sdfg::math::tensor::MatMulNode::quantization ( const data_flow::DataFlowGraph & dataflow ) const

◆ relocalize_operand()

bool sdfg::math::tensor::MatMulNode::relocalize_operand	(	int	input_idx,
		const TensorLayout &	packed )

overridevirtual

Repoint the operand at input_idx to a locally staged buffer.

packed is the operand's own view of that buffer: its shape, the buffer's strides, and the buffer offset of its origin (in elements; may be a sub-window of the staged tile). Updates the node's addressing metadata to match. Returns false if the node cannot express the view (localization is then rejected).

Reimplemented from sdfg::data_flow::LibraryNode.

◆ replace() [1/2]

void sdfg::math::tensor::MatMulNode::replace	(	const symbolic::Expression	old_expression,
		const symbolic::Expression	new_expression )

overridevirtual

Implements sdfg::Element.

◆ replace() [2/2]

void sdfg::math::tensor::MatMulNode::replace ( const symbolic::ExpressionMapping & replacements )

overridevirtual

Reimplemented from sdfg::Element.

◆ set_fixed_quantization()

void sdfg::math::tensor::MatMulNode::set_fixed_quantization ( const QuantizationType quant )

◆ set_types()

void sdfg::math::tensor::MatMulNode::set_types ( const ContractionTypes & types )

◆ supports_integer_types()

bool sdfg::math::tensor::MatMulNode::supports_integer_types ( ) const

inlineoverridevirtual

Check if this operation supports integer types.

Operations like add, sub, mul, div, min, max support both floating-point and integer types. Operations like exp, log, tanh, sigmoid only support floating-point types.

Returns: true if integer types are supported, false otherwise

Implements sdfg::math::tensor::TensorNode.

◆ symbols()

symbolic::SymbolSet sdfg::math::tensor::MatMulNode::symbols ( ) const

overridevirtual

Get all symbols used in this node.

Returns: Set of symbolic expressions used by this node

Implements sdfg::data_flow::LibraryNode.

◆ toStr()

std::string sdfg::math::tensor::MatMulNode::toStr ( ) const

overridevirtual

Convert node to string representation.

Returns: String representation of the node

Reimplemented from sdfg::data_flow::LibraryNode.

◆ types()

const ContractionTypes & sdfg::math::tensor::MatMulNode::types ( ) const

The configured element types; Void entries are inferred, see input_type(), accumulator_type(), output_type().

◆ uniform_quantization()

std::optional< types::PrimitiveType > sdfg::math::tensor::MatMulNode::uniform_quantization ( const data_flow::DataFlowGraph & dataflow ) const

The single type of inputs, accumulator and output, if they all agree.

None if the matmul mixes types, such as Int8 inputs into an Int32 output.

◆ validate()

void sdfg::math::tensor::MatMulNode::validate ( const Function & function ) const

overridevirtual

Validate the code node's configuration.

Parameters

function Function context for validation

Ensures that:

No two access nodes connect to the same data for inputs/outputs

Reimplemented from sdfg::data_flow::CodeNode.

◆ verify_data_types()

void sdfg::math::tensor::MatMulNode::verify_data_types ( const data_flow::DataFlowGraph & graph ) const

overridevirtual

Called as part of validate.

By default checks that all primitive types match (TODO do not do this for all TensorOps, too restrictive)

Parameters

graph

Reimplemented from sdfg::math::tensor::TensorNode.

Member Data Documentation

◆ A_INPUT_IDX

auto constexpr sdfg::math::tensor::MatMulNode::A_INPUT_IDX = 1

staticconstexpr

◆ B_INPUT_IDX

auto constexpr sdfg::math::tensor::MatMulNode::B_INPUT_IDX = 2

staticconstexpr

◆ OPERAND_METADATA_KEY

const char* sdfg::math::tensor::MatMulNode::OPERAND_METADATA_KEY = "ml.matmul.operand"

staticconstexpr

Metadata key tagging a memlet with the matmul operand it carries ("A", "B", or "Y"), attached during expansion so the originating operand can be matched on the expanded dataflow.

◆ Y_INPUT_IDX

auto constexpr sdfg::math::tensor::MatMulNode::Y_INPUT_IDX = 0

staticconstexpr

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

sdfg/include/sdfg/data_flow/library_nodes/math/tensor/matmul_node.h

Public Member Functions

Static Public Member Functions

Static Public Attributes

Additional Inherited Members

Detailed Description

Input/Output Requirements

Broadcasting

Example

Element types

Constructor & Destructor Documentation

◆ MatMulNode() [1/2]

◆ MatMulNode() [2/2]

Member Function Documentation

◆ accumulator_type()

◆ can_relocalize_operand()

◆ clone()

◆ expand()

◆ fixed_quantization()

◆ flop()

◆ get_linear_result_layout()

◆ input_quantization()

◆ input_type()

◆ k()

◆ layout_a() [1/2]

◆ layout_a() [2/2]

◆ layout_b() [1/2]

◆ layout_b() [2/2]

◆ layout_y() [1/2]

◆ layout_y() [2/2]

◆ m()

◆ n()

◆ output_quantization()

◆ output_type()

◆ pointer_access_type()

◆ quantization() [1/2]

◆ quantization() [2/2]

◆ relocalize_operand()

◆ replace() [1/2]

◆ replace() [2/2]

◆ set_fixed_quantization()

◆ set_types()

◆ supports_integer_types()

◆ symbols()

◆ toStr()

◆ types()

◆ uniform_quantization()

◆ validate()

◆ verify_data_types()

Member Data Documentation

◆ A_INPUT_IDX

◆ B_INPUT_IDX

◆ OPERAND_METADATA_KEY

◆ Y_INPUT_IDX