# 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
