Implementing Linear Layers and Memory Optimizations on MLIR

In this article, I am expanding the basic pipeline from the previous article to support a Linear Layer (the foundational building block of modern neural networks). At the hardware level, this workload is driven by a GeMM (General Matrix Multiply) operation, traditionally defined as:

C ← αAB + βC

To achieve this, I added two new operations to the compiler: relu and matmul. Since relu shares a similar element-wise structure to the 'add' operator covered in Part 1, this article will focus primarily on the unique challenges of building the matmul operator. Additionally, I will dive into optimizing performance through cache locality and memory access using MLIR's built-in infrastructure, specifically pattern rewrite engines and custom passes driven by interfaces.

Github repo for the project is available here

MatMul Operation

To create a MatMul op, start by adding it to the Ops.td as outlined in part 1's article for the 'Add' op. It takes 2 arguments for lhs and rhs and returns a single output tensor. I also include a builders block to declare that a custom C++ builder will be written for it. While the 'Add' op can simply copy its output shape directly from its lhs/rhs operand, a matrix multiplication's output shape is mathematically determined by combining the outer dimension of both inputs (MxN and NxO yields MxO). Because MLIR cannot automatically infer this shape from just one operand, a custom C++ builder must be written for it to handle this dimension calculation

def MatMulOp : Dlc_Op<"matmul"> {
	let summary = "matrix multiplication";
	let arguments = (ins F32Tensor:$lhs, F32Tensor:$rhs);
	let results = (outs F32Tensor);
	let hasVerifier = 1;
	let builders = [
		OpBuilder<(ins "Value":$lhs, "Value": $rhs)>
	];
}

Custom Builder

As noted above, the purpose of the custom builder for 'matmul' is to infer the output shape. This implementation specifically checks the tensor rank of both inputs to handle various combinations of 1D vectors and 2D matrices (such as vector-matrix, matrix-vector, matrix-matrix, and vector-vector dot products).

It extracts the outer dimensions (M and N), verifies via an 'assert' that the inner reduction dimension (K) match, and constructs the appropriate output 'RankedTensorType' before packing everything into the operation state.

// tensor-compiler/src/mlir/Dialect.cpp

void MatMulOp::build(OpBuilder &builder, OperationState &state,
					Value lhs, Value rhs) {
	// extract operands' type and primitive data type
	auto lhsType = llvm::cast<RankedTensorType>(lhs.getType());
	auto rhsType = llvm::cast<RankedTensorType>(rhs.getType());
	auto elemType = lhsType.getElementType();

	// MatMul Shape logic: (M x K) @ (K x N) -> (M x N)
	auto lhsShape = lhsType.getShape();
	auto rhsShape = rhsType.getShape();

	bool lhsIsVector = lhsShape.size() == 1;
	bool rhsIsVector = rhsShape.size() == 1;

	// Inner dimension check (K)
	int64_t lhsK = lhsShape.back();
	int64_t rhsK = rhsShape.back();
	assert(lhsK == rhsK && "matmul inner dimensions must match");

	RankedTensorType resultType;

	if (lhsIsVector && !rhsIsVector) {
		// [K] @ [K, N] -> [N]
		resultType = RankedTensorType::get(
			{rhsShape[1]}, elemType
		);
	} else if (!lhsIsVector && rhsIsVector) {
		// [M, K] @ [K] -> [M]
		resultType = RankedTensorType::get(
			{lhsShape[0]}, elemType
		);
	} else if (!lhsIsVector && !rhsIsVector) {
		// [M, K] @ [K, N] -> [M, N]
		resultType = RankedTensorType::get(
			{lhsShape[0], rhsShape[1]}, elemType
		);
	} else {
		// [K] @ [K] -> scalar
		// Fallback for vector-vector
		resultType = RankedTensorType::get(
			{}, elemType
		);
	}

	// Push the calculated result type and the inputs into the state
	state.addTypes(resulType);
	state.addOperands({lhs, rhs});
}

For this project, the current framework restricts its scope to 1D vectors and 2D matrices. To do that, I implement a custom verifier to enforce these layout constraints at compile time, it also acts as a mathematical guard, first rejecting any tensors outside the 1D-to-2D range, and then asserting that the inner contraction dimensions (K) align perfectly.

