bfloat16 lowering from torch MLIR to LinAlg fails for a small representative MLP model

[kumardeepakamd]

Save following as model.py

module {
  func.func @main_graph(%arg0: !torch.vtensor<[8,3],bf16>) -> !torch.vtensor<[8,5],bf16> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.3.0"} {
    %0 = torch.vtensor.literal(dense_resource<_layers.0.weight> : tensor<4x3xbf16>) : !torch.vtensor<[4,3],bf16>
    %1 = torch.vtensor.literal(dense_resource<_layers.0.bias> : tensor<4xbf16>) : !torch.vtensor<[4],bf16>
    %2 = torch.vtensor.literal(dense_resource<_layers.2.weight> : tensor<5x4xbf16>) : !torch.vtensor<[5,4],bf16>
    %3 = torch.vtensor.literal(dense_resource<_layers.2.bias> : tensor<5xbf16>) : !torch.vtensor<[5],bf16>
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %4 = torch.aten.transpose.int %0, %int0, %int1 : !torch.vtensor<[4,3],bf16>, !torch.int, !torch.int -> !torch.vtensor<[3,4],bf16>
    %5 = torch.aten.mm %arg0, %4 : !torch.vtensor<[8,3],bf16>, !torch.vtensor<[3,4],bf16> -> !torch.vtensor<[8,4],bf16>
    %6 = torch.aten.add.Tensor %5, %1, %int1 : !torch.vtensor<[8,4],bf16>, !torch.vtensor<[4],bf16>, !torch.int -> !torch.vtensor<[8,4],bf16>
    %7 = torch.aten.relu %6 : !torch.vtensor<[8,4],bf16> -> !torch.vtensor<[8,4],bf16>
    %int0_0 = torch.constant.int 0
    %int1_1 = torch.constant.int 1
    %8 = torch.aten.transpose.int %2, %int0_0, %int1_1 : !torch.vtensor<[5,4],bf16>, !torch.int, !torch.int -> !torch.vtensor<[4,5],bf16>
    %9 = torch.aten.mm %7, %8 : !torch.vtensor<[8,4],bf16>, !torch.vtensor<[4,5],bf16> -> !torch.vtensor<[8,5],bf16>
    %10 = torch.aten.add.Tensor %9, %3, %int1_1 : !torch.vtensor<[8,5],bf16>, !torch.vtensor<[5],bf16>, !torch.int -> !torch.vtensor<[8,5],bf16>
    %11 = torch.aten.relu %10 : !torch.vtensor<[8,5],bf16> -> !torch.vtensor<[8,5],bf16>
    return %11 : !torch.vtensor<[8,5],bf16>
  }
}

{-#
  dialect_resources: {
    builtin: {
      _layers.0.weight: "0x08000000413CBDBE113F033EA53E2ABE063F9BBDA03ECBBE08BE66BE",
      _layers.0.bias: "0x08000000473E943E6EBEAD3E",
      _layers.2.weight: "0x08000000FA3D3CBD38BC3FBDA0BE89BD76BEF0BEED3C40BE683EFBBCF63E9BBE423E8E3EC33EF3BE203D9D3E",
      _layers.2.bias: "0x0800000044BE34BE363ECFBE923D"
    }
  }
#-}

Run:

<your torch MLIR build>/bin/torch-mlir-opt -convert-torch-onnx-to-torch mlp.bf16.torch-onnx.mlir > mlp.bf16.onnx.torch.mlir 
<your iree build>/tools/iree-compile --iree-hal-target-backends=llvm-cpu mlp.bf16.onnx.torch.mlir > mlp.bf16.vfmb

you will see an error:
:0: error: 'arith.constant' op failed to verify that all of {value, result} have same type
:0: note: see current operation: %3 = "arith.constant"() <{value = dense_resource<_layers.0.bias> : tensor<4xbf16>}> : () -> tensor<4xf32>
mlp.bf16.onnx.torch.mlir:16:10: error: failed to run translation of source executable to target executable for backend #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "x86_64-unknown-unknown-eabi-elf", ukernels = "default"}>
%9 = torch.aten.mm %7, %8 : !torch.vtensor<[8,4],bf16>, !torch.vtensor<[4,5],bf16> -> !torch.vtensor<[8,5],bf16>
^
mlp.bf16.onnx.torch.mlir:2:3: note: called from
func.func @main_graph(%arg0: !torch.vtensor<[8,3],bf16>) -> !torch.vtensor<[8,5],bf16> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.3.0"} {
^
mlp.bf16.onnx.torch.mlir:16:10: note: see current operation: