matmul_node.h File Reference

Tensor matrix multiplication node compatible with ONNX MatMul operator. More...

#include <optional>
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/symbolic/symbolic.h"

Go to the source code of this file.

Classes
class	sdfg::math::tensor::MatMulNode
	Tensor matrix multiplication following ONNX MatMul operator specification. More...
class	sdfg::math::tensor::MatMulNodeSerializer
	Serializer for MatMulNode. More...
class	sdfg::math::tensor::MatMulNaiveExpander
class	sdfg::math::tensor::MatMulBlasExpander

Namespaces
namespace	sdfg
namespace	sdfg::math
namespace	sdfg::math::tensor

Functions
data_flow::LibraryNodeCode	sdfg::math::tensor::LibraryNodeType_MatMul ("ml::MatMul")

Detailed Description

Tensor matrix multiplication node compatible with ONNX MatMul operator.

This file defines the MatMulNode class which implements a tensor matrix multiplication operation following the ONNX MatMul operator specification.

ONNX MatMul Operator Compatibility

The MatMulNode implements the ONNX MatMul operator with the following semantics:

• Input tensor A: [..., M, K] - arbitrary batch dimensions followed by matrix dims
• Input tensor B: [..., K, N] - arbitrary batch dimensions followed by matrix dims
• Output tensor Y: [..., M, N] - broadcasted batch dimensions with result matrix

The operation performs matrix multiplication on the last two dimensions and broadcasts over the batch dimensions following numpy broadcasting rules.

Expansion

The matmul operation is expanded into nested maps:

1. Create outer maps for parallel iteration over batch and output dimensions
2. Create inner loop for sequential accumulation over the K dimension
3. Compute matrix multiplication using FMA (fused multiply-add) operations
4. Write results to output tensor

Example

Creating a batched matrix multiplication:

symbolic::MultiExpression shape_a = {symbolic::symbol("B"), symbolic::symbol("M"), symbolic::symbol("K")};
symbolic::MultiExpression shape_b = {symbolic::symbol("B"), symbolic::symbol("K"), symbolic::symbol("N")};
 
auto& matmul_node = builder.add_library_node<math::tensor::MatMulNode>(
    block, debug_info, shape_a, shape_b
);