[mlir][arith] Disallow zero ranked tensors for select's condition

[manas]

Summary

Zero ranked tensor (say tensor<i1>) when used for arith.select's condition,
crashes optimizer during bufferization. This patch puts a constraint on
condition to be either scalar or of matching shape as to its result.

[mehdi_amini]

I am not sure that the appropriate description is « book-like of non zero rank ».
To me it is rather: « scalar condition or matching shape » here.

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1385

Seems like this is just « no broadcasting » right?

[manas]

This test contains tensor<i1> and I am not sure if I should edit the CHECKs or remove the test altogether (as it will not check anything about linalg.generic, the reason why it was placed there).

[manas]

In D151270#4366445, @mehdi_amini wrote:

I am not sure that the appropriate description is « book-like of non zero rank ».
To me it is rather: « scalar condition or matching shape » here.

That makes sense. I will update.

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1385

Should I remove the suggestion and leave only the disallow info?

[manas]

Rename type to ScalarConditionOrMatchingShape

[manas]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1385

Not sure about no broadcasting, but for tensors I saw there is tensor.extract. So I mentioned according to that.

[mehdi_amini]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1385

The whole part about Zero ranked tensors (tensor<i1>) are disallowed as condition does not seem very accurate to me, or a bit too specific. It seems also unnecessary: this operation (like all the others in this dialect) don't do broadcasting (which would be the only way to make sense of this tensor).

The doc already says

The operation applies to vectors and tensors elementwise given the _shape_ of all operands is identical. [...] If an i1 is provided as the condition, the entire vector or tensor is chosen.

This seems enough to me to forbid rank-0 tensor when the shape does not match (that is a rank-zero tensor as condition seems still "correct" to me when the other operands are also rank-0).

[manas]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1385

Makes sense. Let me remove this from here.

[manas]

Remove irrelevant docs

[mehdi_amini]

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

1404

Now what we need is not a constraint here, the BoolLike was fine. Was we need is a constraint at the op level that spans multiple-operand, I suspect that may require crafting something somewhere in this section: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/OpBase.td#L2494-L2501 ; the exact construct likely does not exist but there should be plenty inspiration about how to design a constraint here.

[manas]

Add op type constraint over SelectOp operands

[mehdi_amini]

LG with some nits

mlir/include/mlir/IR/OpBase.td

1012

Is this used anywhere?

2531

I don't think "Match" makes sense in the name here?

[manas]

Do dead-code elimination

mlir/include/mlir/IR/OpBase.td

1012

Oops. It is the previous diff. Removed.

[mehdi_amini]

LGTM
(do you have commit access?)

[manas]

In D151270#4379692, @mehdi_amini wrote:

LGTM
(do you have commit access?)

I think I do. I will push it. Thanks a lot for reviewing this!

In D151270#4366447, @manas wrote:

This test contains tensor<i1> and I am not sure if I should edit the CHECKs or remove the test altogether (as it will not check anything about linalg.generic, the reason why it was placed there).

Btw what should I do about this failing test?

[mehdi_amini]

In D151270#4366447, @manas wrote:

This test contains tensor<i1> and I am not sure if I should edit the CHECKs or remove the test altogether (as it will not check anything about linalg.generic, the reason why it was placed there).

Btw what should I do about this failing test?

This test looks correct to me, if it fails then you have something incorrect with you check.

mlir/include/mlir/IR/OpBase.td

2686

It's not clear to me why you need Non0RankPred here?

[manas]

In D151270#4379785, @mehdi_amini wrote:

This test looks correct to me, if it fails then you have something incorrect with you check.

I believe we want to prevent any `tensor<i1> in the arith.select's condition.

mlir/include/mlir/IR/OpBase.td

2686

If the condition is not of scalar type, then we ought to prevent it from being a zero ranked type. And AllShapesMatch will match when every operand is of zero rank. For e.g.,

arith.select %cond %true %false : (tensor<i1>, tensor<i32>, tensor<i32>) ->  tensor<i32>

will also match AllShapesMatch.

[mehdi_amini]

In D151270#4379805, @manas wrote:

In D151270#4379785, @mehdi_amini wrote:

This test looks correct to me, if it fails then you have something incorrect with you check.

I believe we want to prevent any `tensor<i1> in the arith.select's condition.

Why would we? As mentioned before, the "no broadcast" rule is enough: that is "either scalar or shape matching" describe what this op should support, I don't quite get why rank-0 should be special? (the problem is indeed that rank-0 is special right now and I believe the point of this fix should be to make it not special)

mlir/include/mlir/IR/OpBase.td

2686

It's not clear to me why

[kuhar]

mlir/test/Dialect/Arith/invalid.mlir

761

Why would we want to disallow this? This seems like a perfectly valid use to me since both the condition and operands match in type/shape and number of elements.

Zero ranked tensor (say tensor<i1>) when used for arith.select's

condition, crashes optimizer during bufferization.

Can we fix the optimizer instead?