Finally, it validates the structural shape of the output tensor against the exact rank combinations of the inputs (specifically handling vector-matrix, matrix-vector and matrix-matrix operations) while explicitly throwing an error for unsupported vector-vector dot products.

// tensor-compiler/src/mlir/Dialect.cpp

LogicalResult MatMulOp::verify() {
	auto lhsShape = llvm::cast<RankedTensorType>(getLhs().getType()).getShape();
	auto rhsShape = llvm::cast<RankedTensorType>(getRhs().getType()).getShape();
	auto resShape = llvm::cast<RankedTensorType>(getResult().getType()).getShape();

	// Input Rank check
	if (lhsShape.size() < 1 || lhsShape.size() > 2 ||
		rhsShape.size() < 1 || rhsShape.size() > 2) {
		return emitOpError("LHS and RHS must be either 1D or 2D tensors");
	}

	// Inner Dimension (K) check
	int64_t lhsK = lhsShape.back();
	int64_t rhsK = rhsShape[0];
	if (lhs != rhsK) {
		return emitOpError("Inner dimensions (K) must match");
	}

	int64_t lhsRank = lhsShape.size();
	int64_t rhsRank = rhsShape.size();
	int64_t resRank = resShape.size();

	// vector @ matrix -> vector [N]
	if (lhsRank == 1 && rhsRank == 2) {
		if (resRank != 1 || resShape[0] != rhsShape[1]) {
			return emitOpError("expected result shape [N] for vector x matrix");
		}
	}
	// matrix @ vector -> vector [M]
	else if (lhsRank == 2 && rhsRank == 1) {
		if (resRank != 1 || resShape[0] != lhsShape[0]) {
			return emitOpError("expected result shape [M] for matrix x vector");
		}
	}
	// matrix @ matrix -> matrix [M, N]
	else if (lhsRank == 2 && rhsRank == 2) {
		if (resRank != 2 || resShape[0] != lhsShape[0] || resShape[1] != lhsShape[1]) {
			return emitOpError("expected result shape [M, N] for matrix x matrix");
		}
	}
	else { return emitOpError("vector x vector matmul not supported"); }
	return success();
}

Transpose right-hand-side matrix

The first optimization I performed in this project is transposing the rhs matrix for better memory access. In a naive Matrix multiplication, the RHS matrix is accessed by column as seen in the 'before' picture above. Since standard C++ and machine learning formats like ONNX represent tensors in a row-major layout, walking down a column forces the processor to jump across large strides in linear memory. This triggers continuous cache misses, fetching a brand-new cache line for every single value and evicting useful data before it can be reused. (Note that if we were working with a column-major layout, like in Fortran or MATLAB, this pattern would invert, and we would transpose the LHS matrix instead)

By transposing the RHS matrix during calculation, the algorithm's access pattern is aligned with how the data is sequentially laid out in physical memory. As shown in the 'After' picture above, a single cache line fetch now provides multiple subsequent values (highlighted in green). For this implementation, this transformation reduced L1 Data Misses by 20x in the benchmark and accelerated execution speed by roughly 6x, representing an 83% reduction in total runtime. (Performance Analysis discussed in more depth in the next article)

To keep the compiler's pipeline streamlined, I implemented the transpose rhs matrix (transposeB) optimization directly within the lowering pattern from dlc to linalg. This allows for the conditional generation of highly optimized linalg.generic indexing maps on the fly using a configuration flag.

While a production compiler would typically isolate this into a separate high-level graph transformation pass to enable optimization sharing across layers (such as hoisting a single transpose out of multiple layers via CSE), the direct-lowering approach provides a clean self-contained method for achieving 6x performance milestone without adding pass overhead.

Start by setting up the conversion pattern like the previous article and store a 'useTransposeB' boolean directly inside the lowering struct so it can be easily toggled on or off during runtime

// tensor-compiler/src/mlir/LowerToTensor.cpp

