[MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided #1742

AmosLewis · 2022-12-21T19:17:57Z

Got this bug in DistilGpt2 to Tosa nod-ai/SHARK-Studio#494

// -----
// CHECK-LABEL:   func.func @torch.aten.as_strided(
// CHECK-SAME:                                     %[[VAL_0:.*]]: !torch.vtensor<[3,3],f32>) -> !torch.vtensor<[2,2],f32> {
// CHECK:           %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[3,3],f32> -> tensor<3x3xf32>
// CHECK:           %[[VAL_2:.*]] = torch.constant.int 2
// CHECK:           %[[VAL_3:.*]] = torch.constant.int 1
// CHECK:           %[[VAL_4:.*]] = torch.constant.int 0
// CHECK:           %[[VAL_5:.*]] = torch.prim.ListConstruct %[[VAL_2]], %[[VAL_2]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK:           %[[VAL_6:.*]] = torch.prim.ListConstruct %[[VAL_3]], %[[VAL_2]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK:           %[[VAL_7:.*]] = "tosa.const"() {value = dense<{{\[\[}}0, 2, 1, 3]]> : tensor<1x4xi32>} : () -> tensor<1x4xi32>
// CHECK:           %[[VAL_8:.*]] = "tosa.reshape"(%[[VAL_1]]) {new_shape = [1, 9, 1]} : (tensor<3x3xf32>) -> tensor<1x9x1xf32>
// CHECK:           %[[VAL_9:.*]] = "tosa.gather"(%[[VAL_8]], %[[VAL_7]]) : (tensor<1x9x1xf32>, tensor<1x4xi32>) -> tensor<1x4x1xf32>
// CHECK:           %[[VAL_10:.*]] = "tosa.reshape"(%[[VAL_9]]) {new_shape = [2, 2]} : (tensor<1x4x1xf32>) -> tensor<2x2xf32>
// CHECK:           %[[VAL_11:.*]] = torch_c.from_builtin_tensor %[[VAL_10]] : tensor<2x2xf32> -> !torch.vtensor<[2,2],f32>
// CHECK:           return %[[VAL_11]] : !torch.vtensor<[2,2],f32>
// CHECK:         }
func.func @torch.aten.as_strided(%arg0: !torch.vtensor<[3,3],f32>) -> !torch.vtensor<[2,2],f32> {
  %int2 = torch.constant.int 2
  %int1 = torch.constant.int 1
  %int0 = torch.constant.int 0
  %0 = torch.prim.ListConstruct %int2, %int2 : (!torch.int, !torch.int) -> !torch.list<int>
  %1 = torch.prim.ListConstruct %int1, %int2 : (!torch.int, !torch.int) -> !torch.list<int>
  %2 = torch.aten.as_strided %arg0, %0, %1, %int0 : !torch.vtensor<[3,3],f32>, !torch.list<int>, !torch.list<int>, !torch.int -> !torch.vtensor<[2,2],f32>
  return %2 : !torch.vtensor<[2,2],f32>
}

AmosLewis · 2022-12-22T00:37:21Z

distillgpt2_torch_delete_decompose_amax_selectint.mlir

Need this refind types patch to fix the decompose amax issue:
#1745

ramiro050

On Github I don't see a lowering for as_strided in this PR. Is there a file missing?

AmosLewis · 2022-12-22T17:30:56Z

On Github I don't see a lowering for as_strided in this PR. Is there a file missing?

I am planning to add lower to tosa in the next PR. This is required for the distilgpt2 model. It might be a long patch. So I just added this e2e first so I can erase the torch.operator.* in mlir file. Before adding lower to tosa, I have to fix the aten.slice.Tensor to tosa first.

ramiro050 · 2022-12-22T17:42:54Z

I am planning to add lower to tosa in the next PR. This is required for the distilgpt2 model.

In general, we should try to avoid adding code that does not get run by the e2e test suite, even if done temporarily. There are several parts in this PR currently don't get executed by the e2e suite.

It might be a long patch.

The changes here are not that much code. They should all be part of a single patch that adds e2e support. Not only does this help avoid having dead code in torch-mlir, but it makes reviewing easier, since the reviewer can see the declaration of the op, as well as the handling of its dtype, shape, and testing.

AmosLewis · 2022-12-22T17:47:56Z

I am planning to add lower to tosa in the next PR. This is required for the distilgpt2 model.

In general, we should try to avoid adding code that does not get run by the e2e test suite, even if done temporarily. There are several parts in this PR currently don't get executed by the e2e suite.

It might be a long patch.

The changes here are not that much code. They should all be part of a single patch that adds e2e support. Not only does this help avoid having dead code in torch-mlir, but it makes reviewing easier, since the reviewer can see the declaration of the op, as well as the handling of its dtype, shape, and testing.

Ok. I will continue iterating on this patch.

lib/Dialect/Torch/Transforms/RefineTypes.cpp

eric-k256

The implementation looks okay to me. Ideally we'd avoid a tosa.gather as it tends to be a slow op for acceleration, but I'm not sure a loop of SLICEs would be significantly better in this case. Looking at the original network, is it doing effectively as_strided(as_strided(as_strided(as_strided(tensor))))?

test/Conversion/TorchToTosa/basic.mlir

lib/Conversion/TorchToTosa/TorchToTosa.cpp

ramiro050 · 2023-01-03T23:12:53Z

python/torch_mlir_e2e_test/test_suite/reshape_like.py

+    ])
+
+    def forward(self, x):
+        return torch.ops.aten.as_strided(x, (2, 2), (1, 2), 1)


