[mlir][transform] Implement `FlattenElementwiseLinalgOp` transform op #81431

srcarroll · 2024-02-11T21:00:23Z

A transform.structured.flatten_elementwise op is implemented for flattening the iteration space and (applicable) operaands/results to a single dimension.

MaheshRavishankar · 2024-02-13T20:32:41Z

This seems to be a duplicate of

llvm-project/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Line 1086 in ac97562

collapseOpIterationDims(LinalgType op,

srcarroll · 2024-02-13T20:41:10Z

Ah interesting. I wasn't aware of this. I wouldn't say this whole PR is a duplicate though. The purpose is to implement a transform op. I can refactor my changes to reuse collapseOpIterationDims but right now it only supports collapsing linalg.generic and linalg.copy. But i could refactor that function to cover more if you think that's a good idea

MaheshRavishankar · 2024-02-14T01:39:03Z

The transformation should work on any LinalgOp. Anything that works on linalg.generic would directly work on LinalgOp you just need to change the type of operation from GenericOp to LinalgOp and it should just work.

srcarroll · 2024-02-14T01:46:20Z

Yes I understand. What I'm saying is that this is just a helper function, not a transform op, as in transform.structured.flatten. This PR defines such a transform op. I do agree that the helper function can be refactored and reused by my op implementation. But I'm just pointing out that it does indeed need a refactor for me to use it.

srcarroll · 2024-02-14T01:50:32Z

Yes I understand. What I'm saying is that this is just a helper function, not a transform op, as in transform.structured.flatten. This PR defines such a transform op. I do agree that this can be refactored and reused by my op implementation. But I'm just pointing out that it does indeed need a refactor for me to use it. And I'm happy to do so, to be clear

MaheshRavishankar · 2024-02-14T03:15:18Z

Yes I understand. What I'm saying is that this is just a helper function, not a transform op, as in transform.structured.flatten. This PR defines such a transform op. I do agree that this can be refactored and reused by my op implementation. But I'm just pointing out that it does indeed need a refactor for me to use it. And I'm happy to do so, to be clear

Sounds good. I am not sure how much refactoring is needed. This entry point has been used downstream without too many issues, but happy to see what you want to do. If it does change, this would be a breaking change.

srcarroll · 2024-02-14T05:26:31Z

I made a first attempt at a refactor and it wasn't so bad. i tried to make the least changes possible but its still a wip and will appreciate feedback on the design aspects.

MaheshRavishankar

Thanks! Just a few comments, but this looks along the right path.

MaheshRavishankar · 2024-02-15T21:37:51Z

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

-  // Get the iterator types for the operand.
-  SmallVector<utils::IteratorType> iteratorTypes =
-      getCollapsedOpIteratorTypes(op.getIteratorTypesArray(), collapsingInfo);
+  Operation *collapsedOp = clone(


Why is the operation being cloned?

Follow up.. I see why you are cloning it.... This is an interesting way of doing it. Does it work?

MaheshRavishankar · 2024-02-15T21:39:08Z

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

+
+  // TODO: Find a more general way to determine if op requires explicit
+  // indexing_maps and iterator_types
+  if (isa<linalg::GenericOp>(op)) {


I dont think you need to do this. All LinalgOps have getIteratorTypesArray (or should have). Also if you dont clone the op you dont need to set the indexing maps etc. explicitly....

Yah i know all LinalgOps have these attrs and the corresponding getters. But most of the structured ops don't have an explicit attr for indexing maps and iterator types. so if i unconditionally try to set them, i will end up with named ops with additional attrs that aren't at all associated with the implicit ones. I'm sure there's a way to do this correctly. I just dont know.

I'm not sure how any of this would work without cloning the op (other than having a switch statement checking every single linalg op that exists). Of course the alternative is to just always convert to a linalg.generic op. But I dont like that. I much prefer keeping the named op and that's why I chose to clone.

hmm. there is this kMemoizedIndexingMapsAttrName that's for named ops. so i could check if the op has this and then replace it. otherwise replace the expicit indexing maps. but i dont see a similar attr name for iterator types

it woudl be nice if there was a general setIndexingMaps and setIteratoryTypes for arbitrary LinalgOps

One alternative is that you could have a interface method to LinalgOp that allows you to clone with a given indexing_maps and iterator_types and just uses the region from the original operation. Each specific op could just implement its own version (including linalg.generic).

Not a bad idea. Id have to see what it looks like to understand the implications fully but I sense that it might be too demanding to require all ops to implement that. Would this work for the ops generated from core_named_ops.py (can't remember the exact file name right now)?

To start with you can just have the method return a LogicalResult and the default implementation return failure()and only add the implementation for ops that you want to support. Those can be filled in over time.

On second thought maybe instead of going full interface method, just add a templated method in this file for cloning? The default can be for handling named ops, and the generic op can get its specialization

Thanks for the feedback. I wont be able to get back to this until this weekend or next, but I'll take a look at your suggestion and keep you updated.

MaheshRavishankar · 2024-02-27T07:04:40Z

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

-  // Get the iterator types for the operand.
-  SmallVector<utils::IteratorType> iteratorTypes =
-      getCollapsedOpIteratorTypes(op.getIteratorTypesArray(), collapsingInfo);
+  if (op->hasAttr("indexing_maps")) {


This seems iffy to me. THis is looking into internal implementation details of the LinalgOp representation in C++ and is hard to maintain. Let me give you an example. I have no idea what the indexing maps or iterator types are stored as. I just use the utility methods on ops to get this information. Side-stepping that can introduce silent bugs if something changes in the implementation of the operation.

I understand you are trying to generalize these things.. One way to do that would be to define a method

template <typename OpTy> Operation * cloneWithIndexingMapsIteratorTypesAndOperands(RewriterBase &rewriter, OpTy origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) { return nullptr; } template <> Operation *cloneWithIndexingMapsIteratorTypesAndOperands<LinalgOp>(RewriterBase &rewriter, LinalgOp origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) { return clone(rewriter, origOp, resultTypes, inputOperands, outputOperands); } template <> Operation *cloneWithIndexingMapsIteratorTypesAndOperands<GenericOp>(RewriterBase &rewriter, GenericOp origOp, TypeRange resultTypes, ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorType, ValueRange inputOperands, ValueRange outputOperands) { SmallVector<utils::IteratorType> iteratorTypes = getCollapsedOpIteratorTypes(op.getIteratorTypesArray(), collapsingInfo); // Get the indexing maps. auto indexingMaps = llvm::map_to_vector(op.getIndexingMapsArray(), [&](AffineMap map) { return getCollapsedOpIndexingMap(map, collapsingInfo); }); Operation *collapsedOp = rewriter.create<linalg::GenericOp>( loc, resultTypes, inputOperands, outputOperands, indexingMaps, iteratorTypes, [](OpBuilder &builder, Location loc, ValueRange args) {}); Block *origOpBlock = &op->getRegion(0).front(); Block *collapsedOpBlock = &collapsedOp->getRegion(0).front(); rewriter.mergeBlocks(origOpBlock, collapsedOpBlock, collapsedOpBlock->getArguments()); }

Then you can

if (auto genericOp = dyn_cast<GenericOp>(op)) { cloneWithIndexingMapsIteratorTypesAndOperands(rewriter, genericOp, ....) } else { cloneWithIndexingMapsIteratorTypesAndOperands(rewriter, cast<LinalgOp>(op), ...) }

I think that gives you what you want and doesnt leak internal implementation details of the op.

yah i totally agree this recent change is garbage. shouldn't have even pushed it :).

i'm going to go with your suggestion. however, just so i understand correctly, this is functionally equivalent to what i had before (checking if generic), but you are suggesting this for cleanliness and maintainability?

Yup. We dont need to look at whether an attribute is present or not.. So a bit more cleaner?

agreed. do you think these cloneWithIndexingMapsIteratorTypesAndOperands should be static functions specific to this file? or as part of some lib?

i pushed a modified version of your suggestion. First, i templated the input and output op to enforce that the cloning should be same named op. Second, i didn't want to have both collapsed operands and collapsingInfo as arguments to a function because that leaves room for inconsistency. And the collapsingInfo is needed for the generic's indexing map. So I did more refactoring with that in mind.

MaheshRavishankar

Nice! Thanks a lot!

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

…llvm#81431) A `transform.structured.flatten_elementwise` op is implemented for flattening the iteration space and (applicable) operands/results to a single dimension.

srcarroll added 2 commits February 11, 2024 02:33

Implement FlattenElementwiseLinalgOp transform

6e05d6a

Add a couple regression tests

aff79ba

srcarroll added 2 commits February 13, 2024 22:27

Refactor collapseOpIterationDims to work for all linalg ops

cd0ebb1

Refactor FlattenElementwiseLinalgOp to use collapseOpIterationDims

780394c

srcarroll force-pushed the linalg-elementwise-flatten-transform branch from 9d0b35a to 780394c Compare February 14, 2024 05:23

srcarroll and others added 4 commits February 13, 2024 23:30

Add EOL

db62df3

Remove restrictions and add tensor test

27fb208

Fix designated initializers warning

6ec2d19

Merge branch 'main' into linalg-elementwise-flatten-transform

9676452

MaheshRavishankar requested changes Feb 15, 2024

View reviewed changes

Fix bug

f6e1eca

MaheshRavishankar reviewed Feb 27, 2024

View reviewed changes

srcarroll added 2 commits February 27, 2024 16:33

apply reviewer suggestion

a350306

remove unnecessary cast

bdd04e2

MaheshRavishankar approved these changes Feb 28, 2024

View reviewed changes

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

Address comment and elaborate summary

12444bc

srcarroll marked this pull request as ready for review February 28, 2024 07:40

srcarroll requested review from nicolasvasilache and dcaballe as code owners February 28, 2024 07:40

Merge branch 'main' into linalg-elementwise-flatten-transform

9b33d45

srcarroll merged commit b6f4dd9 into llvm:main Feb 28, 2024

srcarroll deleted the linalg-elementwise-flatten-transform branch June 5, 2024 02:59

[mlir][transform] Implement FlattenElementwiseLinalgOp transform op #81431

[mlir][transform] Implement FlattenElementwiseLinalgOp transform op #81431

Uh oh!

Conversation

srcarroll commented Feb 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaheshRavishankar commented Feb 13, 2024

Uh oh!

srcarroll commented Feb 13, 2024

Uh oh!

MaheshRavishankar commented Feb 14, 2024

Uh oh!

srcarroll commented Feb 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

srcarroll commented Feb 14, 2024

Uh oh!

MaheshRavishankar commented Feb 14, 2024

Uh oh!

srcarroll commented Feb 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

srcarroll Feb 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[mlir][transform] Implement `FlattenElementwiseLinalgOp` transform op #81431

[mlir][transform] Implement `FlattenElementwiseLinalgOp` transform op #81431

srcarroll commented Feb 11, 2024 •

edited

Loading

srcarroll commented Feb 14, 2024 •

edited

Loading

srcarroll commented Feb 14, 2024 •

edited

Loading

srcarroll Feb 27, 2024 •

edited

Loading