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.