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
472

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

491

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

tpopp added inline comments.Apr 16 2021, 7:16 AM
mlir/lib/Dialect/Shape/IR/Shape.cpp
472

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
472

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
453

Is this supposed to be ShapedType and ShapeType?

536–538

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

540–543

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

547

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

550–551

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

565

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
536–538

Rather being cautious.

565

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