Page MenuHomePhabricator

[MLIR][Shape] Add canonicalizations for `shape.broadcast`

Authored by frgossen on Apr 16 2021, 3:29 AM.



Eliminate empty shapes from the operands, partially fold all constant shape
operands, and fix normal folding.

Diff Detail

Event Timeline

frgossen created this revision.Apr 16 2021, 3:29 AM
frgossen requested review of this revision.Apr 16 2021, 3:29 AM

Nice. Some comments to further generalize.


This could be generalized to broadcasting all constant arguments. It would simplify the IR a little. Not sure it is worth it.


An alternative would be to filter all empty shapes, which simplifies the broadcast, as well.

tpopp added inline comments.Apr 16 2021, 7:16 AM

I mostly agree with this. I think it would also make the above code easier to read by not requiring this stateful onlyPotentiallyNonEmptyShape.

My only concern with doing this in fold though is that I could imagine being confused if I create an op, and operands are missing due to folding upon Op creation.

frgossen updated this revision to Diff 338141.Apr 16 2021, 8:55 AM
frgossen marked 3 inline comments as done.

Address comments


I'm not sure if folding can/should create new operations anyways. If we were to replace a shape.broadcast op with a simplified version that sounds more like a canonicalization.

Arc diff doesn't update the description, so FYI in case you didn't, update to mention the additional features you added.

frgossen retitled this revision from [MLIR][Shape] Generalize `shape.broadcast` folding to [MLIR][Shape] Add canonicalizations for `shape.broadcast`.Apr 21 2021, 12:35 AM
frgossen edited the summary of this revision. (Show Details)

Thanks, updated it.

A few places you checked that *.isa<ShapeType>() == *.isa<ShapeType>(). Out of curiosity, why not just compare the types?


Is this supposed to be ShapedType and ShapeType?


nit: I think you can also replace this with something like nonEmptyShapeOperands = llvm::make_filter_range(op->getOperands(), isNonEmptyShape)


Is this change actually necessary for a reason that you know, or are you just being extra cautious?


nit: maybe a different name like foldedConstantShape? I was confused at first thinking this should be a Vector of Values based on the name.


nit: if (auto constShape = shape.getDefiningOp<ConstShapeOp>()) {


Alternatively you could just check to see if foldedConstantOperands is empty.

frgossen updated this revision to Diff 339133.Apr 21 2021, 1:12 AM
frgossen marked 6 inline comments as done.
frgossen edited the summary of this revision. (Show Details)

Address comments



Rather being cautious.


If there is exactly one constant operand foldedConstantOperands would not be empty but we don't want to rewrite because nothing changed.

tpopp accepted this revision.Apr 21 2021, 6:51 AM
This revision is now accepted and ready to land.Apr 21 2021, 6:51 AM