[MLIR] Allow compatible shapes in `Elementwise` operations

[herhut]

mlir/lib/IR/Operation.cpp

1093

Do you need the getTypeID() here? Can you not compare types directly? Also, do the cheap comparison first maybe?

[silvas]

This is explicitly prohibited. Do not add this (see the rationale: https://github.com/llvm/llvm-project/blob/f0ccdde3c9ab23f997a4a30a0e9ab7d9435ec1c6/mlir/include/mlir/IR/OpDefinition.h#L1250).

[frgossen]

I don't see how the rationale prohibits that and I would like to at least understand why this is important and how to resolve it.

The way I understand constraint 3, it is asking for the *looser* shape equivalence where equal shapes may be differently but compatibly typed.
This is also the notion of shape equivalence that is implemented in, e.g., SameOperandsShape and subsequently in MHLO.
Point 3 says: "must be of the same shape. The shape may be dynamic in which case the op's behaviour is undefined for non-matching shapes."

The reason I would like this change is for HLO where element-wise operations are allowed to have differently concrete but equal shapes (there's an IREE test that that uses this).
The hard constraint of equal-equal shape types also makes rewriting quite difficult where you may want to concretise a type of an operand locally.

How does this affect folding?
For folding you will get attributes which will reveal the very concrete shape anyways, no?
"3. guarantees that folding can be done across scalars/vectors/tensors with the same pattern, as otherwise lots of special handling for type mismatches would be needed."

[silvas]

The key issue is:

%1 = "logical_not"(%0) : tensor<?xi1> -> tensor<?xi1>
%2 = "logical_not"(%1) : tensor<?xi1> -> tensor<3xi1>

With the proposed sematnics, we cannot fold this away by replacing %2 with %0 (e.g. suppose %2 is the return value of a function, or a block argument to a succesor, or iter_arg to scf.for). If we have tensor.cast ops explicitly, then we don't have this problem (and the tensor.cast ops can be eliminated by other means, such as commuting them with elementwise ops).

IMO HLO is broken in this regard. There have been bugs about this. HLO is originally statically shaped in XLA, so many aspects of the design of the dynamic shape handling of the hlo dialect in MLIR is somewhat arbitrary and not fully thought through (roughly "let's allow ? in the types" + cargo cult tf dialect behavior). Note: tf dialect is a very different (higher) abstraction level where it makes more sense to be loose like this, especially given the desire to round-trip.

SameOperandsShape and similar MLIR traits are also from the early days where we were just pattern matching on TF-isms and not fully thought through. For example, the following op satisfies SameOperandShape: "foo.op"(%0, %1, %2) : tensor<?xf32>, tensor<4xf32>, tensor<7xf32> (because verifyCompatibleShape is not transitive). Similar bug exists in SameOperandsAndResultShape. I don't see much difference between what you are trying to change Elementwise into and SameOperandsAndResultShape for instance -- why don't you just use SameOperandsAndResultShape for HLO ops? Without the scalarization traits (which HLO can't use) there isn't much value to Elementwise vs SameOperandsAndResultShape.

The hard constraint of equal-equal shape types also makes rewriting quite difficult where you may want to concretise a type of an operand locally.

Can you give an example?

[frgossen]

I am using SameOperandsAndResultShape already but am specifically interested in element-wise behaviour, to exclude e.g. transposition.
The goal is to move broadcasts over ops to allow for fusion which is only semantics-preserving for these ops.

The problem you describe is one that I ran into (and discussed) already.
My understanding was that the solution to this is to always allow for more specifically shaped types in all the ops.
While that is not the case for every op today, it would solve these issues.
The only remaining case is then return where an explicit cast is then needed (also solvable).

I see that stricter shape equivalence is another solution.
However, that seems to be a big barrier for folding in general:

%0 = constant [1, 2, 3] : tensor<3x..>
%1 = cast %0 : tensor<3x..> to tensor<?x..>

If what you describe is general consensus, then I should probably introduce an Elementwise trait specifically for HLO (or *fix* HLO).

[frgossen]

Address comment

[frgossen]

mlir/lib/IR/Operation.cpp

1093

getTypeID ensures that the operands are either all vectors or all tensors. Comparing the types directly would test for exact equality, not compatibility. Having the cheap comparison first makes sense :)

1093

This suffers from the same issue as the SameOperandsShape verifier. @tpopp is working on a fix which we can then also use here.

[silvas]

Really, it feels a lot better to convert Elementwise ops, broadcasting, etc. into linalg.generic ops, and much of the folding that I describe above can happen there trivially on the scalarized payload without any of these considerations.

Insofar as people want to use this to pursue non-linalg-based transformations like folding broadcasts like you described, and if this makes your life easier, then I guess go for it. I just don't think it's a good design within this space of ops on tensors NOT using linalg for these kinds of things (which I don't think itself is a good design space anyway). I just don't have the energy to argue it. All the flows that I care about really only care about using ElementwiseMappable to convert to linalg on tensors ASAP (which doesn't have the exact shape equality verification), so while I think this is a footgun (and I've hit the same issue with linalg ops too, btw), it doesn't really affect me.

All this is to say "I'm not going to block this", but if you want further discussion on this topic you may want a reviewer more invested in this design space to review the ramifications of allowing this and possibly changing HLO if necessary based on the outcome of that decision.

In D98186#2615849, @frgossen wrote:

The problem you describe is one that I ran into (and discussed) already.
My understanding was that the solution to this is to always allow for more specifically shaped types in all the ops.
While that is not the case for every op today, it would solve these issues.
The only remaining case is then return where an explicit cast is then needed (also solvable).

Just a note, but this change also allows less-specifically shaped types on the input, which can still fold to create less specifically shaped types. You may want to address that if you truly believe in the "folding to less specifically shaped types should be illegal".

Also, I'm not sure about the viability of "folding to more specifically shaped types". For example, that makes it illegal to have an op like "foo"(%0, %1, %2) : (tensor<?x4x?xf32>, index, index) -> () which has "alloc-like" behavior where each ? in the type corresponds to a new operand. In the above, folding that "foo" op's operand to tensor<?x4x4xf32> would then mess up that invariant and create an invalid op. We actually are finding that such ops are useful in IREE (see https://github.com/google/iree/pull/4881). I don't have the energy to have that design discussion right now, but it's not obvious to me that "folding to more specifically shaped types" is an IR property we want to guarantee is legal.

mlir/lib/IR/Operation.cpp

1093

Can you just fix it in this patch? it shouldn't be too difficult, and makes sense to introduce in this patch.

mlir/test/IR/traits.mlir

185–186

remove "failed" from th ename.

[frgossen]

Address comments

Updating D98186: [MLIR] Allow compatible shapes in Elementwise operations

[frgossen]

This comment was removed by frgossen.

[frgossen]

Address comment

mlir/lib/IR/Operation.cpp

if (region.getNumArguments() != 0) {
                << region.getRegionNumber() << " should have no arguments";
       else
         return op->emitOpError("region should have no arguments");
     }
   }
   return success();
 }
 
-/// Checks if two ShapedTypes are the same, ignoring the element type.
-static bool areSameShapedTypeIgnoringElementType(ShapedType a, ShapedType b) {
-  if (a.getTypeID() != b.getTypeID())
-    return false;
-  if (!a.hasRank())
-    return !b.hasRank();
-  return a.getShape() == b.getShape();
-}
-
 LogicalResult OpTrait::impl::verifyElementwise(Operation *op) {
   auto isMappableType = [](Type type) {
     return type.isa<VectorType, TensorType>();
   };
   auto resultMappableTypes = llvm::to_vector<1>(
       llvm::make_filter_range(op->getResultTypes(), isMappableType));
   auto operandMappableTypes = llvm::to_vector<2>(
       llvm::make_filter_range(op->getOperandTypes(), isMappableType));
LogicalResult OpTrait::impl::verifyElementwise(Operation *op) {
 
   if (resultMappableTypes.size() != op->getNumResults())
     return op->emitOpError(
         "if an operand is non-scalar, then all results must be non-scalar");
 
   auto mustMatchType = operandMappableTypes[0].cast<ShapedType>();
   for (auto type :
        llvm::concat<Type>(resultMappableTypes, operandMappableTypes)) {
-    if (!areSameShapedTypeIgnoringElementType(type.cast<ShapedType>(),
-                                              mustMatchType)) {
+    auto shapedTy = type.cast<ShapedType>();
+    if (shapedTy.getTypeID() != mustMatchType.getTypeID() ||
+        failed(verifyCompatibleShape(shapedTy, mustMatchType))) {
       return op->emitOpError() << "all non-scalar operands/results must have "
                                   "the same shape and base type: found "
                                << type << " and " << mustMatchType;
     }
   }
 
   return success();
 }

mlir/test/IR/traits.mlir

func @succeededSameOperandAndResultType(%t10x10 : tensor<10x10xf32>, %t1: tensor<1xf32>, %tr: tensor<*xf32>, %t1d: tensor<?xf32>, %i32 : i32) {
   return
 }
 
 // -----
 
 func @failedSameOperandAndResultType_operand_result_mismatch(%t10 : tensor<10xf32>, %t20 : tensor<20xf32>) {
   // expected-error@+1 {{requires the same type for all operands and results}}
   "test.same_operand_and_result_type"(%t10, %t20) : (tensor<10xf32>, tensor<20xf32>) -> tensor<10xf32>
+  return
 }
 
 // -----
 
 func @failedElementwiseMappable_different_rankedness(%arg0: tensor<?xf32>, %arg1: tensor<*xf32>) {
   // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<*xf32>' and 'tensor<?xf32>'}}
   %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<*xf32>) -> tensor<*xf32>
+  return
 }
 
 // -----
 
 func @failedElementwiseMappable_different_rank(%arg0: tensor<?xf32>, %arg1: tensor<?x?xf32>) {
   // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<?x?xf32>' and 'tensor<?xf32>'}}
   %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+  return
 }
 
 // -----
 
-func @failedElementwiseMappable_different_shape(%arg0: tensor<?xf32>, %arg1: tensor<5xf32>) {
-  // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<5xf32>' and 'tensor<?xf32>'}}
-  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (tensor<?xf32>, tensor<5xf32>) -> tensor<?xf32>
+func @elementwiseMappable_dynamic_shapes(%arg0: tensor<?xf32>,
+    %arg1: tensor<5xf32>) {
+  %0 = "test.elementwise_mappable"(%arg0, %arg1) :
+      (tensor<?xf32>, tensor<5xf32>) -> tensor<?xf32>
+  return
 }
 
 // -----
 
 func @failedElementwiseMappable_different_base_type(%arg0: vector<2xf32>, %arg1: tensor<2xf32>) {
   // expected-error@+1 {{all non-scalar operands/results must have the same shape and base type: found 'tensor<2xf32>' and 'vector<2xf32>'}}
   %0 = "test.elementwise_mappable"(%arg0, %arg1) : (vector<2xf32>, tensor<2xf32>) -> tensor<2xf32>
   return