[Mlir-commits] [mlir] [mlir][arith] Further clean up select op definition (PR #93358)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: Jakub Kuderski (kuhar)

<details>
<summary>Changes</summary>

* Improve the condition type requirement description ('scalar' -> signless i1), to match what is actually verified.
* Use the `I1` type predicate instead of `AnyBooleanTypeMatch`.

Related discussion: https://github.com/llvm/llvm-project/pull/93351#issuecomment-2130453233.

---
Full diff: https://github.com/llvm/llvm-project/pull/93358.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Arith/IR/ArithOps.td (+4-11) 
- (modified) mlir/test/Dialect/Arith/invalid.mlir (+2-2) 
- (modified) mlir/test/IR/invalid-ops.mlir (+2-2) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index ead52332e8eec..3981a46e1f072 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -1540,21 +1540,14 @@ def Arith_CmpFOp : Arith_CompareOp<"cmpf",
 // SelectOp
 //===----------------------------------------------------------------------===//
 
-class AnyBooleanTypeMatch<list<string> names> :
-    AnyMatchOperatorTrait<names, "$_self.getType().isSignlessInteger(1)",
-                          "scalar type">;
-
-class ScalarConditionOrMatchingShape<list<string> names> :
+class BooleanConditionOrMatchingShape<list<string> names> :
     PredOpTrait<
-        !head(names) # " is scalar or has matching shape",
-        Or<[AnyBooleanTypeMatch<[!head(names)]>.predicate,
-            AllShapesMatch<names>.predicate]>> {
-  list<string> values = names;
-}
+        !head(names) # " is signless i1 or has matching shape",
+        Or<[TypeIsPred<!head(names), I1>, AllShapesMatch<names>.predicate]>>;
 
 def SelectOp : Arith_Op<"select", [Pure,
     AllTypesMatch<["true_value", "false_value", "result"]>,
-    ScalarConditionOrMatchingShape<["condition", "result"]>,
+    BooleanConditionOrMatchingShape<["condition", "result"]>,
     DeclareOpInterfaceMethods<InferIntRangeInterface>,
   ] # ElementwiseMappable.traits> {
   let summary = "select operation";
diff --git a/mlir/test/Dialect/Arith/invalid.mlir b/mlir/test/Dialect/Arith/invalid.mlir
index ada849220bb83..815c77a749382 100644
--- a/mlir/test/Dialect/Arith/invalid.mlir
+++ b/mlir/test/Dialect/Arith/invalid.mlir
@@ -792,7 +792,7 @@ func.func @func() {
 // -----
 
 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 failed to verify that condition is scalar or has matching shape}}
+  // expected-error @+1 {{'arith.select' op failed to verify that condition is signless i1 or has matching shape}}
   %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<2xi64>
   return %0 : tensor<2xi64>
 }
@@ -800,7 +800,7 @@ func.func @disallow_zero_rank_tensor_with_ranked_tensor(%arg0 : tensor<i1>, %arg
 // -----
 
 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 failed to verify that condition is scalar or has matching shape}}
+  // expected-error @+1 {{'arith.select' op failed to verify that condition is signless i1 or has matching shape}}
   %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<2x?xi64>
   return %0 : tensor<2x?xi64>
 }
diff --git a/mlir/test/IR/invalid-ops.mlir b/mlir/test/IR/invalid-ops.mlir
index 6ca7035022adb..6672c8840ffde 100644
--- a/mlir/test/IR/invalid-ops.mlir
+++ b/mlir/test/IR/invalid-ops.mlir
@@ -82,7 +82,7 @@ func.func @func_with_ops(i1, i32, i64) {
 
 func.func @func_with_ops(vector<12xi1>, vector<42xi32>, vector<42xi32>) {
 ^bb0(%cond : vector<12xi1>, %t : vector<42xi32>, %f : vector<42xi32>):
-  // expected-error at +1 {{'arith.select' op failed to verify that condition is scalar or has matching shape}}
+  // expected-error at +1 {{'arith.select' op failed to verify that condition is signless i1 or has matching shape}}
   %r = "arith.select"(%cond, %t, %f) : (vector<12xi1>, vector<42xi32>, vector<42xi32>) -> vector<42xi32>
 }
 
@@ -90,7 +90,7 @@ func.func @func_with_ops(vector<12xi1>, vector<42xi32>, vector<42xi32>) {
 
 func.func @func_with_ops(tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) {
 ^bb0(%cond : tensor<12xi1>, %t : tensor<42xi32>, %f : tensor<42xi32>):
-  // expected-error at +1 {{'arith.select' op failed to verify that condition is scalar or has matching shape}}
+  // expected-error at +1 {{'arith.select' op failed to verify that condition is signless i1 or has matching shape}}
   %r = "arith.select"(%cond, %t, %f) : (tensor<12xi1>, tensor<42xi32>, tensor<42xi32>) -> tensor<42xi32>
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/93358