MLIR module generated by torch_xla fails to compile

[aboubezari]

Hey all, I'd like to be able to convert a torch model to MLIR/HLO, then compile it using XLA's internal compiler. Here's what I did:

First, I created a very simple torch model and generated MLIR/HLO according to this guide.. This is my code:

class SimpleLinearModel(nn.Module):
    def __init__(self):
        super(SimpleLinearModel, self).__init__()
        self.a = nn.Parameter(torch.tensor(1.0)) 
        self.b = nn.Parameter(torch.tensor(0.0))  

    def forward(self, x):
        return self.a * x + self.b

model = SimpleLinearModel()
sample_input = (torch.randn(4, 1), )

exported = export(model, sample_input)
stablehlo_program = exported_program_to_stablehlo(exported)

with open("module.txt", "w") as f:
    f.write(stablehlo_program.get_stablehlo_text('forward'))
with open("module.mlir", "wb") as f:
    f.write(stablehlo_program.get_stablehlo_bytecode('forward'))

Then, I built & ran the xla_compile binary and tried to compile the module.

xla_compile --platform=gpu --module_file=module.mlir --output_file=output --result_output_file=/tmp/result_output

I got this failure message:

2024-12-18 19:01:41.055629: E xla/service/xla_compile_main.cc:115] Compilation failed: UNKNOWN: <unknown>:0: error: conversion requires module with `main` function
<unknown>:0: note: see current operation:
"builtin.module"() <{sym_name = "IrToHlo.11"}> ({
  "vhlo.func_v1"() <{arg_attrs = #vhlo.array_v1<[]>, function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<4x1x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x1x!vhlo.f32_v1>>>, res_attrs = #vhlo.array_v1<[]>, sym_name = #vhlo.string_v1<"main">, sym_visibility = #vhlo.string_v1<"">}> ({
  ^bb0(%arg0: !vhlo.tensor_v1<!vhlo.f32_v1>, %arg1: !vhlo.tensor_v1<4x1x!vhlo.f32_v1>, %arg2: !vhlo.tensor_v1<!vhlo.f32_v1>):
    %0 = "vhlo.broadcast_in_dim_v1"(%arg2) <{broadcast_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x1x!vhlo.f32_v1>
    %1 = "vhlo.multiply_v1"(%0, %arg1) : (!vhlo.tensor_v1<4x1x!vhlo.f32_v1>, !vhlo.tensor_v1<4x1x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x1x!vhlo.f32_v1>
    %2 = "vhlo.broadcast_in_dim_v1"(%arg0) <{broadcast_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x1x!vhlo.f32_v1>
    %3 = "vhlo.add_v1"(%1, %2) : (!vhlo.tensor_v1<4x1x!vhlo.f32_v1>, !vhlo.tensor_v1<4x1x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x1x!vhlo.f32_v1>
    "vhlo.return_v1"(%3) : (!vhlo.tensor_v1<4x1x!vhlo.f32_v1>) -> ()
  }) : () -> ()
}) {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} : () -> ()

Any ideas or suggestions on what went wrong? Thanks!

[NaiyerRizz]

Hi @aboubezari

I tried reproducing the issue here is a github gist for the same:
https://colab.sandbox.google.com/gist/NaiyerRizz/205aa3b6802d1cb8448c5d70e748e3f9/issue_20696.ipynb

I converted that module.txt(stableHlo) to HLO using the below cmd:

bazel run xla/hlo/translate:xla-translate -- --stablehlo-to-hlo-text < module.mlir
(your module.txt file with mlir extension)
Above would generate the HLO output something like:

**HloModule IrToHlo.11, entry_computation_layout={(f32[], f32[4,1]{1,0}, f32[])->f32[4,1]{1,0}}

ENTRY %main.8 (Arg_0.1: f32[], Arg_1.2: f32[4,1], Arg_2.3: f32[]) -> f32[4,1] {
%Arg_2.3 = f32[] parameter(2)
%broadcast.4 = f32[4,1] broadcast(f32[] %Arg_2.3), dimensions={}, metadata={source_file="" source_line=3}
%Arg_1.2 = f32[4,1] parameter(1)
%multiply.5 = f32[4,1] multiply(f32[4,1] %broadcast.4, f32[4,1] %Arg_1.2), metadata={source_file="" source_line=4}
%Arg_0.1 = f32[] parameter(0)
%broadcast.6 = f32[4,1] broadcast(f32[] %Arg_0.1), dimensions={}, metadata={source_file="" source_line=5}
ROOT %add.7 = f32[4,1] add(f32[4,1] %multiply.5, f32[4,1] %broadcast.6), metadata={source_file="" source_line=6}
}**
then I just replaced IrToHlo.11 with %main.8 and then exectued, It worked for me.

Thanks

[Zentrik]

In module.mlir you probably need to rename the function to main