Building a Minimal MLIR Pipeline from ONNX to LLVM JIT

14 hours ago
28 min read

MLIR Lowering Pipeline

Lowering Strategy (High-Level Overview)

The lowering pipeline looks like this:

ONNX graph -> custom mlir dialect -> tensor (value) -> memref (memory) -> llvm -> JIT

Full code for this article is available here

This article's focus is only building an end-to-end skeleton that takes a high-level ONNX model and lowers it through progressively more concrete representations until it executes via a JIT engine. I kept the initial scope to scalar and 1D vector models to focus on understanding the underlying workflow of the MLIR lowering infrastructure

The following is the summary for each process:

ONNX Frontend Loading

I start by loading the ONNX model through the frontend (which I've labeled as "parsing" in my project, though "loading" more accurately describes the implementation). Instead of traditional parsing, the logic here manually maps ONNX operations into an internal representation

Custom MLIR Dialect and IR Generation

To define custom operations without drowning in C++ boilerplate, I used ODS (Operation Definition Specification)[1] via Ops.td. This is the core of the dialect; it's where I formally define the operations and their constraints. For this stage, I'm sticking to the fundamentals, which is defining the basic operation set before expanding into more complex traits later

To support the TableGen definitions, I implemented a minimal 'Dialect.cpp'. While this file is where I eventually store my custom builders, my focus at this stage is get a basic verifier running.

Finally, I moved to 'mlirGen.cpp' to handle the actual generation logic. This pass transforms the ONNX data into the custom dialect.

Lowering to Tensor

Next, I lower my custom dialect into a combination of the Linalg and Arith dialects. I chose to lower to the tensor level first because I plan to implement high-level optimizations here in the next stage. I use Linalg for the structured math (like generic map operations) and Arith to handle things like constant definitions. Once the conversion is defined, I wrap it in a Custom Pass and register it in the driver

Bufferization

For the transition from tensors to physical memory, MLIR's One-Shot Bufferizer does the heavy lifting. In my driver.cpp pipeline, I coordinate this by first converting empty tensors into explicit allocations and then applying the bufferization pass.

Under the hood, this pass performs a comprehensive analysis to determine where buffers can be reused safely and where new allocations are required. After this, I applied a deallocation simplification pass to make sure that the resulting memrefs are properly managed and memory-safe for the final JIT execution

LLVM Dialect and JIT Execution

To get the code ready for execution, I lower the memref and scf dialects down to the LLVM dialect. I used Conversion Patterns here. Since I am already at the memory-buffer stage, I lower structured loops into control flow (blocks and branches) because that is the language LLVM speaks

The final stage is the LLVM JIT. This was a major learning point for me regarding the MemRef Descriptor. When passing data from C++ to the JIT-compiled function, the code has to account for the fact that MLIR represents a memref as a specific struct.

Code in this post has been shortened for brevity, for full code, please refer to the repo linked at the beginning of this article. The branch "repo-setup" contains everything needed for this end-to-end skeleton.

ONNX Frontend with MLIR

During this first stage, I want to get a structured form of the ONNX model, the high-level workflow looks like this:

load onnx model -> Set the structure for the model in ModelInfo.h -> parse the model proto and save the information to ModelInfo.h -> dump the proto

Load ONNX Model

The function opens the file in binary mode and uses 'ParseFromIstream' to populate a 'ModelProto' object (from the google protobuf library). This process converts the serialized data into a structured in-memory representation. This allows for verifying graph integrity (such as node count) before initiating the lowering passes to custom MLIR dialect.

// tensor-compiler/src/driver.cpp

static std::unique_ptr<onnx::ModelProto> loadONNXModel(llvm::StringRef filename) {
	// Open the onnx file as a binary stream
	std::ifstream input(filename.str(), std::ios::binary);
	if (!input) { // error handling }

	auto model = std::make_unique<onnx::ModelProto>();

	// Parse the serialized data into the ModelProto structure
	if (!model->ParseFromIstream(&input)) {
		// error handling
	};
	return model;
}

Move loaded data into a structured Model Info

Once the model is loaded, I move the data into a custom 'ModelInfo' structure. While it might seem redundant to copy data from the Protobuf object, this structured approach decouples the compiler from the raw ONNX schema

It ensures that every operation has immediate access to the specific attribute it needs (like shapes and element types) without the overhead of navigating the deep Protobuf tree during every lowering pass.

The extraction follows a top-down hierarchical traversal of the model schema (model -> graph -> nodes -> attributes -> extract data (tensors and initializers)):

Model -> Graph: I start at the top level to grab the IR version and producer metadata
Graph -> Nodes: I iterate through the graph to extract each node's operation type, inputs, and outputs
Nodes -> Attributes: For each node, I parse specific attributes and map them to a custom 'AttributeInfo' struct
Data Extraction: At the end, I extract the tensors and initializers.

Moving the data into the ModelInfo structure decouples the logic from ONNX's internal complexities. It provides a straightforward way to dump the model's information and ensure everything is correct before starting the MLIR lowering.

ModelInfo.h contains the structure of the ONNX model while Parser.cpp converts the model from the ModelProto structure into the structure provided by ModelInfo.h

// tensor-compiler/src/include/dlc/ModelInfo.h code snippet 
// (start at the "leaf" and then move up to the root)

struct TensorInfo {
	std::string name;
	enum class DataType { FLOAT, INT64, INT32, BOOL, UINT8, UNKNOWN };
	std::vector<int64_t> shape;
	DataType elementType;
	std::vector<char> rawData;
	bool isScalar() const { ... };
};

// ValueInfo -> AttributeInfo -> NodeInfo -> GraphInfo

struct ModelInfo {
	int64_t ir_version;
	std::string producer_name;
	GraphInfo graph;
};

// tensor-compiler/src/parser/Parser.cpp code snippet 
// (This follows the same path as ModelInfo.h, but when running the code, it will start
// from the root (graph) and then moves down the leaves (tensors)
// The main purpose is to save the structured ONNX model's information
// into ModelInfo's structure for easy debugging
// and provide stable source of truth for MLIR conversion
static TensorInfo parseTensor(const onnx::TensorProto &t) {
	TensorInfo info;

	// get tensor's name and shape
	info.name = t.name();
	for (int64_t d : t.dims()) {
		info.shape.push_back(d);
	}

	// Map ONNX data type to Enum data type as structured in the ModelInfo
	// ONNX uses int constants (e.g., 1 for FLOAT, 7 for INT64, etc.)
	// this converts those into a more readable C++ type
	switch(t.data_type()) {
		case 1: info.elementType = TensorInfo::DataType::FLOAT; break;
		// the rest of the elementTye -> data type here including default
	}

	// Get raw data
	return info;
}

// parseValueInfo -> parseAttribute -> parseNode -> parseGraph -> parseModelProto

Print the structured model

Once the model's information is saved in the form of ModelInfo, I want to check to make sure everything is working well so I printed them out on the CLI. The code is available on driver.cpp following the same top-down hierarchical traversal like the ModelInfo and Parser.cpp. An example of the dumpPROTO (the printer) for an add operation:

Custom MLIR Dialect and IR generation

Define Dialect Operations with ODS

To generate MLIR, I use ODS (Operation Definition Specification). Operations defined in a TableGen (.td) file generates necessary C++ classes with minimal boilerplate.[2] This ensures a single source of truth for the dialect's structure.

The process begins by defining the dialect and a base operation class. For each specific op, you define the mnemonic (its IR name), traits (behaviors), arguments (inputs/attributes), and results (output types). Additional fields, such as verifiers or custom builders, can be included to handle validation and object construction. MLIR then automatically generates the underlying C++ code for operation handling, attribute accessors, and trait enforcement.

// Example Ops.td
// tensor-compiler/src/include/dlc/Ops.td code snippet
// Start with the definition of the dialect, here, I named my dialect dlc
def Dlc_dialect : Dialect {
	let name = "dlc";
	let summary = "A high-level dialect for onnx IR";
	let cppNamespace = "::mlir::dlc";
}

// Base class for the dialect, this op inherits from the base 'Op' class in Opbase.td
// It takes in a mnemonic (name of your op), and a list of traits (behaviors)
// in this example for Constant op, a "Pure" trait means it has no side effects
class Dlc_Op<string mnemonic, list<Trait> traits = []> :
	Op<Dlc_Dialect, mnemonic, traits>;

// Constant op
def ConsantOp : DlcOp<"constant", [Pure]> {
	let summary = "constant";
	// ... description here

	// argument are the inputs to your op
	// can be operands or attributes
	let arguments = (ins ElementsAttr:$value);

	// TableGen requires things to be explicit
	// the result here can only be a Tensor of data type Float32
	let results = (outsF32Tensor);

	// Tells TableGen that I'll be writing a verifier for this op
	// could be to check shapes, datatypes, etc.
	let hasVerifier = 1;

	// this tells the TableGen how to create a constructor for this op
	// in this case, when given an input of DenseElementsAttr, build an op that
	// returns that datat's type and store the data in the $value attribute.
	let builders = [
		OpBuilder<(ins "DenseElementsAttr":$value), [{
			build($_builder, $_state, value.getType(), value);
		)]>
	];
]

Note: In my current implementation, I strictly define the result type as Float32 to keep the lowering passes predictable. But production-grade compilers like onnx-mlir[3] (Ops.td auto-generated via onnx-mlir/utils/gen_onnx_mlir.py) handle type diversity using UnionType via 'AnyTypeOf'.[4] To act as a universal container, their ConstantOp defines multiple OptionalAttr fields. This allows the same operation to hold various data formats (e.g., ints, floats, etc.) while remaining explicitly typed for the rest of the MLIR pipeline

You can see exactly what MLIR TableGen tool builds by running it manually. This command transforms the declarative .td file into actual C++ source code. Specifically, the -gen-op-defs flag invokes the OpDefinitionsGen[5] backend within the MLIR project to generate the operation definitions (the .cpp.inc file). This file contains the C++ logic for the constructors, accessors, and verifiers defined in ODS

To function correctly, the command uses the -I flag to include the standard MLIR directory, allowing the tool to locate base classes like Op or Dialect.

// Go to llvm-project/build directory and run this on your CLI
// change everything inside <> and remove the <>

./bin/mlir-tblgen -gen-op-defs </path/to/your/ops/td/file> -I </path/to/your>/llvm-project/mlir/include

// example

./bin/mlir-tblgen -gen-op-defs /workspace/tensor-compiler/src/include/dlc/Ops.td -I /workspace/llvm-project/mlir/include

Register dialect and provide logic for custom fields

Once the operations are defined in ODS, the project requires a C++ implementation to register the dialect and provide the logic for custom fields. This is handled in Dialect.cpp

the initialize() method acts as the entry point, where addOperations registers the generated C++ classes with the MLIR system. This ensures that the compiler recognizes the dialect's operations during the lowering passes. Following the registration, the file includes the implementation for manual logic that were declared in the .td file (in this example, I added a custom verifier.

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

// register all operations with the MLIR context
void mlir::dlc::DlcDialect::initialize() {
	addOperations<
#define GET_OP_LIST
#include "dlc/Ops.cpp.inc"
	>();
)

// Pulls in the generated class definitions
#define GET_OP_CLASSES
#include "dlc/Ops.cpp.inc"

// Custom verifier for the ConstantOp
// This implements the 'hasVerifier = 1' promise from Ops.td
LogicalResult ConstantOp::verify() {
	// Get the attribute value and the result type
	auto tensorAttr = llvm::dyn_cast<DenseElementsAttr>(getValue());
	auto tensorType = llvm::cast<RankedTensorType>(getResult().getType());

	// Check if the num of elements matches
	if (tensorAttr.getNumElements() != tensorType.getNumElements()) {
		// error message
	}
	return success();
}

Generate MLIR

The final step in the frontend is the MLIR Generation phase. This process takes the structured ModelInfo above and translates it into a format MLIR understands. To begin this conversion, three core components are required:

MLIRContext[6]: The central registry. It manages the global state, including loaded dialects, type uniquing, and operation metadata
ModuleOp[7]: A top-level container that holds all the functions, constants, and global variables for the model
OpBuilder[8]: The builder is the primary interface for IR construction. It maintains an "insertion point," so when calling 'builder.create<...>', it knows exactly where to place that instruction in the code

// tensor-compiler/src/mlir/MLIRGen.cpp code snippet

// Helper functions created (to see implementation, go to the full code)
// getMlirType: convert the data type ->  Mlir's type, e.g., FLOAT -> Mlir's Float
// getLoc: get the location (node's name)

// Generate MLIR from the Ops defined above
mlir::OwningOpRef<mlir::ModuleOp> mlirGen (mlir::MLIRContext &context,
								::dlc::ModelInfo &model) {
	// Register dialect
	context.getOrLoadDialect<DlcDialect>();
	context.getOrLoadDialect<mlir::func::FuncDialect>();

	// Create module
	auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));

	// Builder
	mlir::OpBuilder builder(&context);

Preparing to walk the graph

Before translating the model logic, I need to establish the function's entry point and signature. This involves the following steps:

Iterate through the ModelInfo inputs to convert them into MLIR RankedTensorType objects. Storing these in argTypes allows the builder to define the formal parameters for the function
I then create a func.func operation named "main". This acts as the primary container for the model's logic. At this stage, the return types are left empty, as they will be populated later during the lowering process
By setting the _mlir_ciface_main attribute, I tell the compiler to generate a wrapper. This is crucial for hardware execution because it creates a stable C-compatible entry point, simplifying how external data is passed to the compiled code
I then add an entry block the function and move the OpBuilder's insertion point to the start of block. The function is then pushed into the ModuleOp container

Note on the mlir_ciface_main wrapper:

During the initial phase of the project, my JIT engine failed to execute the main function despite the IR looking correct.
This is due to MLIR's internal representation of tensors as it uses complex descriptors that aren't natively compatible with the standard C calling convention used by my runner
By adding the _mlir_ciface_main attribute, the compiler is instructed to generate a C-compatible wrapper. This acts as the essential bridge, 'unwrapping' the complex descriptors so the JIT engine can successfully pass data into the compiled hardware logic

// Prepare the function signature (input arguments)
// Converts the custom ModelInfo inputs into a list of MLIR RankedTensorTypes
llvm::SmallVector<mlir::Type, 4> argTypes;
for (const auto &inputInfo : model.graph.inputs) {
	auto elementType = getMlirType(builder, inputInfo.elementType);
	argTypes.push_back(mlir::RankedTensorType::get(inputInfo.shape, elementType));
}

// Create main function with the func dialect (container for the whole model)
auto func = func::FuncOp::create(
	builder.getUnknownLoc(),
	"main",
	builder.getFunctionType(argTypes, {}) // signature: (inputs) -> (outputs)
);

// Tell the compiler to generate the _mlir_ciface_main wrapper
func->setAttr(mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(), builder.getUnitAttr());

Block *entry = func.addEntryBlock();
builder.setInsertionPointToStart(entry);
module.push_back(func);

mlir::ImplicitLocBuilder b(builder.getUnknownLoc(), builder);

To bridge the gap between ModelInfo's name-based data flow and MLIR's value-based system, I initialize a valueMap. I map each input name to its corresponding Block Argument.

In MLIR, Block Arguments are the actual entry points for data within a block; they act as the live wires that carry the input tensors into the function. This mapping creates a source of truth: as the graph is traversed, any operation requiring a specific ModelInfo input can simply look up the correct mlir::Value handle to use as an operand

Since the MLIRContext created these mlir::Value wires when the entry block was added earlier (based on the argTypes provided), the loop below simply maps the ModelInfo names to those existing wires using the same index i to ensure the correct name matches the correct input port

// Initialize value table: ModelInfo's name -> MLIR value
llvm::StringMap<mlir::Value> valueMap;

for (size_t i = 0; i < model.graph.inputs.size(); ++i) {
	const auto &inputInfo = model.graph.inputs[i];
	valueMap[inputInfo.name] = entry->getArgument(i);
}

Walking the Nodes

With the entry point ready, the generator iterates through each node in the graph. For every node, a symbolic lookup resolves its inputs into MLIR operands using a two-tier check:

ValueMap Check: If an input name is already in the valueMap, the data is "live" (it is either a global input or the result of a previous layer). The existing MLIR handle is simply added to the operands list
Initializer Check: If the name is missing from the map, the compiler checks the initializers (where weights and constants reside). When a match is found, a ConstantOp is emitted using the stored shape, type, and raw buffer. This new result is saved back into the valueMap to prevent redundant operations if the same weight is referenced again.

Once operands are resolved, the builder emits the actual math operation. For an Add node, the generator takes the two resolved operands (LHS and RHS) and creates the MLIR AddOp. The resulting value is then stored in the valueMap under the node's output name, allowing subsequent layers to consume it.

For example, let's say we have the following operation:

(X + W1) + W2

The first add Node (add_1)

inputs: ["X", "W1"]
Logic: The code looks up "X" and "W1", puts them in "operands"
Result: "operands[0]" is "%X", "operands[1]" is "%W1"
Action: "b.create<AddOp>(%X, %W1)" creates "%sum1"
Record: "valueMap["sum1_out"] = %sum1"

Second add node (add_2)

inputs: ["sum1_out", "W2"]
Logic: the loop starts over. "operands" is now empty
Lookup: It looks up "sum1_out" (found in 'valueMap') and "W2" (found in initializers)
Result: Now, "operands[0]" is "%sum1" and 'operands[1]' is "%W2"
Action: "b.create<AddOp>(%sum1, %W2)" creates the final result

// Walk nodes
for (const NodeInfo &node : model.graph.nodes) {
	mlir::Location loc = getLoc(node, builder);
	b.setLoc(loc);

	llvm::SmallVector<mlir::Value, 2> operands;
	for (const std::string &inName : node.inputs) {
		if (valueMap.count(inName)) {
			operands.push_back(valueMap[inName]);
		} else if (model.graph.initializers.count(inName)) {
			const auto &tensor = model.graph.initializers[inName];
			auto type = RankedTensorType::get(tensor.shape,
								getMlirType(builder, tensor.elementType));

			auto denseAttr = DenseElementsAttr::getFromRawBuffer(
				type, llvm::ArrayRef<char>(tensor.rawData.data(),
									tensor.rawData.size())
			);

			auto ConstOp = b.create<ConstantOp>(type, denseAttr);
			valueMap[inName] = constOp.getResult();
			operands.push_back(constOp.getResult());
		}
	}
	if (node.op_type == "Add") {
		valueMap[node.outputs[0]] = b.create<AddOp>(operands[0], operands[1]);
	}
}

After iterating through all nodes, the generator handles the function's return values by looping through the model's outputs. For each output name, a lookup is performed in the valueMap to find the corresponding MLIR handle

If the value is present, it is added to the returnValues and returnTypes lists
If a required output is missing from the map, it indicates a break in the graph's data flow, requiring error handling

Finally, a func::ReturnOp is emitted using these values, and the function's signature is updated to reflect the finalized return types.

// Handle return values after the loop is done
SmallVector<mlir::Value, 4> returnValues;
SmallVector<mlir::Type, 4> returnTypes;

for (const auto &outputInfo : model.graph.outputs) {
	if (valueMap.count(outputInfo.name) {
		mlir::Value val = valueMap[outputInfo.name];
		returnValues.push_back(val);
		returnTypes.push_back(val.getType);
	} else {
		// error handling here
	}
}
b.create<func::ReturnOp>(returnValues));
func.setType(builder.getFunctionType(argTypes, returnTypes));

The final step is a structural audit to ensure the generated IR adheres to all dialect specific constraints, e.g., matching tensor shapes for arithmetic operations. Once the ModuleOp passes this integrity check, it is returned for further lowering or optimization.

if (failed(mlir::verify(module))) {
	// error message here
	return nullptr;
}
return module;

Production Grade Frontends

In this project, I chose to manually map ONNX operations to a custom dialect to gain a deep understanding of attribute-to-operand translation. While this manual approach is invaluable for learning, it becomes a significant bottleneck for production-grade compilers that must support hundreds of evolving operators across the entire stack

Instead of manual mapping, production-grade compilers (like onnx-mlir) use a more automated frontend strategy.[9] They use scripts to parse official ONNX operator schemas and automatically generate MLIR TableGen (.td) definitions. This automation ensures the dialect remains synchronized with the latest ONNX versions and eliminates the repetitive boilerplate code for every new operation

Another example is torch-mlir[10], which takes a similar approach. Since PyTorch is defined by its C++ registry, torch-mlir runs a script that queries the active PyTorch JIT operator registry to see every available function and then generates the MLIR ODS.

Furthermore, the field seems to be increasingly exploring AOT (Ahead-of-Time) capture methods. Rather than just scraping a registry, FxImporter[11] uses TorchDynamo and TorchFX to intercept Python bytecode and capture a normalized graph before lowering.

For generating MLIR, my implementation above manually walked the ONNX graph and emit MLIR operations through an if-else dispatch logic. Production grade compilers like onnx-mlir adopt a more declarative, table-driven architecture.[12][13]

Instead of manually writing emission logic for each operator, onnx-mlir uses a generated dispatch table (OpBuildTable.inc). This table acts as a bridge between ONNX strings and C++ templates. When the compiler encounters an ONNX string like "Gemm" or "Conv", it performs a constant-time map lookup to trigger a template-instantiated handler (buildOperation<T>)

By offloading the heavy lifting of attribute parsing, variadic operand management, and result type inference to this automated pipeline, the production method ensures the compiler remains maintainable even as the ONNX specification evolves. My approach simplifies these complexities to focus on the core mechanics of the lowering pipeline.

DUMP MLIR

To check the dlc dialect and print out the dlc IR: load the model proto, run mlirGen and dump the IR to console

// driver.cpp

// Parse the internal model representation
dlc::ModelInfo modelInfo = dlc::parseModelProto(model);

// Generate the MLIR module
auto module = dlc::mlirGen(context, modelInfo);

// Dump to console
module->print(llvm::outs());

Lowering To Tensor

After generating the dlc dialect, the next step involves lowering it to standard MLIR dialects. For this phase, I target the arith, tensor, and linalg dialects to move the IR closer to executable code.

Constant Lowering

Lowering dlc.constant to arith.constant is a direct transformation. Both operations use a DenseElementsAttr to store underlying data. DenseElementsAttr is a specific MLIR attribute designed to store multi-dimensional data in contiguous, dense memory block. Because both dialects share this storage format, I can simply move the attribute from the source to the destination operation

To handle this, I use an 'OpConversionPattern'.[14] This class is the standard choice for dialect conversion because it is aware of type changes and integrates with the 'ConversionTarget' to make sure that the IR is legalized properly. Unlike a greedy rewriter that loops until nothing changes, this pattern-driven approach is goal oriented: it specifically looks for illegal ops (in this case, dlc.constant) and uses the rewriter to swap it for a legal op (in this case, arith.constant)

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

struct ConstantOpLowering : public OpConversionPattern<dlc::ConstantOp> {
	using OpConversionPattern<dlc::ConstantOp>::OpConversionPattern;

	LogicalResult matchAndRewrite(dlc::ConstantOp op, 	// takes in original dlc op
							OpAdaptor adaptor, 	// provides converted operands
							ConversionPatternRewriter &rewriter) const final {

		// Create the new arith.constant operation
		// Pull the location and the actual tensor data (op.getValue())
		// directly from the original dlc.constant
		auto newOp = arith::ConstantOp::create(rewriter, op.getLoc(), op.getValue());

		// Replace the old op
		// The rewriter handles the heavy lifting of updating the IR to point to the
		// new operation
		rewriter.replaceOp(op, newOp);
		return success();
	}
};

Add Operation

Lowering dlc.add is more involved because linalg operation on tensors follow Destination Passing Style (DPS).[15]

Unlike standard arithmetic where an operation implicitly returns a new value, DPS requires an explicit "destination" tensor to hold the result. This design enables the compiler to perform advanced memory optimizations and buffer reuse later in the pipeline. My implementation handles this by:

Creating an uninitialized tensor via tensor.empty based on the result shape
Providing this tensor as the output operand to linalg.add

auto loc = op.getLoc();
auto resultType = cast<RankedTensorType>(op.getType());

// Create the destination tensor
auto dest = tensor::EmptyOp::create(
				rewriter,
				loc,
				resultType.getShape(),
				resultType.getElementType());

auto addOp = linalg::AddOp::create(
		rewriter,
		loc,
		TypeRange{resultType},					// result types
		ValueRange{adaptor.getLhs(), adaptor.getRhs()},	// ins
		ValueRange{dest}						//outs
);
rewriter.replaceOp(op, addOp->getResults());

Finally, I wrap these patterns into a 'ConversionTarget'. By marking standard dialect as Legal and the dlc dialect as illegal, I tell the rewriter to exhaustively apply these patterns until every custom operation is lowered. To execute this, I add the custom pass to a 'PassManager' alongside a canonicalization pass to clean up the resulting IR

// loader: ensures dialects are initialized in the MLIR Context
void getDependentDialects(DialectRegistry &registry) const override {
	registry.insert<arith::ArithDialect, linalg::LinalgDialect,
				tensor::TensorDialect, func::FuncDialect>();
}

void runOnOperation() final {
	ConversionTarget target(getContext());

	// gatekeepers: defines which dialects this pass is permitted to output
	target.addLegalDialect<arith::ArithDialect, linalg::LinalgDialect,
					tensor::TensorDialect, func::FuncDialect>();

	// Mark dlc dialect is illegal
	target.addIllegalDialect<dlc::DlcDialect>();

	RewritePatternSet patterns(&getContext());
	patterns.add<ConstantOpLowering, AddOpLowering>(&getContext());

	// apply partial conversion
	if (failed(applyPartialConversion(getOperation(), target, std::move(patterns)))) {
		signalPassFailure();
	}
};

Finally, add the custom pass to the pass manager to lower to tensor:

pm.addPass(mlir::dlc::createLowerToTensorPass());

Bufferization

After tensor, the next step is to lower to MemRef, which is when the IR stops being a math equation and starts looking like actual execution. In MLIR, Tensors are immutable, they don't exist in memory, they are just values. But hardware needs Memrefs as they are essentially the pointers to physical memory buffers.

MLIR has done most of the work for lowering to Memref in this project with the One-Shot Bufferize pass. It performs a comprehensive analysis of Use-Def chains to identify data hazards. It treats bufferization as a resource allocation problem: if a tensor's buffer can be safely overwritten without violating dependencies (no Read-after-Write conflicts), the pass marks the operation as in-place, mapping multiple tensor values to a single underlying memref. If a conflict is detected, it materializes a new allocation to preserve data integrity.[16]

To transition from tensors to physical memory, I implemented a bufferization pipeline through the following stages:

Preparation: I used the EmptyTensorToAllocTensor pass to convert tensor.empty ops into alloc_tensor ops. This aligns with the Destination-Passing Style, essentially telling the compiler exactly where memory might need to be carved out.
Analysis: The One-Shot bufferize pass performs a top-down analysis to map these tensors to buffers. By enabling bufferizeFunctionBoundaries, it ensures the final function signatures use pointers (memref) rather than astract tensors, making the compiled kernel callable from a C++ driver
Memory Lifecycle management: Finally, since memref represents raw memory, the lifecycle of these pointers need to be handled. The BufferDeallocationSimplification pass automatically inserts deallocation instructions at the moment a buffer is no longer needed, preventing memory leaks in the final execution

// tensor-compiler/src/driver.cpp

pm.addPass(mlir::bufferization::createEmptyTensorToAllocTensorPass());
mlir::bufferization::OneShotBufferizePassOptions passOptions;

passOptions.allowReturnAllocFromLoops = true;
passOptions.bufferizeFunctionBoundaries = true;

pm.addPass(mlir::bufferization::createOneShotBufferizePass(passOptions));
pm.addPass(mlir::bufferization::createBufferDeallocationSimplificationPass());

// CLEAN UP

Clean Up

Once everything is in memory, linalg is lowered to standard loops.Canonicalizer and CSE (Common Subexpression Elimination) is immediately run after because Bufferization often generates a lot of redundant pointer arithmetic and dead code, so these passes sweep that up before hitting the LLVM backend.

At this point, the compiler part is mostly done. The output is a structured loop nest operating on raw memory, the next step is to get it ready for execution.

LLVM and JIT Execution

Lower To LLVM Dialect

To get the IR ready for execution on a CPU, the high-level operations are first lowered to the LLVM Dialect within MLIR. This process begins by declaring the necessary dependent dialects and defining an LLVMConversionTarget. The ModuleOp is marked as a legal operation because a full conversion is being performed. If the top-level ModuleOp container were not marked legal, the conversion would fail, as the engine would find no valid pattern to transform the module structure itself.

The LLVMTypeConverter[17] handles the complex tasks of translating high-level types into the LLVM type system. Under the hood, it acts as a dictionary and set of rules for translating high-level types into LLVM's type system. For example, when converting an index type (commonly used for loop iterations), the converter queries the hardware's data layout to determine whether to lower it to an i32 or i64 integer

Finally, collect all the conversion patterns for the dialects that have been used so far, in this case: Arith, MemRef, SCF, Func, and lower them into a combination of LLVM and Control Flow IR. As a final cleanup step, the ReconcileUnrealizedCasts pass is executed to strip away any redundant cast pairs that were generated during the lowering process.

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

LLVMConversionTarget target(getContext());
target.addLegalOp<mlir::ModuleOp>();

// Define the Type Converter (How do types change)
LLVMTypeConverter typeConverter(&getContext());

// Collect all the patterns
RewritePatternSet patterns(&getContext());

// Patterns for the dialects used
populateSCFToControlFlowConversionPatterns(patterns);
// the rest of the conversion pattern here

// Apply the conversion to the module
auto module = getOperation();
applyFullConversion(module, target, move(patterns)

Lower to LLVM

The next step is to lower to LLVM before compiling the code, MLIR has made it easy to do that with the registerBuiltInDialectTranslation and registerLLVMDialectTranslation, which provide the compiler with the logic required to map the LLVM Dialect in MLIR to actual LLVM IR

It then converts the MLIR Module that was passed into LLVM IR with the translateModuleToLLVMIR function. Under the hood, it walks the MLIR tree and for every op it sees, it uses the LLVM C++ API to create the LLVM op of the MLIR op.

The next part is to understand the targets, specifically for data layouts, for example, it needs to know how big a pointer is or how to align a double for the specific hardware. The workflow initializes the specific backend for the host machine (e.g., x86 or ARM), and queries the OS to determine the target triple[18] (e.g., x86_64-apple-darwin23.4.0). This information is used by createTargetMachine() to generate a configuration object that encapsulates the specific features of the CPU (including the number of available registers and specialized instruction sets)

This hardware model is critical for backend algorithms like register allocation and instruction selection. Finally, setupTargetTripleAndDataLayout acts as the final preparation step before execution. It attaches essential hardware metadata (the target triple and data layout) to the LLVM module. This ensures that the abstract instructions are calibrated to the physical constrains of the host CPU, such as pointer bit-width and memory alignment rules

// tensor-compiler/src/driver.cpp

static int dumpLLVMIR(mlir::ModuleOp module) {
	// Register translations to LLVM IR
	mlir::registerBuiltinDialectTranslation(*module->getContext());
	mlir::registerLLVMDialectTranslation(*module->getContext());

	// Convert MLIR Module to LLVM IR module
	llvm::LLVMContext llvmContext;
	auto llvmModule = mlir::translateModuleToLLVMIR(module, llvmContext);

	// Initialize LLVM targets
	llvm::InitializeNativeTarget();
	llvm::InitializeNativeTargetAsmPrinter();

	// Set up data layout and target triple
	auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
	auto tmOrError = tmBuilderOrError->createTargetMachine();
	mlir::ExecutionEngine::setTargetTripleAndDataLayout(
		llvmModule.get(), tmOrError.get().get()
	);

	// output the result
}

Run JIT

The final step is to execute the LLVM IR that was just generated. I used a JIT (Just-In-Time) runner, which compiles the generated LLVM IR into machine instructions at runtime.

The following explanation is for running JIT for dynamic input type. Initially, I only accept 2 arguments to get started (in repo-setup branch) and later accommodated dynamic input types (in main branch)

To support dynamic input types, the runner requires specific metadata: the ranks and shapes of the model's inputs and outputs. This metadata must be captured while the IR is still at a high level of abstraction. In this pipeline, I collect the shape and rank information before lowering to the Tensor/Linalg dialects.

This collection must occur before bufferization, this is because bufferization converts abstract tensors into concrete memory buffers. Once bufferization occurs, the explicit shape of the operation is lowered into the MemRef descriptor's metadata, making it much harder to extract.

// Find the main function created by mlirGen (still at dlc dialect at this point)
auto mainFunc = module->lookupSymbol<mlir::func::FuncOp>("main");

// Get input ranks and shapes
std::vector<int64_t> inRanks;
std::vector<std::vector<int64_t>> inShapes;
for (auto type : mainFunc.getArgumentTypes()) {
	auto tensorType = llvm::cast<mlir::RankedTensorType>(type);
	inRanks.push_back(tensorType.getRank());
	inShapes.push_back(tensorType.getShape().vec();
}

// .. Get output rank and shape (similar to input Args but get result types instead of
// args types instead)

// Prepare data (explained below)
auto inputDataStore = prepareInputData(inShapes, inputDataList, inputFiles);

processMLIR(context, *module);
return runJIT(*module, outRank, inRanks, outShape, inShapes, inputDataStore);

Prepare Input Data

To prepare the input data, first check to make sure that the input shapes collected from the high-level IR are correct compared to the CLI data's size (collected through the inputDataList from the command line). The next part is to pass the data into the input data store. It accepts data from the CLI directly, a binary, or random data (for testing).

To prepare the input data to accept files, the code creates an Input File Stream (ifs)--essentially a data pipe--to read from binary files on disk. It calculates the total byte size required by multiplying the number of elements by the size of a float (4 bytes). It then performs a binary read, which treats the floating-point data as raw bytes just long enough to copy them directly from the disk into the data vector in RAM.

The driver also accepts data via the CLI. If the provided data doesn't match the expected input size, the code zeros out the missing elements. While a production version would return a hard error to ensure total accuracy, this padding was used during development to make it easier to test MLIR passes without being blocked by strict input requirements.

Finally, there's a random number generator if nothing is given, which is mainly for testing purposes such as basic execution or performance benchmarking.

// prepare inputDataStore
static std::vector<std::vector<float>> prepareInputData(
	const std::vector<std::vector<int64_t>>& inShapes,
	const llvm::cl::list<float>& cliData,
	const llvm::cl::list<std::string>& cliFiles) {
	// .. Check to make sure that the input == total expected

	std::vector<std::vector<float>> inputDataStore;
	size_t fileIdx = 0;
	size_t cliOffset = 0;

	// Fixed seed for random generator
	std::mt19937 gen(42);
	std::uniform_real_distribution<float>dist(0.0f, 1.0f);

	for (const auto& shape : inShapes) {
		int64_t numElements = 1;
		for (int64_t dim : shape) numElements *= dim;
		std::vector<float> data(numElements);

		if (fileIdx < cliFiles.size()) {
			// Load from binary file if reading from files
			std::ifstream ifs(cliFiles[fileIdx++], std::ios::binary);
			ifs.read(reinterpret_cast<char*>(data.data()), numElements * sizeof(float));
		} else if (cliOffset < cliData.size()) {
			// Load from CLI directly
			for (int64_t j = 0; j < numElements; ++j) {
				if (cliOffset < cliData.size()) {
					data[j] = cliData[cliOffset++];
				} else {
					data[j] = 0.0f;
				}
			}
		} else {
			// Random data
			for (int64_t j = 0; j < numElements; ++j) {
				data[j] = dist(gen);
			}
		}
		inputDataStore.push_back(std::move(data));
	}
	return inputDataStore;
}

Engine Setup

The JIT execution begins with environment setup, where the CPU is prepared to interpret the specific instructions generated by the compiler. Once it's ready, the process moves to engine creation by building an MLIR Execution Engine. For the execution to actually process the information, raw C++ arrays are converted into a specific descriptor format (MemRef Descriptor) that MLIR understands. Finally, the system runs and perform cleanup, where the operations are executed, timed and cleaned up.

The MLIR engine knows how to lower to LLVM as it does the heavy lifting of checking the target machine for us (as opposed to the method shown during lowering to LLVM above where there's some steps explicitly mentioned such as checking target triples, I created specific lower to LLVM step above because I wanted to show the IR directly which could be useful when needing to debug)

To set up the engine, I first initialized the native LLVM targets and register the necessary dialect translations so MLIR knows how to talk to the host CPU. Set the optimization level and then construct the ExecutionEngine which is what handles the JIT compilation of the module into machine code.

Under the hood, the Execution engine[19] acts as an orchestrator for the LLVM ORC (On-Request Compilation)[20] JIT stack. It begins by translating the MLIR module into LLVM IR and then use the LLJIT builder to construct a layered compilation pipeline. This pipeline consists of a compile layer, which transforms IR into machine-code object files, and an object linking layer, which manages executable memory allocation.

To bridge the gap between the C++ code and the compiled JIT code, the engine performs Argument Packing. Instead of the programmer having to match the exact hardware requirements of every function, it creates a simplified wrapper interface by wrapping the function a type-erased interface that always looks the same: a single list of generic pointers (void**)

Finally, it resolves external symbols (such as runtime print utilities) by dynamically searching the provided shared library paths, effectively linking the JITed binary to its environment at runtime.

I then added the following libraries to the Engine:

libmlir_runner: Descriptors are complex for dynamic shapes (varying ranks, strides, and sizes). This library contains the logic to make sure that even if the shape changes at runtime, the JITed code correctly understands the memory layout provided by the code.
libmlir_c_runner: This library contains the implementation of common runtime operations

// Initialize LLVM targets and register translations to LLVM IR
llvm::InitializeNativeTarget();
llvm::InitializeNativeTargetAsmPrinter();
mlir::registerBuiltInDialectTranslation(*module->getContext());
mlir::registerLLVMDialectTranslation(*module->getContext());

// Create an MLIR ExecutionEngine
auto optPipeline = mlir::makeOptimizingTransformer(3, 0, nullptr);
mlir::ExecutionEngineOptions engineOptions;
engineOptions.transformer = optPipeline;
engineOptions.jitCodeGenOptLevel = llvm::CodeGenOptLevel::Aggressive;
engineOptions.sharedLibPaths = {
	"path/to/libmlir_runner_utils.so",
	// .. other utils needed here
};

auto execEngine = mlir::ExecutionEngine::create(module, engineOptions);
auto &engine = execEngine.get();

MemRef Preparation

After setting up the ExecutionEngine, the next step is to prepare the input and output data for invocation. MLIR represents a tensor as a MemRef Descriptor.[21] It's not just a flat array; it's a struct that looks like this:

Allocated Pointer: The raw memory address
Aligned Pointer: The address used for data access
Offset: Where the data actually starts
Size: The dimensions
Strides: How many steps to take to move to the next row

For example, for a rank r tensor, you need 3 fixed slots (Alloc, Aligned, Offset) plus 2 slots per dimension (Size and Stride)

Because I want to accept inputs with varying shapes and ranks, the structure cannot be hard-coded. Instead, the logic iterates through each input and calculates the required slots in the descriptor array dynamically using the formula as described above: 3 + (2 x ranks)

For every dimension, the code stores the size and calculates the stride. In a contiguous row-major array, the stride serves as the crucial instruction for the JIT, defining exactly how many memory addresses to skip to reach the next element in a specific dimension.

While MLIR defines the descriptor structure, the code is responsible for initializing its metadata. In this implementation, I mapped the allocated and aligned pointers to the same base address and setting the offset to zero because the standard std::vector guarantees natural alignment for its underlying data type on modern 64-bit systems.

While input data is stored in existing std::vector containers, a specific buffer must be allocated for the output. The total element count is derived by multiplying the output dimensions, followed by a malloc call to reserve the necessary space in memory. An output Descriptor is then initialized to point to this outData pointer, giving the JIT the memory address needed to write the computation results.

// Prepare input data
std::vector<std::vector<int64_t>> inputDescriptors;
inputDescriptors.reserve(inRanks.size());

size_t DataOffset = 0;

// Dynamic Shape Handling for inputs
for (size_t i = 0; i < inRanks.size(); ++i) {
	int64_t rank = inRanks[i];
	const auto& shape = inShapes[i];

	// Build the MemRef Descriptor
	size_t slots = 3 + (2 * rank);
	inputDescriptors.emplace_back(slots, 0);
	auto& desc = inputDescriptors.back();

	desc[0] = reinterpret_cast<int64_t>(inputDataStore[i].data());	// Allocated
	desc[1] = desc[0];
	desc[2] = 0;

	// Fill sizes and strides
	for (int64_t j = 0; j < rank; ++j) {
		desc[3 + j] = shape[j];	// size
		int64_t stride = 1;
		for (int64_t k = j + 1; k < rank; ++k) {
			stride *= shape[k];
		}
		desc[3 + rank + j] = stride;
	}
}

// The output descriptor process is very similar to the input, but requires manual
// memory reservation since a destination buffer does not yet exist.
// Also, unlike the multiple input tensors, this implementation only handles a single
// output tensor

// Example of the memory allocation for output data
int64_t totalOutElements = 1;
for (int64_t dim : outShape) totalOutElements *= dim;
float* outData = (float*)malloc(totalOutElements * sizeof(float));

The ExecutionEngine uses a type-erased 'void' interface, which demands a list of memory addresses. This introduces a critical requirement for memory stability:

An args vector holds the addresses of the descriptors
To prevent pointers from becoming dangling due to vector reallocations, descriptors are stored in a way that ensures their memory addresses remain fixed
Storing pointers to each descriptor in a dedicated 'inputPtrs' vector ensures that the addresses passed to invokePacked remain valid until the JIT completes its task.
So instead of passing '&inputDescriptors[i].data()' directly into 'args', I created 'inputPtrs' to create a stable intermediate layer. (While JIT is running, the inputDescriptors may decide to resize itself, in which the address initially given to the JIT would now point to garbage memory)

// Once memory is allocated, pack the inputs and outputs in the 'args' Array
// By MLIR calling convention, output is expected to be the first argument in the packed
// array followed by inputs

// Pack output: Store a pointer to the descriptor data to provide the stable address
// (&outDescPtr) required by the JIT
std::vector<void*> args;
void* outDescPtr = outDesc.data();
args.push_back(&outDescPtr);

// input pointers
// same as output but with vectors (there's multiple inputs and only one output tensor)
std::vector<int64_t*> inputPtrs;
for (auto& desc : inputDescriptors) {
	inputPtrs.push_back(desc.data());
}

for (size_t i = 0; i < inputPtrs.size(); ++i) {
	args.push_back(&inputPtrs[i]);
}

The array of pointers is passed to the _mlir_ciface_main function. The ciface suffix indicates that the JIT-compiled function follows the C-interface convention for receiving descriptors. The engine unpacks these addresses, maps them to the internal MLIR representations, and executes the machine code on the allocated memory. It is the necessary entry point for any function using the invokePacked mechanism with MemRef descriptors. (a generic main function in C++ usually expects standard arguments like 'int argc, char argv'. However, the JIT-compiled function in this project expects a list of MemRef Descriptors)

Finally, call the engine using 'engine->invokePacked', passing the output descriptor pointer followed by the input pointers. This follows the standard MLIR calling convention where the result buffer is passed as the first argument

// invoke
engine->invokePacked("_mlir_ciface_main", args);

After the call, I reconstruct the result by reading from the output descriptor, extract the aligned pointer and offset, then calculate the total number of elements by multiplying the sizes stored in the descriptor before printing the results to the console (and don't forget to free the allocated memory)

// Read and save data
float *allocated = reinterpret_cast<float*>(outDesc[0]);
float *aligned = reinterpret_cast<float*>(outDesc[1]);
std::ofstream ofs("output.bin", std::ios::binary);
if (ofs) {
	// ofs.write accepts a const char* as the first argument
	// treat the float array as a sequence of raw bytes
	ofs.write(reinterpret_cast<char*>(aligned), totalOutElements * sizeof(float));
	ofs.close());
}

Note on packing input and output descriptors:

When packing the desc[i] for input and output, the original data is stored as a float, but notice how it’s casted as int64_t when storing the data in the descriptors, and then casted back as float when reading/saving data. So why not just use float to begin with?
This is because the MemRef descriptor needs to hold two very different types of information:
- pointers: memory addresses of the data (64-bit addresses)
- integers: the sizes, offsets, and strides (64-bit numbers)
In C++, if you want a single array to hold an address and a size, all elements in that array must be the same size. Since the host machine is a 64-bit system, both a pointer and an int64_t are 8 bytes (64 bits). So int64_t was used as a common denominator (as opposed to float which would be 32-bits, which would chop off half the address)
The JIT compiler also expects the descriptor to be a contiguous block of 64-bit slots. When the JIT-compiled code runs, it doesn’t know it’s looking at a C++ vector. It just sees a pointer to a block of memory and starts reading the first 8 bytes as the data address, the next 8 bytes as the aligned address, and so on. (if float was used—which is 32-bits, the addresses would be corrupted)
I could potentially have used void like ‘std::vector<void*>’ but would run into opposite problems when storing sizes and strides
- Since “size” is a number, not an address
- While technically a number can be forced into a pointer type, it makes the code very confusing to read and makes the math for the stride calculations much more difficult because C++ does not allow standard multiplication on ‘void*’ types

// .. print logic for the code here (available on the github codebase linked at the beginning of the article)
// free the memory
if (allocated) free(allocated);

Conclusion

In this phase of the project, I built an end-to-end skeleton for a tensor compiler to lower high-level ONNX models to a JIT execution engine. The pipeline loads serialized ONNX models into a custom internal representation (ModelInfo), generates a custom MLIR dialect (dlc) via ODS and progressively lowers operations into standard dialects like linalg, arith, and tensor.

From there, MLIR's One-Shot Bufferizer manages the transition from tensors to memref physical memory pointers, which are then translated into standard LLVM IR calibrated to the host machine's target triple and data layout. Finally, an ExecutionEngine packing 64-bit descriptors compiles and invokes the executable machine code at runtime.

The next phase focuses entirely on maximizing performance on the matmul workloads by transposing the right-hand side matrix and creating a custom tiling + fusion pass.