Page MenuHomePhabricator

[mlir][Linalg] Generalize identity generic op removal to handle trivial broadcasts.
Needs ReviewPublic

Authored by mravishankar on Feb 15 2022, 4:22 PM.

Details

Summary

The existing pattern to remove identity generic ops when input is just
copied to the output looks for strict identity of the indexing
maps. This could be generalized to handle another case where the
indexing maps for the operand and result match, except for existence
of unit dimensions in either. In such cases use reshapes to replace
the generic op.
The existing pattern for removing identity generic op is superceded by
this pattern for ops with tensor semantics. The existing pattern is
now only used for the case where the op has buffer semantics.

Depends on D119365

Diff Detail

Event Timeline

mravishankar created this revision.Feb 15 2022, 4:22 PM
mravishankar requested review of this revision.Feb 15 2022, 4:22 PM

Rebase on dependent patch.

gysit added a comment.Feb 17 2022, 7:15 AM

Looks good to me.

I have one question though. What happens if the operand and the result contain unit-dims. My understanding is the code generates either an expand shape or a collapse shape. But could there be operand result pairs that collapse and expand at the same time (meaning there are unit-dims in both operand and result shape)? Shouldn't this be checked some how, for example, by counting the number of collapsed and expanded dims (one of them should always be zero). Or does this already exist and I did not see it?

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

Agreed :)

957

is this needed due to the expand / collapse shape? copy / broadcast generic should not require a cast?

967

nit: on buffers?

972

maybe EraseBufferCopyOp or similar would by a better name here?

mlir/test/Dialect/Linalg/canonicalize.mlir
786

nit: newline

Looks good to me.

I have one question though. What happens if the operand and the result contain unit-dims. My understanding is the code generates either an expand shape or a collapse shape. But could there be operand result pairs that collapse and expand at the same time (meaning there are unit-dims in both operand and result shape)? Shouldn't this be checked some how, for example, by counting the number of collapsed and expanded dims (one of them should always be zero). Or does this already exist and I did not see it?

Yeah, you are probably right. Maybe I can split the reassociation analysis to first try to fold just the operand, or just the result cause the mixing in with preserved 1-extent dims in result and operand can be painful. Ill test this pattern a bit more before submitting.

mravishankar marked 3 inline comments as done.

Address comments.

Minor fixes to comments.

Upload missed changes.

gysit accepted this revision.Feb 19 2022, 3:16 AM

LGTM!

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
874–876

The linalg op interface's getRank(OpOperand*) also returns rank 0 for scalar operands.

This revision is now accepted and ready to land.Feb 19 2022, 3:16 AM
gysit resigned from this revision.Oct 16 2022, 12:20 AM
This revision now requires review to proceed.Oct 16 2022, 12:20 AM
Herald added a project: Restricted Project. · View Herald TranscriptOct 16 2022, 12:20 AM
Herald added a subscriber: bzcheeseman. · View Herald Transcript