Does this implementation work if you do

return torch.ops.aten.as_strided(x.t(), (2, 2), (1, 2), 1)

I.e. pass the transpose of x as the argument

To add to this, I think adding support for this op in torch-mlir will be very tricky. The reason is that this op depends on knowledge about the storage used by the input tensor, and at the torch dialect level in torch-mlir there is no notion of storage. In the example I gave above, x.t() does not change the storage of the tensor, so PyTorch returns the same output as when x is passed. However, I expect torch-mlir will output a tensor as if x.t().contiguous() had been passed instead because it does not know that x.t() and x share the same storage.

Where is it that you're seeing this op used? Given the warning in the documentation of the op, I would expect that this op is not explicitly used in the definition of a model, but rather it is being generated by PyTorch when turning the model to JIT IR. If this is the case, then maybe we can find a way to fold it back.

With x.t(), it fail.

➜ torch-mlir git:(as_stride) ✗ python -m e2e_testing.main -c tosa -f "AsStridedStaticModule_basic" FAIL - "AsStridedStaticModule_basic" Unexpected outcome summary: ****** Failed tests - 1 tests FAIL - "AsStridedStaticModule_basic" Summary: Failed: 1

Where is it that you're seeing this op used?
I got it from transformer distilgpt2 model.
Here is the python patch I use: distillgpt2.py
Here is the torchscript ID I got: distillgpt2_torchscript.mlir
Here is the debug torch mlir I got:distilgpt_lambda.mlir
Here is the final tosa file I generated for distilgpt2:distilgpt2_tosa.mlir

it is being generated by PyTorch when turning the model to JIT IR.
Not sure how to find if it is generated by JIT IR. And furthermore, if it is, how to fix it?

AmosLewis · 2023-01-04T23:43:06Z

The implementation looks okay to me. Ideally we'd avoid a tosa.gather as it tends to be a slow op for acceleration, but I'm not sure a loop of SLICEs would be significantly better in this case. Looking at the original network, is it doing effectively as_strided(as_strided(as_strided(as_strided(tensor))))?

I thought of using tosa::SliceOp but didn't figure out a way. In the generated distillgpt2_torchscript.mlir, I didn't find as_strided(as_strided(as_strided(as_strided(tensor)))), but find the as_strided(view(tensor)), as_strided(view(tensor)), as_strided(view(tensor)). Could you post the link to the code you find it?

eric-k256 · 2023-01-06T00:08:12Z

