[SYCL-MLIR] Fix SelectI1Ext canonicalize pattern (#8290)

For the code below, `foo(true)` should return `1`, and `foo(false)`
should return `2`.
```
func.func @foo(%arg0: i1) -> i32 {
  %c1_i32 = arith.constant 1 : i32
  %c2_i32 = arith.constant 2 : i32
  %0 = arith.select %arg0, %c1_i32, %c2_i32 : i32
  return %0 : i32
}
```
Without this patch, the `SelectI1Ext` canonicalization pattern
(incorrectly) would change it to the following code:
```
  func.func @foo(%arg0: i1) -> i32 {
    %0 = arith.extui %arg0 : i1 to i32
    return %0 : i32
  }
```
The canonicalization pattern is only legal if the 2 operands of the
select operation are representable by a `i1` type (so only for the
constant 0 and 1).

Signed-off-by: Tsang, Whitney <[email protected]>

Author: whitneywhtsang

polygeist/lib/polygeist/Ops.cpp

@@ -2045,33 +2045,33 @@ class SelectI1Ext final : public OpRewritePattern<arith::SelectOp> {
 
   LogicalResult matchAndRewrite(arith::SelectOp op,
                                 PatternRewriter &rewriter) const override {
-    auto ty = op.getType().dyn_cast<IntegerType>();
-    if (!ty)
-      return failure();
-    if (ty.getWidth() == 1)
-      return failure();
-    IntegerAttr lhs, rhs;
-    Value lhs_v = nullptr, rhs_v = nullptr;
-    if (auto ext = op.getTrueValue().getDefiningOp<arith::ExtUIOp>()) {
-      lhs_v = ext.getIn();
-      if (lhs_v.getType().cast<IntegerType>().getWidth() != 1)
-        return failure();
-    } else if (matchPattern(op.getTrueValue(), m_Constant(&lhs))) {
-    } else
+    // Cannot extui i1 to i1, or i1 to f32
+    if (!op.getType().isa<IntegerType>() || op.getType().isInteger(1))
       return failure();
 
-    if (auto ext = op.getFalseValue().getDefiningOp<arith::ExtUIOp>()) {
-      rhs_v = ext.getIn();
-      if (rhs_v.getType().cast<IntegerType>().getWidth() != 1)
-        return failure();
-    } else if (matchPattern(op.getFalseValue(), m_Constant(&rhs))) {
-    } else
-      return failure();
+    // Determines whether the given value fits into a boolean type.
+    auto getI1 = [&op, &rewriter](Value val) -> Value {
+      constexpr int typeWidth = 1;
+      if (matchPattern(val, m_Op<arith::ExtUIOp>())) {
+        Value result = val.getDefiningOp()->getOperand(0);
+        if (result.getType().cast<IntegerType>().isInteger(typeWidth))
+          return result;
+      }
+
+      IntegerAttr intAttr;
+      if (matchPattern(val, m_Constant(&intAttr))) {
+        if (intAttr.getInt() == 0 || intAttr.getInt() == 1)
+          return rewriter.create<ConstantIntOp>(op.getLoc(), intAttr.getInt(),
+                                                typeWidth);
+      }
 
-    if (!lhs_v)
-      lhs_v = rewriter.create<ConstantIntOp>(op.getLoc(), lhs.getInt(), 1);
-    if (!rhs_v)
-      rhs_v = rewriter.create<ConstantIntOp>(op.getLoc(), rhs.getInt(), 1);
+      return nullptr;
+    };
+
+    Value lhs_v = getI1(op.getTrueValue());
+    Value rhs_v = getI1(op.getFalseValue());
+    if (!lhs_v || !rhs_v)
+      return failure();
 
     rewriter.replaceOpWithNewOp<ExtUIOp>(
         op, op.getType(),

polygeist/test/polygeist-opt/selecti1ext.mlir

@@ -0,0 +1,28 @@
+// RUN: polygeist-opt --canonicalize %s | FileCheck %s 
+  
+// CHECK-LABEL: func.func @test1(%arg0: i1) -> i32 {
+// CHECK:         %0 = arith.select %arg0, %c1_i32, %c2_i32 : i32
+// CHECK-NEXT:    return %0 : i32
+// CHECK-NEXT:  }
+
+func.func @test1(%arg0: i1) -> i32 {
+  %c1_i32 = arith.constant 1 : i32
+  %c2_i32 = arith.constant 2 : i32
+  %0 = arith.select %arg0, %c1_i32, %c2_i32 : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: func.func @test2(%arg0: i1) -> i32 {
+// CHECK:         %0 = arith.xori %arg0, %true : i1
+// CHECK-NEXT:    %1 = arith.extui %0 : i1 to i32
+// CHECK-NEXT:    return %1 : i32
+// CHECK-NEXT:  }
+
+func.func @test2(%arg0: i1) -> i32 {
+  %true = arith.constant true
+  %false = arith.constant false
+  %0 = arith.extui %true : i1 to i32
+  %1 = arith.extui %false : i1 to i32
+  %2 = arith.select %arg0, %1, %0 : i32
+  return %2 : i32
+}

polygeist/tools/cgeist/Test/Verification/sycl/structvec.cpp

@@ -79,7 +79,9 @@ SYCL_EXTERNAL structvec test_store(structvec sv, int idx, char el) {
 // CHECK-NEXT:  }
 
 // CHECK-LABEL: func.func @_ZN9structvecC1ESt16initializer_listIcE(%arg0: !llvm.ptr<struct<(vector<2xi8>)>, 4> {llvm.align = 2 : i64, llvm.dereferenceable_or_null = 2 : i64, llvm.noundef}, %arg1: !llvm.ptr<!llvm.struct<(memref<?xi8, 4>, i64)>> {llvm.align = 8 : i64, llvm.byval = !llvm.struct<(memref<?xi8, 4>, i64)>, llvm.noundef})
-// CHECK-NEXT:    %c0_i8 = arith.constant 0 : i8
+// CHECK-DAG:     %c-1_i32 = arith.constant -1 : i32
+// CHECK-DAG:     %c0_i32 = arith.constant 0 : i32
+// CHECK-DAG:     %c0_i8 = arith.constant 0 : i8
 // CHECK-NEXT:    %0 = llvm.addrspacecast %arg1 : !llvm.ptr<!llvm.struct<(memref<?xi8, 4>, i64)>> to !llvm.ptr<!llvm.struct<(memref<?xi8, 4>, i64)>, 4>
 // CHECK-NEXT:    %1 = llvm.getelementptr %arg0[0, 0] : (!llvm.ptr<struct<(vector<2xi8>)>, 4>) -> !llvm.ptr<vector<2xi8>, 4>
 // CHECK-NEXT:    affine.for %arg2 = 0 to 2 {
@@ -88,7 +90,7 @@ SYCL_EXTERNAL structvec test_store(structvec sv, int idx, char el) {
 // CHECK-NEXT:      %4 = arith.index_castui %2 : i32 to index
 // CHECK-NEXT:      %5 = memref.load %3[%4] : memref<?xi8, 4>
 // CHECK-NEXT:      %6 = arith.cmpi ne, %5, %c0_i8 : i8
-// CHECK-NEXT:      %7 = arith.extui %6 : i1 to i32
+// CHECK-NEXT:      %7 = arith.select %6, %c-1_i32, %c0_i32 : i32
 // CHECK-NEXT:      %8 = arith.trunci %7 : i32 to i8
 // CHECK-NEXT:      %9 = llvm.load %1 : !llvm.ptr<vector<2xi8>, 4>
 // CHECK-NEXT:      %10 = vector.insertelement %8, %9[%2 : i32] : vector<2xi8>