This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Linalg] Combine canonicalizers that deal with removing dead/redundant args.
ClosedPublic

Authored by mravishankar on Apr 12 2022, 1:43 PM.

Details

Summary

linalg.generic ops have canonicalizers that either remove arguments
not used in the payload, or redundant arguments. Combine these and
enhance the canonicalization to also remove results that have no use.
This is effectively dead code elimination for Linalg ops.

Diff Detail

Event Timeline

mravishankar created this revision.Apr 12 2022, 1:43 PM
Herald added a project: Restricted Project. · View Herald TranscriptApr 12 2022, 1:43 PM
mravishankar requested review of this revision.Apr 12 2022, 1:43 PM
nirvedhmeshram requested changes to this revision.Apr 12 2022, 4:58 PM
nirvedhmeshram added inline comments.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
16

Do you need this?

21

Same question as above.

705

There is a helper function called hasUniqueDim on line 114 that is not needed with this logic, so you can delete that.

710

is is -> is
maybe you want to mention these two conditions are "or" and not "and"?

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir
94

I believe as part of adding the tests here, you can delete two tests in canonicalize.mlir
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/canonicalize.mlir#L346-L392

This revision now requires changes to proceed.Apr 12 2022, 4:58 PM
okkwon added inline comments.Apr 12 2022, 6:24 PM
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
700

I think this can be put in the function.

705

inversePermutation() asserts on that the indexing map does not have a symbolic expression. We might need an explicit check for that.

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir
108

Q: do we expect a race condition as a desirable outcome?

mravishankar marked 3 inline comments as done.

Address comments and disallow dropping of dead outputs when it is needed for loop dim computation.

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

Ah thanks. Dont need it anymore. I thought I did.

21

This I do need.

700

Wont it result in a allocation on every invocation?

705

LinalgOp verify checks for no symbols. So it can be an assumed property. In general this computation is how loop bounds are determined based on indexing maps.

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir
94

Should be fine. More tests are not a problem :)

108

Not sure I follow the question. How is there a race condition here?

okkwon added inline comments.Apr 12 2022, 10:54 PM
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
700

Good point!

705

Thanks for the info!

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir
108

Oh, I think I was confused because %arg0 is used for all outputs. But now it looks fine because only %b0 is used. Sorry for the noise.

gysit added inline comments.Apr 13 2022, 5:54 AM
mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
722

All the uses I found are actually calling getValue immediately to get the AffineMap. I would thus suggest to use the existing getTiedIndexingMap function to avoid growing the interface? If an attribute is needed then AffineMapAttr::get() is pretty compact to get it?

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
168

nit: resultTensorTypses -> resultTensorTypes

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

nit: in most places we prefix the operands of the newly created op with new instead of modified. Feel free if you think modified is a better fit here!

721

I think it would be nice if we could avoid the assumption and rely on the structured op interface only?

We could for example use a slightly different algorithm that for every loop dimension counts the number of indexing map result appearances. E.g. if an op has the indexing maps (d0, d1) -> (d0, d1) and (d0, d1) -> (d0) we could get the vector dimCounts = [2, 1] saying d0 appears twice and d1 once. We can then iterate opOperands and and try to subract the indexingMap tied to the current opOperand. If the dimCounts vector contains a zero after the subtraction we cannot delete it. In the example, only the operand tied to (d0, d1) -> (d0) can be removed.

I think if we have two lambdas/functions to add and subtract the result dim counts for a specific indexing map the change may even simplify the pattern a little bit. Admittedly, it also implements some logic that is usually hidden in getShapesToLoopsMap (the real one in the structured op interface).

An alternative could be to have a method on the structured op interface that takes a list of dropped operands and returns if the remaining ones suffice to define the iteration domain E.g.:

genericOp.isIterationDomainDefinedBy(/*droppedOperands*/=SmallVector<OpOperand*>{opOperandToDelete1, opOperandToDelete2});

which could then be used in the verification as well. That could fit well with the current algorithm.

845

I think it should be possible to use rewriter.replaceOp(genericOp, newResults) by setting the dropped values to the old result or nullptr? Or is there some check that prevents this?

nirvedhmeshram accepted this revision.Apr 13 2022, 6:45 AM

LGTM for my comments.

This revision is now accepted and ready to land.Apr 13 2022, 6:45 AM
mravishankar marked 5 inline comments as done.

Address most comments.

mravishankar marked an inline comment as done.Apr 13 2022, 1:52 PM
mravishankar added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
722

You are right. I dont need this anymore.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
168

Seems like Gollum came and wrote that signature :)

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

Good idea. Smaller names.

721

The first approach will require walking the indexing maps, and maintaining involved data structures to get it current. Seems more invasive and error prone.

The second approach looks feasible. Let me try that.

Use interface method to check if operands can be dropped while still being able to compute loop bounds.

Fix a bug in the canonicalizer and add test that exercises that path.

This revision was landed with ongoing or failed builds.May 11 2022, 10:23 PM
This revision was automatically updated to reflect the committed changes.