I was looking at your code, and this looks like a sequence of nested as_strided:

    %395 = torch.prim.ListConstruct %int1048576, %int1048576, %int1024, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %396 = torch.operator "aten.as_strided"(%393, %394, %395, %int0) : (!torch.tensor, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.tensor loc(#loc182)
    %397 = torch.prim.ListConstruct %int1, %int1, %int1024, %int1024 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %398 = torch.prim.ListConstruct %int1048576, %int1048576, %int1024, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %399 = torch.operator "aten.as_strided"(%396, %397, %398, %int0) : (!torch.tensor, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.tensor loc(#loc183)
    %400 = torch.prim.ListConstruct %int1, %int1, %int128, %int1024 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %401 = torch.prim.ListConstruct %int1048576, %int1048576, %int1024, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %402 = torch.operator "aten.as_strided"(%399, %400, %401, %int0) : (!torch.tensor, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.tensor loc(#loc184)
    %403 = torch.prim.ListConstruct %int1, %int1, %int128, %int128 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %404 = torch.prim.ListConstruct %int1048576, %int1048576, %int1024, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int> loc(#loc)
    %405 = torch.operator "aten.as_strided"(%402, %403, %404, %int0) : (!torch.tensor, !torch.list<int>, !torch.list<int>, !torch.int) -> !torch.tensor loc(#loc185)

Looking at this, I agree with Ramiro's concerns, the warning in the documentation implies some behavior that may not be expressed in the captured MLIR. This is almost certainly some effect of the JIT tracing, and although we may be able to get it to work for this case, if we understand how the JIT tracing is generating this sequence, we may be able to map to a better set of operators.

AmosLewis · 2023-01-24T05:47:32Z

@ramiro050 @eric-k256 The as_strided is from the decomposition of torch.ops.aten.slice.Tensor when I use make_fx in python code afte the distilgpt2 model is imported. Just deleting slice decompose will get rid of the as_strided code.
https://github.com/pytorch/pytorch/blob/8f3600b966d896986e334b9a22c43e937ee0169d/torch/_decomp/decompositions.py#L663

ramiro050 · 2023-01-24T19:02:33Z

@ramiro050 @eric-k256 The as_strided is from the decomposition of torch.ops.aten.slice.Tensor when I use make_fx in python code afte the distilgpt2 model is imported. Just deleting slice decompose will get rid of the as_strided code. https://github.com/pytorch/pytorch/blob/8f3600b966d896986e334b9a22c43e937ee0169d/torch/_decomp/decompositions.py#L663

Awesome! Thanks for looking into it

AmosLewis force-pushed the as_stride branch 2 times, most recently from bf5a1a2 to 89ec255 Compare December 22, 2022 00:35

AmosLewis mentioned this pull request Dec 22, 2022

DistilGPT2 to TOSA nod-ai/SHARK-Studio#494

Closed

AmosLewis requested a review from vivekkhandelwal1 December 22, 2022 00:45

This comment was marked as outdated.

Sign in to view

AmosLewis force-pushed the as_stride branch from 89ec255 to 84571e0 Compare December 22, 2022 01:26

AmosLewis marked this pull request as ready for review December 22, 2022 02:40

AmosLewis changed the title ~~[MLIR][TORCH] Add e2e support for aten.as_stride~~ [MLIR][TORCH] Add e2e support for aten.as_strided Dec 22, 2022

vivekkhandelwal1 mentioned this pull request Dec 22, 2022

[Tosa] Add Torch to Tosa support for aten::as_strided #1710

Closed

AmosLewis requested a review from ramiro050 December 22, 2022 06:13

ramiro050 reviewed Dec 22, 2022

View reviewed changes

AmosLewis marked this pull request as draft December 22, 2022 17:48

AmosLewis force-pushed the as_stride branch 3 times, most recently from 1a43554 to b6caa9b Compare January 3, 2023 07:01

AmosLewis marked this pull request as ready for review January 3, 2023 07:01

AmosLewis requested review from ramiro050 and eric-k256 January 3, 2023 07:02

AmosLewis changed the title ~~[MLIR][TORCH] Add e2e support for aten.as_strided~~ [MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided Jan 3, 2023

AmosLewis force-pushed the as_stride branch 2 times, most recently from 5869097 to 9e031e6 Compare January 3, 2023 18:37

ramiro050 requested changes Jan 3, 2023

View reviewed changes

lib/Dialect/Torch/Transforms/RefineTypes.cpp Outdated Show resolved Hide resolved

AmosLewis force-pushed the as_stride branch 2 times, most recently from 38c1771 to 2a45694 Compare January 3, 2023 22:33

AmosLewis requested a review from ramiro050 January 3, 2023 22:34

eric-k256 approved these changes Jan 3, 2023

View reviewed changes

ramiro050 requested changes Jan 3, 2023

View reviewed changes

AmosLewis force-pushed the as_stride branch 2 times, most recently from 3635d78 to f94c723 Compare January 4, 2023 23:37

AmosLewis requested review from ramiro050 and eric-k256 January 4, 2023 23:45

AmosLewis force-pushed the as_stride branch from f94c723 to 99f2d5f Compare January 5, 2023 00:26

AmosLewis force-pushed the as_stride branch 3 times, most recently from edc78eb to 6d1b2f1 Compare January 16, 2023 06:04

[MLIR][TORCH] Add e2e support for aten.as_stride

96ba02d

AmosLewis force-pushed the as_stride branch from 6d1b2f1 to 96ba02d Compare January 16, 2023 07:17

AmosLewis closed this Jan 24, 2023

Vremold mentioned this pull request Jul 3, 2023

[torch-dialect] emit aten.as_strided op #2280

Open

AmosLewis deleted the as_stride branch January 19, 2024 19:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided #1742

[MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided #1742

AmosLewis commented Dec 21, 2022 •

edited

Loading

AmosLewis commented Dec 22, 2022 •

edited

Loading

This comment was marked as outdated.

ramiro050 left a comment

AmosLewis commented Dec 22, 2022 •

edited

Loading

ramiro050 commented Dec 22, 2022 •

edited

Loading

AmosLewis commented Dec 22, 2022

eric-k256 left a comment

ramiro050 Jan 3, 2023

ramiro050 Jan 3, 2023

AmosLewis Jan 4, 2023 •

edited

Loading

AmosLewis Jan 4, 2023

AmosLewis commented Jan 4, 2023

eric-k256 commented Jan 6, 2023

AmosLewis commented Jan 24, 2023

ramiro050 commented Jan 24, 2023

[MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided #1742

[MLIR][TORCH][TOSA] Add e2e Tosa support for aten.as_strided #1742

Conversation

AmosLewis commented Dec 21, 2022 • edited Loading

AmosLewis commented Dec 22, 2022 • edited Loading

This comment was marked as outdated.

ramiro050 left a comment

Choose a reason for hiding this comment

AmosLewis commented Dec 22, 2022 • edited Loading

ramiro050 commented Dec 22, 2022 • edited Loading

AmosLewis commented Dec 22, 2022

eric-k256 left a comment

Choose a reason for hiding this comment

ramiro050 Jan 3, 2023

Choose a reason for hiding this comment

ramiro050 Jan 3, 2023

Choose a reason for hiding this comment

AmosLewis Jan 4, 2023 • edited Loading

Choose a reason for hiding this comment

AmosLewis Jan 4, 2023

Choose a reason for hiding this comment

AmosLewis commented Jan 4, 2023

eric-k256 commented Jan 6, 2023

AmosLewis commented Jan 24, 2023

ramiro050 commented Jan 24, 2023

AmosLewis commented Dec 21, 2022 •

edited

Loading

AmosLewis commented Dec 22, 2022 •

edited

Loading

AmosLewis commented Dec 22, 2022 •

edited

Loading

ramiro050 commented Dec 22, 2022 •

edited

Loading

AmosLewis Jan 4, 2023 •

edited

Loading