Failed to legalize operation 'torch.aten.batch_norm' that was explicitly marked illegal

[jxhekang]

Got the Exception: Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR failed with the following diagnostics: error: failed to legalize operation 'torch.aten.batch_norm' that was explicitly marked illegal note: see current operation: %6 = "torch.aten.batch_norm"(%arg0, %1, %0, %0, %1, %2, %3, %4, %5) : (!torch.vtensor<[1,32],f32>, !torch.vtensor<[32],f32>, !torch.vtensor<[32],f32>, !torch.vtensor<[32],f32>, !torch.vtensor<[32],f32>, !torch.bool, !torch.float, !torch.float, !torch.bool) -> !torch.vtensor<[1,32],f32>

By running the code below:

import torch
from torch import nn
import torch_mlir

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.bn = nn.BatchNorm1d(32)
    
    def forward(self, x):
        out = self.bn(x)
        return out


def main():
    model = SimpleModel().eval()

    example_input = torch.rand((1, 32), dtype=torch.float32)
    print(model(example_input))

    traced = torch.jit.trace(model, example_input)
    linalg_on_tensors_mlir = torch_mlir.compile(traced, example_input,
                                                output_type=torch_mlir.OutputType.LINALG_ON_TENSORS)
    print(linalg_on_tensors_mlir)


if __name__ == "__main__":
    main()

[ramiro050]

Hi @jxhekang,

Please take a look at https://github.com/llvm/torch-mlir/issues/1042, and let me know if that fixes it.

[jxhekang]

Hi @jxhekang,

Please take a look at #1042, and let me know if that fixes it.

Yes, it fixes my bug! I pull the latest torch-mlir code, and test the nn.BatchNorm1d op, no execption any more. Thanks!

[silvasean]

Fixed by https://github.com/llvm/torch-mlir/issues/1042