[MLIR] Generalize expand_shape to take shape as explicit input #69267

Shukla-Gaurav · 2023-10-17T00:07:04Z

This patch generalizes tensor.expand_shape and memref.expand_shape to consume the output shape as a list of SSA values. This enables us to implement generic reshape operations with dynamic shapes using collapse_shape/expand_shape pairs.

The output_shape input to expand_shape follows the static/dynamic representation that's also used in tensor.extract_slice.

Differential Revision: https://reviews.llvm.org/D140821

MaheshRavishankar

Ok, I think this is not looking too bad. Adding the helper build method here would make the diff much smaller I think.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

github-actions · 2023-10-19T16:14:23Z

✅ With the latest revision this PR passed the C/C++ code formatter.

MaheshRavishankar · 2023-10-20T00:13:14Z

Hey, I left a few comments. Let me know when this is ready for review.

MaheshRavishankar

Next rounds of comments. Thanks!

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

Shukla-Gaurav · 2023-10-30T16:18:55Z

@MaheshRavishankar Thanks for reviewing, have addressed all the comments, and is ready for re-review.
@ramiro050 @sanjoy Thanks for implementing this, it was needed for a long time. As I am continuing this, would love to have your review on this.

MaheshRavishankar

ok, this looks good to me. Lets test this with IREE. Ill let you know how to do that.

ramiro050

Thanks for getting this across the finish line! LGTM

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

*DO NOT SUBMIT* This commit is to test and fix the bugs introduced by importing tensor.expand_shape changes in IREE. (check here llvm/llvm-project#69267) Signed-Off-by: Gaurav Shukla <gashukla@amd.com>

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

Shukla-Gaurav · 2024-04-17T18:21:14Z

@MaheshRavishankar Could you please re-review it? There is still one failure in reshape_fusion(as discussed), meanwhile I can fix that. Thanks :)

aartbik · 2024-04-17T21:25:34Z

I added a comment that has not been addressed yet?

MaheshRavishankar

Thanks! Left a few mostly minor comments. Nothing that I see as blocker for landing.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

MaheshRavishankar

Ok, looks ogod to me. Left a single comment but it is not a blocker and is OK to do it after this PR as well since that part needs to be fixed after this PR lands.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

bjacob · 2024-04-19T17:33:13Z

Ok, looks ogod to me. Left a single comment but it is not a blocker and is OK to do it after this PR as well since that part needs to be fixed after this PR lands.

yes please please do as follow-up and not lose the progress on (very slow) CI here :-)

bjacob · 2024-04-21T01:19:38Z

Windows CI failed with

fatal error C1060: compiler is out of heap space

Probably unrelated to this PR, but also not completely impossible that something about this PR increases compiler memory usage and pushes it over a threshold, so we need to get a green CI run before merging. Unfortunately I don't see a "retrigger" button here. Can you push an empty commit to this PR to cause a retrigger maybe?

This patch generalizes tensor.expand_shape and memref.expand_shape to consume the output shape as a list of SSA values. This enables us to implement generic reshape operations with dynamic shapes using collapse_shape/expand_shape pairs. The output_shape input to expand_shape follows the static/dynamic representation that's also used in `tensor.extract_slice`. Differential Revision: https://reviews.llvm.org/D140821

joker-eph · 2024-04-21T11:56:37Z

I think this broke the bot: https://lab.llvm.org/buildbot/#/builders/61/builds/57064

#69267)" This reverts commit e095d97.

…t" (#89540) Reverts #69267 this broke some bots.

joker-eph · 2024-04-21T12:34:19Z

Reverted in #89540

…69267) This patch generalizes tensor.expand_shape and memref.expand_shape to consume the output shape as a list of SSA values. This enables us to implement generic reshape operations with dynamic shapes using collapse_shape/expand_shape pairs. The output_shape input to expand_shape follows the static/dynamic representation that's also used in `tensor.extract_slice`. Differential Revision: https://reviews.llvm.org/D140821 Co-authored-by: Ramiro Leal-Cavazos <ramiroleal050@gmail.com>

…t" (llvm#89540) Reverts llvm#69267 this broke some bots.

llvmbot added mlir:linalg mlir:sparse Sparse compiler in MLIR mlir mlir:tensor mlir:memref labels Oct 17, 2023

MaheshRavishankar mentioned this pull request Oct 17, 2023

Adding reshape fall back case for unimplemented cases in aten view for time being llvm/torch-mlir#2515

Closed

MaheshRavishankar requested changes Oct 17, 2023

View reviewed changes

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from 11c8940 to e8ac533 Compare October 19, 2023 16:12

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from e8ac533 to 3961b9a Compare October 19, 2023 19:25

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from 3961b9a to 3f96df1 Compare October 26, 2023 16:38

MaheshRavishankar requested changes Oct 26, 2023

View reviewed changes

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td Outdated Show resolved Hide resolved

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp Outdated Show resolved Hide resolved

Shukla-Gaurav force-pushed the arcpatch-D140821 branch 2 times, most recently from dec4859 to f9d0f60 Compare October 30, 2023 16:09

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from f9d0f60 to 2e27433 Compare October 30, 2023 17:15

