Page MenuHomePhabricator

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

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

Details

Summary

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.

mlir/lib/Dialect/Shape/IR/Shape.cpp
490

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

492

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
mlir/lib/Dialect/Shape/IR/Shape.cpp
492

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

mlir/lib/Dialect/Shape/IR/Shape.cpp
492

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?

mlir/lib/Dialect/Shape/IR/Shape.cpp
473

Is this supposed to be ShapedType and ShapeType?

539–542

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

561–563

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

572

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

575–576

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

590

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

Thanks

mlir/lib/Dialect/Shape/IR/Shape.cpp
561–563

Rather being cautious.

590

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