[TORCH] Add decomposition of `aten.linear` op #862

Shukla-Gaurav · 2022-05-18T16:31:30Z

This commit adds decomposition of aten.linear op.

TODO: remove existing implementation and add a simple test case. Also handle the (3D,2D) case of matmul.

Signed-Off-by: Gaurav Shukla

silvasean · 2022-05-19T10:39:36Z

Remove the TorchToLinalg code as well please since it is not not useful.

Shukla-Gaurav · 2022-06-21T19:50:15Z

Once #951 is merged, this PR will be unblocked wrt (3D, 2D) case of matmul.

Shukla-Gaurav · 2022-07-01T16:33:25Z

@sjarus Could you please take a look into this, ResNet18Static is failing in torch to tosa lowering. Thanks!

sjarus · 2022-07-02T01:20:44Z

I'll attempt to repro this early next week . Right now I'm trying to get a pipeline of pending BERT related legalizations out. Do you have detailed verbose output from the test ?

Shukla-Gaurav · 2022-07-18T15:17:36Z

@sjarus There is a case where aten.matmul fails to lower to tosa. I quickly looked into it but did not get the solution. Could you please look into it?
%3 = torch.aten.matmul %arg0, %2 : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[3,5],f32> -> !torch.vtensor<[?,?,5],f32>.

Test case:

class Matmul3D(torch.nn.Module):               
    def __init__(self):                                                                                  
        super().__init__()                          
                                                                                                         
    @export                                         
    @annotate_args([                                
        None,                                                                                            
        ([-1, -1, -1], torch.float32, True),
        ([3, 5], torch.float32, True),                                                                   
    ])                               
    def forward(self, lhs, rhs):                                                                         
        return torch.matmul(lhs, rhs)               
                                                                                                         
                                                                                                         
@register_test_case(module_factory=lambda: Matmul3D())                                                  
def Matmul_3d(module, tu: TestUtils):              
    module.forward(tu.rand(3, 4, 3), tu.rand(3, 5)) ```

sjarus · 2022-07-18T17:22:42Z

What happens when you use 3, 4, 3 instead of -1, -1, -1 in the annotate args ?

Shukla-Gaurav · 2022-07-19T14:30:35Z

What happens when you use 3, 4, 3 instead of -1, -1, -1 in the annotate args ?

It passed successfully!
But dynamic dimensions causes runtime error:
python: /home/gaurav/MLIRepos/torch-mlir/build/tools/mlir/include/mlir/IR/BuiltinTypeInterfaces.h.inc:163: bool mlir::ShapedType::isDynamicDim(unsigned int) const: Assertion idx < getRank() && "invalid index for shaped type"' failed.
`

Shukla-Gaurav · 2022-07-20T16:30:21Z

@sjarus Did you get a chance to look at it?

sjarus · 2022-07-21T02:13:20Z

Right, sorry for the delay. The experiment result appears aligned to my expectations, but I agree that if dynamic dims are unsupported then it should fail gracefully and not assert.

It seems the best way to unblock is to keep the test static and refer in a comment to me to update the code so it doesn't assert, or better yet can pass, tho the latter is typically harder. Would that work ?

Shukla-Gaurav · 2022-08-18T16:31:23Z

@sjarus I have made the test case partially static for now in order to pass the tosa test. Should I create a new issue for the dynamic case handling in tosa matmul?
FYI @silvasean

silvasean · 2022-08-18T18:56:58Z

We definitely don't want to regress our testing of dynamic shapes for all backends due a limitation in the TOSA backend, so I would avoid changing the test. We can use the functionality in #1242 to control the decompositions, but I need to make a PR to centralize that logic first -- will try to do that today.

@Shukla-Gaurav can you also file a bug for the TOSA assertion failure? Ideally include a torch-mlir-opt command line invoking the failing pass and a reduced IR snippet.

We were already hitting many cases where backends different in terms of the legal ops that they wanted. This caused unnecessary coupling between the backends. Examples: - llvm#1161 - llvm#862 This PR centralizes all compilation to go through `torch_mlir.compile` so that we can keep the logic centralized there. We should move these lists closer to each backend. Especially cases like llvm#862 where blocking a decomposition is necessary to avoid a crash emphasize that the set of decompositions is tightly coupled to the backend, and should be "controlled by the backend" and not something arbitrarily tweakable. Also: - Fix a small bug in the way we passed through the backendLegalOps option. - Add better error messages in `torch_mlir.compile` for import errors.

silvasean · 2022-08-19T00:12:10Z

We definitely don't want to regress our testing of dynamic shapes for all backends due a limitation in the TOSA backend, so I would avoid changing the test. We can use the functionality in #1242 to control the decompositions, but I need to make a PR to centralize that logic first -- will try to do that today.

See #1248

We were already hitting many cases where backends different in terms of the legal ops that they wanted. This caused unnecessary coupling between the backends. Examples: - llvm#1161 - llvm#862 This PR centralizes all compilation to go through `torch_mlir.compile` so that we can keep the logic centralized there. We should move these lists closer to each backend. Especially cases like llvm#862 where blocking a decomposition is necessary to avoid a crash emphasize that the set of decompositions is tightly coupled to the backend, and should be "controlled by the backend" and not something arbitrarily tweakable. Also: - Fix a small bug in the way we passed through the backendLegalOps option. - Add better error messages in `torch_mlir.compile` for import errors.