MaheshRavishankar reviewed Oct 31, 2023

View reviewed changes

ramiro050 approved these changes Nov 3, 2023

View reviewed changes

newling mentioned this pull request Nov 10, 2023

Support all cases of lowering of ConvertAtenViewOp llvm/torch-mlir#2567

Open

newling reviewed Nov 11, 2023

View reviewed changes

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp Show resolved Hide resolved

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from 2e27433 to 273d71d Compare November 16, 2023 16:02

Shukla-Gaurav changed the title ~~[DRAFT] Generalize expand_shape to take shape as explicit input~~ [MLIR] Generalize expand_shape to take shape as explicit input Nov 30, 2023

newling mentioned this pull request Feb 7, 2024

More TorchToLinalg view failure cases llvm/torch-mlir#2877

Open

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from 273d71d to 24ee8aa Compare April 3, 2024 10:23

Shukla-Gaurav requested review from aartbik, PeimingLiu, yinying-lisa-li, matthias-springer and hanhanW as code owners April 3, 2024 10:23

ScottTodd reviewed Apr 16, 2024

View reviewed changes

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp Show resolved Hide resolved

MaheshRavishankar requested changes Apr 17, 2024

View reviewed changes

bjacob reviewed Apr 18, 2024

View reviewed changes

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td Outdated Show resolved Hide resolved

bjacob mentioned this pull request Apr 18, 2024

Let memref.{expand,collapse}_shape implement ReifyRankedShapedTypeOpInterface #89111

Closed

Shukla-Gaurav force-pushed the arcpatch-D140821 branch 5 times, most recently from 76caac7 to f352a94 Compare April 19, 2024 14:34

MaheshRavishankar approved these changes Apr 19, 2024

View reviewed changes

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp Show resolved Hide resolved

Shukla-Gaurav force-pushed the arcpatch-D140821 branch from f352a94 to 5ea9b9a Compare April 21, 2024 10:33

bjacob merged commit e095d97 into llvm:main Apr 21, 2024
4 checks passed

joker-eph added a commit that referenced this pull request Apr 21, 2024

Revert "[MLIR] Generalize expand_shape to take shape as explicit input (

b8bffad

#69267)" This reverts commit e095d97.

joker-eph mentioned this pull request Apr 21, 2024

Revert "[MLIR] Generalize expand_shape to take shape as explicit input" #89540

Merged

joker-eph added a commit that referenced this pull request Apr 21, 2024

Revert "[MLIR] Generalize expand_shape to take shape as explicit inpu…

8c0341d

…t" (#89540) Reverts #69267 this broke some bots.

aniplcc pushed a commit to aniplcc/llvm-project that referenced this pull request Apr 21, 2024

Revert "[MLIR] Generalize expand_shape to take shape as explicit inpu…

f32304c

…t" (llvm#89540) Reverts llvm#69267 this broke some bots.

This was referenced Apr 22, 2024

[libc] Make fenv and math tests preserve fenv_t state #89658

Merged

[libc] Clean up alternate test framework support #89659

Merged

Shukla-Gaurav mentioned this pull request Apr 25, 2024

[MLIR] Generalize expand_shape to take shape as explicit input #90040

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR] Generalize expand_shape to take shape as explicit input #69267

[MLIR] Generalize expand_shape to take shape as explicit input #69267

Shukla-Gaurav commented Oct 17, 2023 •

edited

MaheshRavishankar left a comment

github-actions bot commented Oct 19, 2023 •

edited

MaheshRavishankar commented Oct 20, 2023

MaheshRavishankar left a comment

Shukla-Gaurav commented Oct 30, 2023

MaheshRavishankar left a comment

ramiro050 left a comment

Shukla-Gaurav commented Apr 17, 2024

aartbik commented Apr 17, 2024

MaheshRavishankar left a comment

MaheshRavishankar left a comment

bjacob commented Apr 19, 2024

bjacob commented Apr 21, 2024

joker-eph commented Apr 21, 2024

joker-eph commented Apr 21, 2024

[MLIR] Generalize expand_shape to take shape as explicit input #69267

[MLIR] Generalize expand_shape to take shape as explicit input #69267

Conversation

Shukla-Gaurav commented Oct 17, 2023 • edited

MaheshRavishankar left a comment

Choose a reason for hiding this comment

github-actions bot commented Oct 19, 2023 • edited

MaheshRavishankar commented Oct 20, 2023

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Shukla-Gaurav commented Oct 30, 2023

MaheshRavishankar left a comment

Choose a reason for hiding this comment

ramiro050 left a comment

Choose a reason for hiding this comment

Shukla-Gaurav commented Apr 17, 2024

aartbik commented Apr 17, 2024

MaheshRavishankar left a comment

Choose a reason for hiding this comment

MaheshRavishankar left a comment

Choose a reason for hiding this comment

bjacob commented Apr 19, 2024

bjacob commented Apr 21, 2024

joker-eph commented Apr 21, 2024

joker-eph commented Apr 21, 2024

Shukla-Gaurav commented Oct 17, 2023 •

edited

github-actions bot commented Oct 19, 2023 •

edited