[manas]

In D151270#4379861, @mehdi_amini wrote:

In D151270#4379805, @manas wrote:

In D151270#4379785, @mehdi_amini wrote:

This test looks correct to me, if it fails then you have something incorrect with you check.

I believe we want to prevent any `tensor<i1> in the arith.select's condition.

Why would we? As mentioned before, the "no broadcast" rule is enough: that is "either scalar or shape matching" describe what this op should support, I don't quite get why rank-0 should be special? (the problem is indeed that rank-0 is special right now and I believe the point of this fix should be to make it not special)

I believe I mistook your and @springerm 's suggestions on discord as preventing tensor<i1> altogether. I see your point. But only putting "no broadcasting" constraints over condition will not solve the bufferization crash. So I propose that we keep this constraint (removing Non0RankPred ofc) and to fix the bufferization crash, I can put a check in SelecOp.bufferize() (or its caller) which will not invoke a bufferization on zero rank tensors.

mlir/test/Dialect/Arith/invalid.mlir

761

I see that fixing bufferizer will require to keep %cond: tensor<i1> as it and not try to bufferize it.

[mehdi_amini]

In D151270#4383826, @manas wrote:

I believe I mistook your and @springerm 's suggestions on discord as preventing tensor<i1> altogether. I see your point. But only putting "no broadcasting" constraints over condition will not solve the bufferization crash. So I propose that we keep this constraint (removing Non0RankPred ofc) and to fix the bufferization crash, I can put a check in SelecOp.bufferize() (or its caller) which will not invoke a bufferization on zero rank tensors.

The bufferization crash wouldn't be solved by this change anyway would it? Would bufferization handles %r = "arith.select"(%cond, %t, %f) : (tensor<42xi1>, tensor<42xi32>, tensor<42xi32>) -> tensor<42xi32> ?
`

[manas]

Remove Non0RankPred

[manas]

In D151270#4385283, @mehdi_amini wrote:

