Summary of problems during iteratively executing TorchSimplificationPipeline

[Vremold]

Recently, I'm working on iteratively executing TorchSimplificationPipeline in LowerBackendContract pass. The motivation of doing so mainly comes from two aspects:

1. SatifyBackendContract may needs multiple execution of TorchSimplificationPipeline. This is what it now looks like. But in my experiments, TorchSimplificationPipeline usually run only once, for the reason that SatisfyBackendContract has been met.
2. We need to run the pipeline multiple times to make sure the code is optimized into a converged state. An example of suboptimization is in the PR Iteratively run SimplificationPipeline until code optimization converges #1261. These optimizations are usually related to shape inference, and the main passes involved are canonicalizer, shape-simiplification-pipeline and refine-rypes.

However, if executing TorchSimplificationPipeline more than one times. The following problems will arise.

1. Dtype mismatch problems in memory-alloc ops. In the first iteration, we get the output tensor's type. If it mismatches the dtype param of the op (for example, dtype is torch.none), the refine-types of the next iteration will fail. An example can be seen in PR Fix a bug about torch-refine-types pass when the dtype of output tensor is known #1280. Thanks a lot for the suggestions of @ramiro050. It indeed solves part of the problems. However, the following e2e tests still fail, namely RandLikeModule_basic，FullLikeModuleInt2DStatic_basic，FullLikeModuleInt2D_basic and FullModuleInt3D_basic. That's because that the dtype of initial captured aten ops is torch.none. Take RandLikeModule_basic as example, the initial IR is:

%1 = torch.aten.rand_like %arg1, %none_0, %none_0, %none_0, %none_0, %none_0 : !torch.tensor, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.tensor

2. In the decompose-complex-ops pass, AtenRollOp will be split into several AtenSliceTensorOps and AtenCatOp. Here is an example IR:

%3 = torch.aten.slice.Tensor %arg0, %int0, %2, %none, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.none, !torch.int -> !torch.vtensor<[?,1,2],f32>
%4 = torch.aten.slice.Tensor %arg0, %int0, %int0_0, %2, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?,1,2],f32>
%5 = torch.prim.ListConstruct %3, %4 : (!torch.vtensor<[?,1,2],f32>, !torch.vtensor<[?,1,2],f32>) -> !torch.list<vtensor<[?,1,2],f32>>
%6 = torch.aten.cat %5, %int0 : !torch.list<vtensor<[?,1,2],f32>>, !torch.int -> !torch.vtensor<[3,1,2],f32>

If run TorchSimplificationPipeline second time, %3 and %4 will be inferred to have static shape and the verification of ListConstructOp will fail, because it finds out that its two operands have different shape.

