Simple [0,1,2,3]*2 model using iree

mul_2.py

import torch
import torch.nn as nn
import os
import numpy as np
import shark_turbine.aot as aot

torch.manual_seed(42)
torch.cuda.manual_seed(42)
np.random.seed(42)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor):
        x = x*2
        return x

model = MLP()
example_x = torch.empty(4, dtype=torch.float32)
exported = aot.export(model, example_x)
exported.print_readable()
out_dir = "notes/mlir/mul_2"
os.makedirs(out_dir, exist_ok=True)
exported.save_mlir(f"{out_dir}/mul_2.mlir")
compiled_binary = exported.compile(save_to=None)

def infer(): 
    import iree.runtime as rt

    config = rt.Config("local-task")
    vmm = rt.load_vm_module(
        rt.VmModule.wrap_buffer(
            config.vm_instance, compiled_binary.map_memory()),
        config,
    )
    x = np.arange(4).astype(np.float32) # definitive inputs
    y = vmm.main(x)
    np.save(f"{out_dir}/input.npy", x)
    np.save(f"{out_dir}/output.npy", y.to_host())
    
    x_torch = torch.from_numpy(x)
    y_torch = model(x_torch)
    torch.testing.assert_close(y.to_host(), y_torch.detach().numpy())
    print("y_torch = ", y_torch)
    print("y_iree_cpu = ", y.to_host())
    if os.path.exists(f"{out_dir}/output_cuda.npy"):
        with open(f'{out_dir}/output_cuda.npy', 'rb') as f:
            y_iree_cuda = np.load(f)
        print("y_iree_cuda = ", y_iree_cuda)