The bufferization crash wouldn't be solved by this change anyway would it? Would bufferization handles %r = "arith.select"(%cond, %t, %f) : (tensor<42xi1>, tensor<42xi32>, tensor<42xi32>) -> tensor<42xi32> ?
`

That's correct. I will patch a fix for the crash after this. Yes the bufferization will crash for above example too.

Right now, I have removed the Non0RankPred's usage, so arith.select %c %t %f : tensor<i1>, tensor<i32> are allowed as they were previously. CI should be happy.

[mehdi_amini]

LG then!

mlir/docs/Bufferization.md

 result, bufferization has to choose a memref buffer in which the result can be
 stored. It is always safe to allocate a brand new buffer, but such a
 bufferization strategy would be unacceptable for high-performance codegen. When
 choosing an already existing buffer, we must be careful not to accidentally
 overwrite data that is still needed later in the program.
 
 To simplify this problem, One-Shot Bufferize was designed for ops that are in
 *destination-passing style*. For every tensor result, such ops have a tensor
-operand, who's buffer could be for storing the result of the op in the absence
-of other conflicts. We call such tensor operands the *destination*.
+operand, whose buffer could be utilized for storing the result of the op in the
+absence of other conflicts. We call such tensor operands the *destination*.
 
 As an example, consider the following op: `%0 = tensor.insert %cst into
 %t[%idx] : tensor<?xf32>`
 
 `%t` is the destination in this example. When choosing a buffer for the result
 `%0`, One-Shot Bufferize considers only two options:
 
 1.  buffer(`%0`) = buffer(`%t`).

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

let description = [{
     The second and the third operand must have the same type.
 
     The operation applies to vectors and tensors elementwise given the _shape_
     of all operands is identical. The choice is made for each element
     individually based on the value at the same position as the element in the
     condition operand. If an i1 is provided as the condition, the entire vector
     or tensor is chosen.
 
+    Zero ranked tensors (`tensor<i1>`) are disallowed as condition. Such cases
+    should extract i1 first.
+
     Example:
 
     ```mlir
     // Custom form of scalar selection.
     %x = arith.select %cond, %true, %false : i32
 
     // Generic form of the same operation.
     %x = "arith.select"(%cond, %true, %false) : (i1, i32, i32) -> i32
 
     // Element-wise vector selection.
     %vx = arith.select %vcond, %vtrue, %vfalse : vector<42xi1>, vector<42xf32>
 
     // Full vector selection.
     %vx = arith.select %cond, %vtrue, %vfalse : vector<42xf32>
     ```
   }];
 
-  let arguments = (ins BoolLike:$condition,
+  let arguments = (ins ScalarConditionOrMatchingShape:$condition,
                        AnyType:$true_value,
                        AnyType:$false_value);
   let results = (outs AnyType:$result);
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 
   // FIXME: Switch this to use the declarative assembly format.
   let hasCustomAssemblyFormat = 1;
 }
 
 #endif // ARITH_OPS

mlir/include/mlir/IR/OpBase.td

 
 // Temporary constraint to allow gradual transition to supporting 0-D vectors.
 // TODO: Remove this when all ops support 0-D vectors.
 class TypeOrContainerOfAnyRank<Type allowedType, string name> : TypeConstraint<Or<[
   allowedType.predicate, VectorOfAnyRankOf<[allowedType]>.predicate,
   TensorOf<[allowedType]>.predicate]>,
   name>;
 
+// Temporary constraint to disallow zero ranked tensors in conditions in
+// arith.select
+class TypeOrContainerOfScalarConditionOrMatchingShape<Type allowedType,
+                                                      string name> :
+    TypeConstraint<Or<[
+  allowedType.predicate, VectorOfAnyRankOf<[allowedType]>.predicate,
+  Non0RankedTensorOf<[allowedType]>.predicate,
+  UnrankedTensorOf<[allowedType]>.predicate]>,
+  name>;
 
 // Type constraint for bool-like types: bools, vectors of bools, tensors of
 // bools.
 def BoolLike : TypeOrContainer<I1, "bool-like">;
 
 def BoolLikeOfAnyRank : TypeOrContainerOfAnyRank<I1, "bool-like">;
 
+// Similar to bool-like, but disallows zero ranked tensors
+def ScalarConditionOrMatchingShape :
+    TypeOrContainerOfScalarConditionOrMatchingShape<I1, "bool-like">;
+
 // Type constraint for signless-integer-like types: signless integers, indices,
 // vectors of signless integers or indices, tensors of signless integers.
 def SignlessIntegerLike : TypeOrContainer<AnySignlessIntegerOrIndex,
     "signless-integer-like">;
 
 def SignlessIntegerLikeOfAnyRank : TypeOrContainerOfAnyRank<
     AnySignlessIntegerOrIndex,
     "signless-integer-like">;
class Shape<string name> :
     StrFunc<"::llvm::cast<::mlir::ShapedType>($" # name # ".getType()).getShape()">;
 
 class ElementCount<string name> :
   StrFunc<"llvm::cast<::mlir::ShapedType>($" # name # ".getType())"
                                  ".getNumElements()">;
 
 class ElementType<string name> : StrFunc<"getElementTypeOrSelf($" # name # ")">;
 
 class AllMatchPred<list<string> values> :
   CPred<!if(!lt(!size(values), 2),
             "true",
             !foldl("(" # !head(values) # ")", !tail(values), acc, v,
                    acc # " == (" # v # ") && (" # v # ")")
               # " == (" # !head(values) # ")")>;
 
 class AllMatch<list<string> values, string summary> :
     PredOpTrait<summary, AllMatchPred<values>>;

 // 1) all operands involved are of shaped type and
 // 2) the indices are not out of range.
 class TCopVTEtAreSameAt<list<int> indices> : CPred<
   "::llvm::all_equal(::llvm::map_range("
       "::mlir::ArrayRef<unsigned>({" # !interleave(indices, ", ") # "}), "
       "[this](unsigned i) { return getElementTypeOrSelf(this->getOperand(i)); "
       "}))">;
 
 #endif // OP_BASE

mlir/test/Dialect/Arith/invalid.mlir

 
 // -----
 
 func.func @func() {
   %c0 = arith.constant  // expected-error {{expected attribute value}}
 
   %x = arith.constant 1 : i32
 }
+
+// -----
+
+func.func @disallow_zero_rank_tensor(%arg0 : tensor<i1>, %arg1 : tensor<i32>, %arg2 : tensor<i32>) -> tensor<i32> {
+  // expected-error @+1 {{'arith.select' op operand #0 must be bool-like, but got 'tensor<i1>'}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32>
+  return %0 : tensor<i32>
+}
+
+// -----
+
+func.func @disallow_zero_rank_tensor_with_ranked_tensor(%arg0 : tensor<i1>, %arg1 : tensor<2xi64>, %arg2 : tensor<2xi64>) -> tensor<2xi64> {
+  // expected-error @+1 {{'arith.select' op operand #0 must be bool-like, but got 'tensor<i1>'}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<2xi64>
+  return %0 : tensor<2xi64>
+}
+
+// -----
+
+func.func @disallow_zero_rank_tensor_with_unranked_tensor(%arg0 : tensor<i1>, %arg1 : tensor<2x?xi64>, %arg2 : tensor<2x?xi64>) -> tensor<2x?xi64> {
+  // expected-error @+1 {{'arith.select' op operand #0 must be bool-like, but got 'tensor<i1>'}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<2x?xi64>
+  return %0 : tensor<2x?xi64>
+}
+
+// -----
+
+func.func @disallow_zero_rank_tensor_with_vector(%arg0 : tensor<i1>, %arg1 : vector<i64>, %arg2 : vector<i64>) -> vector<i64> {
+  // expected-error @+1 {{'arith.select' op operand #0 must be bool-like, but got 'tensor<i1>'}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, vector<i64>
+  return %0 : vector<i64>
+}