module {
func.func @forward(%arg0: !torch.vtensor<[?,32,64,1,64],f32>, %arg1: !torch.vtensor<[?,3,128,256],f32>) -> !torch.vtensor<[?,1,1,9],f32> {
%int0 = torch.constant.int 0
%int32 = torch.constant.int 32
%int64 = torch.constant.int 64
%int1 = torch.constant.int 1
%int2 = torch.constant.int 2
%int4 = torch.constant.int 4
%int3 = torch.constant.int 3
%int768 = torch.constant.int 768
%none = torch.constant.none
%cpu = torch.constant.device "cpu"
%cpu_0 = torch.constant.device "cpu"
%false = torch.constant.bool false
%0 = torch.prim.ListConstruct %int0, %int1, %int2, %int4, %int3 : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten.permute %arg0, %0 : !torch.vtensor<[?,32,64,1,64],f32>, !torch.list<int> -> !torch.vtensor<[?,32,64,64,1],f32>
%2 = torch.prim.ListConstruct %int768, %int32, %int64, %int64 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%3 = torch.aten.view %1, %2 : !torch.vtensor<[?,32,64,64,1],f32>, !torch.list<int> -> !torch.vtensor<[?,32,64,64],f32>
%4 = torch.aten.size.int %arg1, %int0 : !torch.vtensor<[?,3,128,256],f32>, !torch.int -> !torch.int
%5 = torch.aten.mul.int %4, %int32 : !torch.int, !torch.int -> !torch.int
%6 = torch.prim.ListConstruct %5, %int1, %int1, %int32 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%7 = torch.aten.view %3, %6 : !torch.vtensor<[?,32,64,64],f32>, !torch.list<int> -> !torch.vtensor<[?,1,1,64],f32>
%8 = torch.aten.arange.start_step %int0, %5, %int1, %none, %none, %cpu, %false : !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.Device, !torch.bool -> !torch.vtensor<[?],si64>
%9 = torch.aten.view %8, %6 : !torch.vtensor<[?],si64>, !torch.list<int> -> !torch.vtensor<[?,1,1,1],si64>
%10 = torch.aten.arange.start_step %int0, %int1, %int1, %none, %none, %cpu_0, %false : !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.Device, !torch.bool -> !torch.vtensor<[1],si64>
%11 = torch.prim.ListConstruct %int1, %int1, %int1, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%12 = torch.aten.view %10, %11 : !torch.vtensor<[1],si64>, !torch.list<int> -> !torch.vtensor<[1,1,1,1],si64>
%13 = torch.aten.arange.start_step %int0, %int2, %int3, %none, %none, %cpu_0, %false : !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.Device, !torch.bool -> !torch.vtensor<[?],si64>
%14 = torch.prim.ListConstruct %int2, %int2, %int2, %int2 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%15 = torch.aten.view %13, %14 : !torch.vtensor<[?],si64>, !torch.list<int> -> !torch.vtensor<[?,1,1,9],si64>
%16 = torch.prim.ListConstruct %9, %12, %15, %15 : (!torch.vtensor<[?,1,1,1],si64>, !torch.vtensor<[1,1,1,1],si64>, !torch.vtensor<[?,1,1,9],si64>, !torch.vtensor<[?,1,1,9],si64>) -> !torch.list<vtensor>
%17 = torch.aten.index.Tensor_hacked_twin %7, %16 : !torch.vtensor<[?,1,1,64],f32>, !torch.list<vtensor> -> !torch.vtensor<[?,1,1,9],f32>
return %17 : !torch.vtensor<[?,1,1,9],f32>
}
}
Torch-mlir fails with following error when running the below test case:
error: failed to legalize operation 'torch.aten.index.Tensor_hacked_twin' that was explicitly marked illegal %17 = torch.aten.index.Tensor_hacked_twin %7, %16 : !torch.vtensor<[?,1,1,64],f32>, !torch.list<vtensor> -> !torch.vtensor<[?,1,1,9],f32>Attaching the test and the Command to reproduce the failure:
$ torch-mlir-opt -pass-pipeline='builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline{allow-non-finites=True})' test.mlir