Pytorch model to mlir -> llvm -> executable file on Mac book m1

# Step 1: Define and train a simple PyTorch CNN model

import torch

import torch.nn as nn

import torch.optim as optim

import torchvision

import torchvision.transforms as transforms

# Define a simple CNN

class SimpleCNN(nn.Module):

def __init__(self):

super(SimpleCNN, self).__init__()

self.conv1 = nn.Conv2d(1, 32, 3, 1)

self.conv2 = nn.Conv2d(32, 64, 3, 1)

self.dropout1 = nn.Dropout2d(0.25)

self.dropout2 = nn.Dropout2d(0.5)

self.fc1 = nn.Linear(9216, 128)

self.fc2 = nn.Linear(128, 10)

def forward(self, x):

x = self.conv1(x)

x = nn.functional.relu(x)

x = self.conv2(x)

x = nn.functional.relu(x)

x = nn.functional.max_pool2d(x, 2)

x = self.dropout1(x)

x = torch.flatten(x, 1)

x = self.fc1(x)

x = nn.functional.relu(x)

x = self.dropout2(x)

x = self.fc2(x)

output = nn.functional.log_softmax(x, dim=1)

return output

# Train the model (simplified for brevity)

model = SimpleCNN()

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters())

# Assume we've trained the model...

# Save the trained model

torch.save(model.state_dict(), "simple_cnn.pth")

# Step 2: Compile the model with torch-mlir

import torch_mlir

# Load the trained model

model = SimpleCNN()

model.load_state_dict(torch.load("simple_cnn.pth"))

model.eval()

# Create an example input tensor

example_input = torch.randn(1, 1, 28, 28)

# Compile the model to MLIR

mlir_module = torch_mlir.compile(model, example_input, output_type="linalg-on-tensors")

# Save the MLIR module to a file

with open("simple_cnn.mlir", "w") as f:

f.write(str(mlir_module))

# Step 3: Lower MLIR to LLVM IR

# This step typically requires using the MLIR tools from the command line

# mlir-opt simple_cnn.mlir --convert-linalg-to-loops --convert-scf-to-cf --convert-vector-to-llvm --convert-memref-to-llvm --convert-func-to-llvm --reconcile-unrealized-casts | mlir-translate --mlir-to-llvmir > simple_cnn.ll

# Step 4: Compile LLVM IR to machine code

# Use Clang to compile for M1 Mac (arm64 architecture)

# clang -O3 -march=arm64 simple_cnn.ll -o simple_cnn_exec

# The result is an executable file named 'simple_cnn_exec'

# Step 5 (optional): Create a C++ wrapper to use the compiled model

#include <iostream>

#include <vector>

// Declare the function generated from our PyTorch model

extern "C" void simple_cnn(float* input, float* output);

int main() {

// Prepare input (28x28 image flattened to 1D array)

std::vector<float> input(784, 0.0f); // Initialize with zeros for simplicity

// Prepare output (10 classes for MNIST)

std::vector<float> output(10, 0.0f);

// Call the compiled model

simple_cnn(input.data(), output.data());

// Print the output (class probabilities)

for (int i = 0; i < 10; ++i) {

std::cout << "Class " << i << " probability: " << output[i] << std::endl;

}

return 0;

}

# Compile the C++ wrapper with the compiled model

# clang++ -O3 wrapper.cpp simple_cnn_exec -o final_executable