Skip to content

[TorchToLinalg] Failed to legalize operation 'torch.aten.index.Tensor_hacked_twin'. #4495

@nikhil-m-k

Description

@nikhil-m-k

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

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>
  }
}

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions