Page MenuHomePhabricator

Add ElementwiseMappable trait and apply it to std elementwise ops.
ClosedPublic

Authored by silvas on Nov 3 2020, 4:56 PM.

Details

Summary

This patch adds an ElementwiseMappable trait as discussed in the RFC
here:
https://llvm.discourse.group/t/rfc-std-elementwise-ops-on-tensors/2113/23

This trait can power a number of transformations and analyses.
A subsequent patch adds a convert-elementwise-to-linalg pass exhibits
how this trait allows writing generic transformations.
See https://reviews.llvm.org/D90354 for that patch. I have split it out for clarity.

This trait slightly changes some verifier messages, but the diagnostics
are usually about as good. I fiddled with the ordering of the trait in
the .td file trait lists to minimize the changes here.

Diff Detail

Event Timeline

silvas created this revision.Nov 3 2020, 4:56 PM
silvas requested review of this revision.Nov 3 2020, 4:56 PM
silvas updated this revision to Diff 302725.Nov 3 2020, 6:25 PM

split out the convert-elementwise-to-linalg pass

silvas edited the summary of this revision. (Show Details)Nov 3 2020, 6:26 PM
silvas edited the summary of this revision. (Show Details)

@bondhugula I'd love to get your review on the definition of the ElementwiseMappable trait.

stellaraccident accepted this revision.Nov 8 2020, 4:44 PM
stellaraccident added a subscriber: stellaraccident.

This is very nice, and the corresponding RFC seems to have gotten to a nice conclusion.

mlir/include/mlir/IR/OpBase.td
1753–1754

I feel like this description could stand some truing up, at least having the words that make up the trait:

Op can be systematically interconverted between scalar and vector/tensor
form [by mapping elementwise based on the type].

mlir/include/mlir/IR/OpDefinition.h
1209

s/shall be same/shall be the same/

This revision is now accepted and ready to land.Nov 8 2020, 4:44 PM
silvas updated this revision to Diff 303982.Nov 9 2020, 1:42 PM
silvas marked 2 inline comments as done.
silvas edited the summary of this revision. (Show Details)

address comments

mehdi_amini accepted this revision.Nov 9 2020, 7:57 PM
mehdi_amini added inline comments.
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

How does SameOperandsAndResultShape relates to ElementwiseMappable? Is it a subset by any chance?

mlir/include/mlir/IR/OpDefinition.h
1265

Concrete instead of Concrent?

rriddle accepted this revision.Nov 9 2020, 9:45 PM

Please add documentation to https://mlir.llvm.org/docs/Traits/

mlir/lib/IR/Operation.cpp
1082

nit: isa is variadic: isa<TensorType, VectorType>().

bondhugula added inline comments.Nov 9 2020, 10:02 PM
mlir/include/mlir/IR/OpDefinition.h
1244

Side note: would ElementwiseMappableTypeInterface == isa<ShapedType>()? - the way it is now (dense shaped types). One could perhaps introduce a new type in the hierarchy sitting between ShapedType and Type to handle dense + sparse like the way ElementsAttribute works. What interface methods do you specifically see being added to ElementwiseMappableTypeInterface? If there are none, it could just sit in the type inheritance hierarchy?

silvas marked 2 inline comments as done.Nov 10 2020, 1:43 PM
silvas added inline comments.
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

ElementwiseMappable permits select %scalar_pred, %true_tensor, %false_tensor, but SameOperandsAndResultShape does not.

SameOperandsAndResultShape actually doesn't check for "same", only "compatible" (which includes broadcast compatible). So SameOperandsAndResultShape allows "foo"(%0) : tensor<2xf32> -> tensor<?xf32>, but ElementwiseMappable does not (as described in the rationale, this makes folding not generalize well from scalars/vectors to tensors).

So neither is strictly redundant from the other one in general. I think in this case, due to there being only a single operand/result, ElementwiseMappable imposes strictly stronger invariants.

I didn't want to touch the SameOperandsAndResultShape/SameOpearndsAndResultTypes situation in this patch; there's a lot to unwind there. See e.g. this thread https://llvm.discourse.group/t/sameoperandandresulttypes-name-confusing/1313/4

mlir/include/mlir/IR/OpDefinition.h
1244

ElementwiseMappableTypeInterface is not just ShapedType. In general, it could support things like "database columns" or abstract "streams".

So far, one interface method that seems like it will be needed is bool canSpeculateTraversal(). E.g. addf %lhs, %rhs : vector<4xf32> can be speculated with standard poison-like reasoning that LLVM uses. But addf %lhs, %rhs : tensor<?xf32> cannot be speculated, because the traversal itself might have undefined behavior if %lhs and %rhs dynamically have mismatching sizes.

silvas updated this revision to Diff 304322.Nov 10 2020, 1:45 PM

Address comments.

This revision was landed with ongoing or failed builds.Nov 10 2020, 1:46 PM
This revision was automatically updated to reflect the committed changes.
silvas added inline comments.Nov 10 2020, 2:14 PM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

Correction, SameOperandsAndResultShape doesn't allow broadcast compatible. But it does allow e.g. 2 vs ? like in my example.

mehdi_amini added inline comments.Nov 10 2020, 2:30 PM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

Basically what I was after was: "If an operation already defines ElementwiseMappable, what extra verification does it get from SameOperandsAndResultShape which isn't already covered" ?

silvas added inline comments.Nov 10 2020, 2:35 PM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

All I can think of is that SameOperandsAndResultShape would prohibit ops like select %scalar_pred, %true_tensor, %false_tensor.

mehdi_amini added inline comments.Nov 10 2020, 2:50 PM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

OK! Should we remove SameOperandsAndResultShape in cases like here then?

silvas added inline comments.Nov 10 2020, 3:17 PM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
3800

Will do.