We were already hitting many cases where backends different in terms of the legal ops that they wanted. This caused unnecessary coupling between the backends. Examples: - #1161 - #862 This PR centralizes all compilation to go through `torch_mlir.compile` so that we can keep the logic centralized there. We should move these lists closer to each backend. Especially cases like #862 where blocking a decomposition is necessary to avoid a crash emphasize that the set of decompositions is tightly coupled to the backend, and should be "controlled by the backend" and not something arbitrarily tweakable. Also: - Fix a small bug in the way we passed through the backendLegalOps option. - Add better error messages in `torch_mlir.compile` for import errors.

Shukla-Gaurav · 2022-08-31T16:22:01Z

@silvasean Thanks for adding support for more control over decomposition based on backend. Could you please review this again?

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

This commit adds decomposition of `aten.linear` op. Due to limited support at tosa backend in case of dynamic dimensions, this decomposition is currently disabled for tosa backend. Signed-Off-by: Gaurav Shukla <[email protected]>

Shukla-Gaurav mentioned this pull request May 26, 2022

Avoid doing bias-add by setting the bias value as the outs operand #836

Open

Shukla-Gaurav force-pushed the gaurav/decompose_linear branch 2 times, most recently from d87086f to 9b7e965 Compare June 28, 2022 16:23

Shukla-Gaurav changed the title ~~[WIP][TORCH] Add decomposition of aten.linear op~~ [TORCH] Add decomposition of aten.linear op Jul 1, 2022

Shukla-Gaurav mentioned this pull request Jul 6, 2022

Torch to TOSA conversion fails to legalize 'torch.constant.int' #961

Closed

Shukla-Gaurav force-pushed the gaurav/decompose_linear branch 3 times, most recently from c79cc85 to 5cc818b Compare July 18, 2022 09:36

ramiro050 mentioned this pull request Jul 18, 2022

Failed to legalize operation 'torch.aten.linear' that was explicitly marked illegal #1076

Open

Shukla-Gaurav force-pushed the gaurav/decompose_linear branch 3 times, most recently from 3b8a6f9 to 7c3d8cf Compare August 18, 2022 15:51

silvasean mentioned this pull request Aug 19, 2022

Add a way for backends to control which ops are legal for them. #1248

Merged

Shukla-Gaurav force-pushed the gaurav/decompose_linear branch from 7c3d8cf to 6d0c6bc Compare August 31, 2022 16:14

Shukla-Gaurav requested a review from silvasean August 31, 2022 16:21

silvasean requested changes Aug 31, 2022

View reviewed changes

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

silvasean approved these changes Aug 31, 2022

View reviewed changes

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

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

[TORCH] Add decomposition of aten.linear op

8cf264f

This commit adds decomposition of `aten.linear` op. Due to limited support at tosa backend in case of dynamic dimensions, this decomposition is currently disabled for tosa backend. Signed-Off-by: Gaurav Shukla <[email protected]>

Shukla-Gaurav force-pushed the gaurav/decompose_linear branch from 6d0c6bc to 8cf264f Compare September 6, 2022 12:32

Shukla-Gaurav merged commit 99093d0 into llvm:main Sep 7, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TORCH] Add decomposition of `aten.linear` op #862

[TORCH] Add decomposition of `aten.linear` op #862

Shukla-Gaurav commented May 18, 2022 •

edited

Loading

silvasean commented May 19, 2022

Shukla-Gaurav commented Jun 21, 2022

Shukla-Gaurav commented Jul 1, 2022

sjarus commented Jul 2, 2022

Shukla-Gaurav commented Jul 18, 2022 •

edited

Loading

sjarus commented Jul 18, 2022

Shukla-Gaurav commented Jul 19, 2022

Shukla-Gaurav commented Jul 20, 2022

sjarus commented Jul 21, 2022

Shukla-Gaurav commented Aug 18, 2022 •

edited

Loading

silvasean commented Aug 18, 2022

silvasean commented Aug 19, 2022

Shukla-Gaurav commented Aug 31, 2022

[TORCH] Add decomposition of aten.linear op #862

[TORCH] Add decomposition of aten.linear op #862

Conversation

Shukla-Gaurav commented May 18, 2022 • edited Loading

silvasean commented May 19, 2022

Shukla-Gaurav commented Jun 21, 2022

Shukla-Gaurav commented Jul 1, 2022

sjarus commented Jul 2, 2022

Shukla-Gaurav commented Jul 18, 2022 • edited Loading

sjarus commented Jul 18, 2022

Shukla-Gaurav commented Jul 19, 2022

Shukla-Gaurav commented Jul 20, 2022

sjarus commented Jul 21, 2022

Shukla-Gaurav commented Aug 18, 2022 • edited Loading

silvasean commented Aug 18, 2022

silvasean commented Aug 19, 2022

Shukla-Gaurav commented Aug 31, 2022

[TORCH] Add decomposition of `aten.linear` op #862

[TORCH] Add decomposition of `aten.linear` op #862

Shukla-Gaurav commented May 18, 2022 •

edited

Loading

Shukla-Gaurav commented Jul 18, 2022 •

edited

Loading

Shukla-Gaurav commented Aug 18, 2022 •

edited

Loading