%6 = "torch.aten.slice.Tensor"(%arg0, %4, %2, %3, %1) : (!torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.none, !torch.int) -> !torch.vtensor<[2,1,2],f32>
...
%7 = "torch.aten.slice.Tensor"(%arg0, %4, %4, %2, %1) : (!torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.vtensor<[1,1,2],f32>

All the problems are found in current e2e tests. Any idea to solve the above problems? At my point of view, the most straight solution is reorder and reorganization of the passes in TorchSimplificationPipeline

[ramiro050]

Thanks @Vremold for the nice summary of the issues! Iteratively running TorchSimplificationPipeline is uncovering some interesting bugs.

For 1., the issue with aten.rand_like is that one of the ops it gets decomposed into is aten.empty_like, which has a decomposition that is not entirely correct. Namely, the decomposition for aten.emtpy_like should check if op.dtype() is None, and if so, it should pass the dtype of the input tensor, rather than the op.dtype() argument, to the AtenMemoryFormatOp:

torch-mlir/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

Line 1705 in e52e886

op, op.getType(), sizeList, op.dtype(), op.layout(), op.device(),

The fix should be similar to the one you made for the createInitTensor helper function.

[silvasean]

For 2. shape inference should never create invalid IR. If it does then that is a bug. Can you make a reproducer? (a .mlir file and a torch-mlir-opt command line with the shape inference pass that creates the illegal op). I'm happy to take a look at this.

Reorganization of TorchSimplificationPipeline is not a good solution. It can only hide the problems. It does not fix them. It will always be possible to write a new program that hits these issues. These are all legitimate bugs in the relevant passes and should be individually fixed.

[Vremold]

For 2. shape inference should never create invalid IR. If it does then that is a bug. Can you make a reproducer? (a .mlir file and a torch-mlir-opt command line with the shape inference pass that creates the illegal op). I'm happy to take a look at this.

Reorganization of TorchSimplificationPipeline is not a good solution. It can only hide the problems. It does not fix them. It will always be possible to write a new program that hits these issues. These are all legitimate bugs in the relevant passes and should be individually fixed.

The original python script comes from RollModule_basic e2e test case. It looks like this:

class RollModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x.roll([2, -1], [0, 2])
x = torch.rand(3, 1, 2)
module = torch_mlir.compile(RollModule(), x, output_type=torch_mlir.OutputType.TORCH)

I recorded the log in the process.
IR before the first decompose-complex-ops:

// To reproduce this error, run the following commands sequentially:
// torch-mlir-opt -torch-decompose-complex-ops this.mlir
// torch-mlir-opt -pass-pipeline='torch-shape-refinement-pipeline' this.mlir
func.func @forward(%arg0: !torch.vtensor<[3,1,2],f32>) -> !torch.vtensor<[3,1,2],f32> {
  %int0 = torch.constant.int 0
  %int2 = torch.constant.int 2
  %int-1 = torch.constant.int -1
  %0 = torch.prim.ListConstruct %int2, %int-1 : (!torch.int, !torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int0, %int2 : (!torch.int, !torch.int) -> !torch.list<int>
  %2 = torch.aten.roll %arg0, %0, %1 : !torch.vtensor<[3,1,2],f32>, !torch.list<int>, !torch.list<int> -> !torch.vtensor<[3,1,2],f32>
  return %2 : !torch.vtensor<[3,1,2],f32>
}

IR after the first decompose-complex-ops:

...
%2 = torch.aten.neg.int %int2 : !torch.int -> !torch.int
%3 = torch.aten.slice.Tensor %arg0, %int0, %2, %none, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.none, !torch.int -> !torch.vtensor<[?,1,2],f32>
%4 = torch.aten.slice.Tensor %arg0, %int0, %int0_0, %2, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?,1,2],f32>
%5 = torch.prim.ListConstruct %3, %4 : (!torch.vtensor<[?,1,2],f32>, !torch.vtensor<[?,1,2],f32>) -> !torch.list<vtensor<[?,1,2],f32>>
%6 = torch.aten.cat %5, %int0 : !torch.list<vtensor<[?,1,2],f32>>, !torch.int -> !torch.vtensor<[3,1,2],f32>
%7 = torch.aten.neg.int %int-1 : !torch.int -> !torch.int
%8 = torch.aten.slice.Tensor %6, %int2, %7, %none, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.none, !torch.int -> !torch.vtensor<[3,1,?],f32>
%9 = torch.aten.slice.Tensor %6, %int2, %int0_0, %7, %int1 : !torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[3,1,?],f32>
%10 = torch.prim.ListConstruct %8, %9 : (!torch.vtensor<[3,1,?],f32>, !torch.vtensor<[3,1,?],f32>) -> !torch.list<vtensor<[3,1,?],f32>>
%11 = torch.aten.cat %10, %int2 : !torch.list<vtensor<[3,1,?],f32>>, !torch.int -> !torch.vtensor<[3,1,2],f32>
...

IR after the second torch-simplify-shape-calculations, this is where it fails:

%6 = "torch.shape.calculate"() ({
  %14 = "torch.aten.slice.Tensor"(%arg0, %4, %2, %3, %1) : (!torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.none, !torch.int) -> !torch.vtensor<[2,1,2],f32>
  "torch.shape.calculate.yield"(%14) : (!torch.vtensor<[2,1,2],f32>) -> ()
}, {
  %14 = "torch.prim.ListConstruct"(%5, %1, %5) : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
  "torch.shape.calculate.yield.shapes"(%14) : (!torch.list<int>) -> ()
}) : () -> !torch.vtensor<[2,1,2],f32>
%7 = "torch.shape.calculate"() ({
  %14 = "torch.aten.slice.Tensor"(%arg0, %4, %4, %2, %1) : (!torch.vtensor<[3,1,2],f32>, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.vtensor<[1,1,2],f32>
  "torch.shape.calculate.yield"(%14) : (!torch.vtensor<[1,1,2],f32>) -> ()
}, {
  %14 = "torch.prim.ListConstruct"(%1, %1, %5) : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
  "torch.shape.calculate.yield.shapes"(%14) : (!torch.list<int>) -> ()
}) : () -> !torch.vtensor<[1,1,2],f32>
%8 = "torch.prim.ListConstruct"(%6, %7) : (!torch.vtensor<[2,1,2],f32>, !torch.vtensor<[1,1,2],f32>) -> !torch.list<vtensor<[?,1,2],f32>>

The error message is:

raw.mlir:68:10: error: operand types should have the same type as the list contained type
    %3 = torch.aten.roll %arg1, %1, %2 : !torch.tensor, !torch.list<int>, !torch.list<int> -> !torch.tensor
         ^
raw.mlir:68:10: note: see current operation: %8 = "torch.prim.ListConstruct"(%6, %7) : (!torch.vtensor<[2,1,2],f32>, !torch.vtensor<[1,1,2],f32>) -> !torch.list<vtensor<[?,1,2],f32>>

I understand your concerns and I will follow the same line of thinking to deal with these issues later.

[silvasean]

Closing this because the reduced test case passes now:

class RollModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x.roll([2, -1], [0, 2])
x = torch.rand(3, 1, 2)
module = torch_mlir.compile(RollModule(), x, output_type=torch_mlir.OutputType.TORCH)