IREE test code and explanation

.

from iree import compiler, runtime

import numpy as np

import sys

def print_step(step):

print(f'Step: {step}', file=sys.stderr)

# MLIR code as a string

module_str = '''

func.func @simple_add(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {

%0 = arith.addf %arg0, %arg1 : tensor<4xf32>

return %0 : tensor<4xf32>

}

'''

print_step('Compiling module')

compiled_module = compiler.compile_str(module_str, target_backends=['llvm-cpu'])

print_step('Creating runtime config')

config = runtime.Config('local-task')

print_step('Creating system context')

ctx = runtime.SystemContext(config=config)

print_step('Creating VM instance')

vm_instance = runtime.VmInstance()

print_step('Creating VM module')

vm_module = runtime.VmModule.from_flatbuffer(vm_instance, compiled_module, warn_if_copy=False)

print_step('Adding VM module to context')

ctx.add_vm_module(vm_module)

print_step('Getting device')

device = runtime.get_driver('local-task').create_default_device()

print(f'Device: {device}', file=sys.stderr)

print_step('Getting function')

f = ctx.modules.module.simple_add

print_step('Creating device arrays')

arg1 = runtime.asdevicearray(device, np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32))

arg2 = runtime.asdevicearray(device, np.array([5.0, 6.0, 7.0, 8.0], dtype=np.float32))

print_step('Calling function')

result = f(arg1, arg2)

print_step('Getting result')

print(result.to_host())

print_step('Script completed successfully')

..

To run this code:

1. Save it to a file, e.g., test_iree.py.
2. Make sure you have IREE and its Python bindings installed and properly set up in your environment.
3. Run the script using Python:

   python test_iree.py

This script will:

1. Define a simple MLIR function that adds two 4-element float32 tensors.
2. Compile this MLIR code to an IREE module.
3. Set up the IREE runtime environment.
4. Create input data as NumPy arrays.
5. Execute the compiled function with the input data.
6. Print the result.

The output should show each step of the process and finally print the result, which should be [ 6. 8. 10. 12.].

This example demonstrates the basic workflow for testing MLIR code with IREE using Python. You can modify the MLIR code string and input data to test different functions and operations as needed.

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