struct MatMulOpLowering : public OpConversionPattern<dlc::MatMulOp> {
	bool useTransposeB;
	MatMulOpLowering(MLIRContext *context, bool useTransposeB) :
		OpConversionPattern<dlc::MatMulOp>(context), useTransposeB(useTransposeB) {}

As with previous lowerings (from the previous article), extract the source location 'loc' for debugging and pulls out the LHS and RHS operands.

Before handling the transposition, the tensor ranks need to be normalized. To avoid writing separate execution pipelines for Vector-Matrix and Matrix-Matrix multiplication, promote any rank-1 vector operands to rank-2 matrices early.

To do this, extract the vector length 'shape[0]' and define a new target type of '[1, K]'. Then create a 'tensor::ExpandShapeOp' using reassociation map of {{0, 1}}. This acts as an instruction layout telling MLIR to stretch the single input dimension across both dimensions of the new 2D view. Finally, overwrite the operand's variable with this newly generated operation, passing a clean 2D tensor down the pipeline

// tensor-compiler/src/mlir/LowerToTensor.cpp

if (llvm::cast<RankedTensorType>(operand.getType()).getRank() == 1) {
	// get the original 1D shape [K]
	auto shape = llvm::cast<RankedTensorType>(operand.getType()).getShape();

	// define target 2D type [1, K]
	auto newType = RankedTensorType::get({1, shape[0]},
				llvm::cast<RankedTensorType>(operand.getType()).getElementType());

	// Map original Dim 0 to new Dims 0 and 1
	SmallVector<ReassociationIndices> reassoc = {{0, 1}};
	operand = tensor::ExpandShapeOp::create(rewriter, loc, newType, operand, reassoc);
}

On a high-level, the structural lowering strategy remains identical to the process established in the previous article: extract the shapes, element types, and source locations of the original operands, allocate output buffer tracking Destination Passing Style (DPS), and safely replace the original operation. Crucially, I zero-initialize this destination buffer using 'linalg::FillOp'; without explicitly resetting the memory, the operation would accumulate out-of-date data from previous runs, corrupting the final calculation with garbage values. Because the basic mechanics of DPS allocation and buffer filling has been discussed previously, this article will focus closely on the new mechanics inside the 'useTransposeB' optimization block: specifically how to structurally manipulate the indexing maps inside a custom 'linalg::generic' operation.

With C++ pseudocode, I'm performing the following:

// A @ B = C
// MxK @ KxN = MxN
// transposed: MxK @ [NxK]T = MxN

// transpose the B-matrix
for (int i = 0; i < K; ++i) {
	for (int j = 0; j < N; ++j) {
		b_transpose[j][i] = matrix_b[i][j];
	}
}

// Calculate with transpose
for (int i = 0; i < M; ++i) {
	for (int j = 0; j < N; ++j) {
		for (int k = 0; k < K; ++k) {
			// row-major caching: both matrices are accessed via k at the end
			matrix_C[i][j] += matrix_A[i][k] * b_transpose[j][k];
		}
	}
}

// vs calculating without transpose
for (int i = 0; i < M; ++i) {
	for (int j = 0; j < N; ++j) {
		for (int k = 0; k < K; ++k) {
			// notice how matrix_B's k index is on the left,
			// forcing memory jumps across rows
			matrix_C[i][j] += matrix_A[i][k] * matrix_B[k][j];
		}
	}
}

To transpose within the MLIR lowering pipeline, first allocate an output destination buffer for the RHS matrix matching the shape of the transposed result (NxK). Then define a permutation array (set to {1, 0} for a standard 2D transpose), which serves as a reordering blueprint that maps the input's columns (dim 1) into the output's 0th dimension slot, and vice versa. Finally, create the linalg transpose op with that.

// tensor-compiler/src/mlir/LowerToTensor.cpp
// this is the part where KxN turns into NxK in the RHS matrix

SmallVector<int64_t> transpShape = {N, K};
auto transpInit = tensor::EmptyOp::create(rewriter, loc, transpShape, rhsType.getElementType());
SmallVector<int64_t> perm = {1, 0};
auto transposedB = linalg::TransposeOp::create(rewriter, loc, rhs, transpInit.getResult(), perm);

The next step is define the indexing map (which dictate how the loops navigate through the matrices) and the iteration types. The Affine maps define exactly how to index into each tensor for any given iteration of the loop. Instead of writing explicit loop variables for each operand, MLIR establishes a global space of iteration dimension. For a matmul, there are 3 global loop variables (d0, d1, d2), which correspond directly to (m, n, k)

For instance, look at how matrix A uses this space:

• Matrix A (m, k): It maps to dimension 0 (m) and dimension 2 (k), meaning it is shaped and accessed as [m][k] while completely ignoring n

Following the indexing maps, iterTypes are set to define how those loops actually behave when the indices change. Looking at the C++ code layout above, notice that the output tensor matrix_C is indexed strictly by [i][j] (representing the parallel dimensions, m and n). Because each unique coordinate pair points to an entirely isolated memory location, these calculations can execute simultaneously in parallel without any risk of data interference. Conversely, the inner dimension k increments dynamically while the output address remains locked (meaning it continuously updates and reduces a stream of data into that single, fixed slot, classifying it as a reduction iteration dimension.)

// tensor-compiler/src/mlir/LowerToTensor.cpp

// Map 0 (A): (m, n, k) -> (m, k)
auto mapA = AffineMap::get(3, 0, {rewriter.getAffineDimExpr(0),
						rewriter.getAffineDimExpr(2), rewriter.getContext());
// .. mapB and mapC here
SmallVector<AffineMap> maps = {mapA, mapB, mapC};

// Iteration types
SmallVector<utils::IteratorType> iterTypes = {
	utils::IteratorType::parallel,	// m
	utils::IteratorType:: parallel, 	// n
	utils::IteratorType::reduction	// k
};

Ultimately, transposing the RHS matrix structurally changes memory access patterns under the hood. By ensuring the innermost loop walks sequentially across the rows of both tensors simultaneously, hardware cache locality is maximized, eliminating expensive memory jumps.

To see the impact of this optimization, look at the execution results below for a 2048 x 2048 matmul before and after applying the transpose. The full performance analysis will be discussed in the next article.

Naive IR:

Transposed IR:

Custom TilingPass

The next optimization I did was implementing a custom tiling and fusion pass to target cache eviction and memory round-trips. (Note that I constantly update the codebase and the implementation of this pass shown in this article may not be the most up-to-date, if you want to see this exact implementation, you may need to go through the commit history)

The pipeline before tiling and fusion

Without tiling, a matmul loop nest is forced to process an entire dimension from start to finish before moving to the next. For a large matrix (in this example, is 2048x2048), the processor walks all the way to the end of a massive 2048-element row. Because processor caches are incredibly small, loading the data for the end of that row completely evicts the data from the beginning of the row. When the loop steps down to start the subsequent row, the hardware cannot reuse any of the data it just had; it has to stall and fetch it all over again from slow main RAM

Furthermore, without fusion, the pipeline completes the entire matrix multiplication, writes all 4 million elements out to main memory, and then immediately boots up a second loop nest to read that data right back in just to apply the activation layer (like an Add or ReLU)

The pipeline after tiling and fusion

With the custom tiling and fusion pass, instead of sweeping across the entire matrix width and causing mass cache evictions, the iteration space is constrained to tight, hardware friendly 64 x 64 tiles. The execution pipeline processes a local block of data, moving from left to right, completing a horizontal band before stepping down to the next block.

The fusion engine then pushes the matmul directly inside these local blocks alongside the activation layer. This structure ensures that as soon as a tiny 64 x 64 patch of data is calculated, the activation layer processes it immediately while it is still warm in the L1/L2 cache. The intermediate values are consumed in registers, completely bypassing the need to write millions of elements out to main RAM.

The tiling and fusion pass in this project improved the performance with a 33% speedup compared to a Transpose optimization alone. By stopping the loop from overrunning the cache and forcing evictions, the DLmr (Last-Level Data Cache Misses) dropped from 5.1M down to 4.6M.

With C++ pseudocode, I'm performing the following optimization:

// Before Tiling + Fusing
// A: MxK, B: KxN, B_T: NxK, C: MxN

// Calculate transposed matmul
for (int i = 0; i < M; ++i) {
	for (int j = 0; j < N; ++j) {
		for (int k = 0; k < K; ++k) {
			C[i][j] += A[i][k] * B_t[j][k];
		}
	}
}

// Consumer operation (2 loops, element-wise add or relu)
for (int i = 0; i < M; ++i) {
	for (int j = 0; j < N; ++j) {
		Out[i][j] = C[i][j] + bias[i][j];	// or relu(C[i][j])

// After Tiling + Fusing
// Tiling M and N by a tile size of x (let's say 64)
for (int i_outer = 0; i_outer < M; i_outer += 64) {
	for (int j_outer = 0; j_outer < N; j_outer += 64) {

		// instead of calculating ALL of C first, the matmul producer
		// is fused directly inside this 64x64 block
		for (int k_outer = 0; k_outer < K; k_outer += 64) {

			// innermost execution loop (operating on small L1/L2 cache blocks)
			for (int i = i_outer; i < std::min(i_outer + 64, M); ++i) {
				for (int j = j_outer; j < std::min(j_outer + 64, N); ++j) {
					for (int k = k_outer; k < std::min(k_outer + 74, K); ++k) {
						C[i][j] += A[i][k] B_t[j][k];
					}
				}
			}
		}

		// As soon as the tile of C is finished with all K reductions,
		// the element-wise addition runs immediately while data is still in cache
		for (int i = i_outer; i < std::min(i_outer + 64, M); ++i) {
			for (int j = j_outer; j < std::min(j_outer + 64, N); ++j) {
				Out[i][j] = C[i][j] + bias[i][j];
			}
		}
	}
}

The tiling + fusion pass start by looking at the very end of the execution pipeline and works backwards. If MatMul operation were tiled first, it would be very hard to figure out how to handle the activation layer downstream. By targeting the 2D consumer (Bias/Activation) first, the compiler knows exactly what the final output memory footprint looks like.

On a high-level, the pass isolates the 2D consumer block, ignoring the 3D matmul for the initial phase, and wraps the 2D consumer in 64 x 64 loops. MLIR then inspects the input operands of that consumer. Recognizing that the input is fed by the output of the MatMul, the compiler goes upstream to intercept the matmul, dragging it directly inside those newly created 64 x 64 spatial loops.

Once the MatMul is pulled inside, the pass immediately intercepts it again to tile its remaining K-reduction dimension into 64-element chunks. Finally, the framework rewires the data graph to route all downstream operations to these newly optimized loop outputs, completely replaced the original operations.

To write the MLIR custom tiling + fusion pass, I start by initializing a small vector of targets with Tilinginterface. Interfaces[1] provide a generic way of interacting with the IR. The goal is to be able to express transformations/analyses in terms of these interfaces without encoding specific knowledge about the exact operation or dialect involved. Without interfaces, the passes would need separate logic for every operation or dialect involved e.g., separate logic for linalg.matmul, linalg.add, etc. An OpInterface acts as a C++ compile-time contract (implemented using the CRTP[2]). If a newly invented operation implements the interface, all existing transformation passes can automatically optimize it.

The intent of the TilingInterface[3] is to separate the generation of the loop structure (and constructs used for it) from the information needed from the operation to be able to tile them. As a result, an implementation of the tiling algorithm can generate the inter-tile loop structure, and call into the methods of the interface to be able to tile any operation that implements the interface.

Instead of hardcoding loop-slicing logic for individual matrix or tensor variants, the custom pass targets TilingInterface. This allows the pass to eventually generate structural outer loops by querying the operation itself to determine exactly how to segment its specific iteration space.

In this initial search phase, the module walks through the IR and traverses every single operation in the IR module. By passing the TilingInterface op as the argument to the lambda function, the walker will skip completely over low-level, non-tileable operations (like constants or shape casts) and only execute the inner lambda block if the current operation implements the TilingInterface.

When it meets a generic 'Operation*' pointer, it performs an MLIR-style dynamic cast to check if that operation is explicitly a 'linalg.generic' operation. If it's not (e.g., if it's a structural linalg.transpose), the cast fails, returns a nullptr, and skips it. This ensures that structured, loop-based compute blocks are isolated.

The next filter queries the linalg.generic operation to see how many iteration loops it represents. A standard matmul op requires 3 loops (M, N, K), whereas an element-wise activation layer or bias addition operating on the 2D output matrix requires only 2 loops (M, N). This filters for exactly 2 loops to bypass heavy matmul op itself and find the 2D consumer operation that takes the matmul's output and processes it

// tensor-compiler/src/mlir/TilingPass.cpp

auto module = getOperation();
IRRewriter rewriter(&getContext());

SmallVector<TilingInterface> targets;
module.walk([&](TilingInterface op) {
	// Find the relu/add (2d consumer)
	if (auto genericOp = dyn_cast<linalg::GenericOp>(op.getOperation())) {
		if (genericOp.getNumLoops() == 2) targets.push_back(op);
	}
});

The next step is to loop through the isolated 2D consumer anchors (targets) and configure the outer iteration space (M x N) of these consumers into 64 x 64 tiles. Once that is done, the pass intercepts upstream producers and approve fusion for everything except the transpose operation, as only the transposed RHS matrix data is needed to run the matmul for tiling and fuse purposes (not the layout-shifting operation itself).

// tensor-compiler/src/mlir/TilingPass.cpp

for (TilingInterface consumer : targets) {
	scf::SCFTileAndFuseOptions fuseOptions;
	scf::SCFTilingOptions tilingOptions;

	// Tile MxN
	tilingOptions.setTileSizes({rewriter.getIndexAttr(64),
						rewriter.getIndexAttr(64)});

	fuseOptions.setFusionControlFn([](tensor::ExtractSliceOp candidateSliceOp,
							OpResult originalProducer,
							bool isDestination)
					-> std::optional<scf::SCFTileAndFuseOptions::ControlFnResult> {
		if (isa<linalg::TransposeOp>(originalProducer.getOwner()))
			return std::nullopt;
		return scf::SCFTileAndFuseOptions::ControlFnResult{};
	});

Invoking the scf::tileConsumerAndFuseProducersUsingSCF applies this configuration to materialize the outer spatial loops, clone the consumer block, and drag the upstream MatMul operation inside. Finally, the pass inspects the return receipt to locate the freshly cloned MatMul and immediately tiles its remaining K-reduction dimension into 64-element chunks, automatically replacing the original unoptimized operations with the newly generated loop nest.

// tensor-compiler/src/mlir/TilingPass.cpp

fuseOptions.setTilingOptions(tilingOptions);
rewriter.setInsertionPoint(consumer);

auto result = scf::tileConsumerAndFuseProducersUsingSCF(rewriter, consumer, fuseOptions);

if (succeeded(result)) {
	// Find the fused MatMul inside the new loops and tile K
	for (Operation *tiledOp : result->tiledAndFusedOps) {
		// Cast the interface required by the tiling utility
		auto tilingInterfaceOp = dyn_cast<TilingInterface>(tiledOp);

		if (auto gen = dyn_cast<linalg::GenericOp>(tiledOp)) {
			if (gen.getNumLoops() == 3) {
				scf::SCFTilingOptions kOptions;
				kOptions.setTileSizes({rewriter.getIndexAttr(0),
								rewriter.getIndexAttr(0),
								rewriter.getIndexAttr(64)});

				rewriter.setInsertionPoint(tiledOp);
				// scf::tileUsingSCF only accepts a TilingInterface op
				auto kResult = scf::tileUsingSCF(rewriter, tilingInterfaceOp,
										kOptions);

				if (succeeded(kResult)) {
					rewriter.replaceOp(tiledOp, kResult->replacements);
				}}}}

// Route downstream operations to the new loop nest outputs
for (const auto &it : result->replacements) {
	rewriter.replaceAllUsesWith(it.first, it.second);
}
rewriter.eraseOp(consumer);

IR After Tiled + Fused:

Conclusion

In this phase of the project, I expanded the compiler pipeline to support linear layers by adding matmul and relu operations to the custom dialect. To improve the performance of the compiler, I integrated a RHS matrix transposition directly into the IR loweing pipeline using custom affine indexing maps and permutation arrays to align memory access with strict row-major physical layout.

I also created a custom tiling and fusion pass with MLIR's TilingInterface. This pass isolates the 2D activation consumer, uses SCF's tileConsumerAndFuseProducers to create 64 x 64 spatial loops while pulling the upstream matrix multiplication inside, and applies tileUsingSCF to segment the remaining K-reduction loop. This was done to keep data in cache registers and stop round-trips to the main RAM.

The next phase focuses on verifying these optimization and evaluating efficiency through performance analysis with Valgrind and testing methods.

References

[1] MLIR Doc: Interfaces

[2] Wikipedia: Curiously Recurring Template Pattern

[3] MLIR Codebase: TilingInterface