Support more Arith integer binary ops

Author: cferry-AMD

mlir/lib/Conversion/ArithToEmitC/ArithToEmitC.cpp

@@ -443,6 +443,46 @@ class IntegerOpConversion final : public OpConversionPattern<ArithOp> {
   }
 };
 
+template <typename ArithOp, typename EmitCOp, bool booleansLegal>
+class BitwiseOpConversion final : public OpConversionPattern<ArithOp> {
+public:
+  using OpConversionPattern<ArithOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ArithOp op, typename ArithOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    Type type = this->getTypeConverter()->convertType(op.getType());
+    if (!isa_and_nonnull<IntegerType, emitc::SignedSizeTType, emitc::SizeTType>(
+            type)) {
+      return rewriter.notifyMatchFailure(
+          op, "expected integer or size_t/ssize_t type");
+    }
+
+    if (type.isInteger(1)) {
+      if (!booleansLegal)
+        return rewriter.notifyMatchFailure(op, "i1 type is not implemented");
+
+      rewriter.replaceOpWithNewOp<EmitCOp>(op, type, adaptor.getLhs(),
+                                           adaptor.getRhs());
+      return success();
+    }
+
+    Type arithmeticType = adaptIntegralTypeSignedness(type, true);
+
+    Value lhs = adaptValueType(adaptor.getLhs(), rewriter, arithmeticType);
+    Value rhs = adaptValueType(adaptor.getRhs(), rewriter, arithmeticType);
+
+    Value arithmeticResult = rewriter.template create<EmitCOp>(
+        op.getLoc(), arithmeticType, lhs, rhs);
+
+    Value result = adaptValueType(arithmeticResult, rewriter, type);
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
 class SelectOpConversion : public OpConversionPattern<arith::SelectOp> {
 public:
   using OpConversionPattern<arith::SelectOp>::OpConversionPattern;
@@ -581,6 +621,11 @@ void mlir::populateArithToEmitCPatterns(TypeConverter &typeConverter,
     IntegerOpConversion<arith::AddIOp, emitc::AddOp>,
     IntegerOpConversion<arith::MulIOp, emitc::MulOp>,
     IntegerOpConversion<arith::SubIOp, emitc::SubOp>,
+    BitwiseOpConversion<arith::AndIOp, emitc::BitwiseAndOp, true>,
+    BitwiseOpConversion<arith::OrIOp, emitc::BitwiseOrOp, true>,
+    BitwiseOpConversion<arith::XOrIOp, emitc::BitwiseXorOp, true>,
+    BitwiseOpConversion<arith::ShLIOp, emitc::BitwiseLeftShiftOp, false>,
+    BitwiseOpConversion<arith::ShRSIOp, emitc::BitwiseRightShiftOp, false>,
     CmpFOpConversion,
     CmpIOpConversion,
     SelectOpConversion,

mlir/test/Conversion/ArithToEmitC/arith-to-emitc.mlir

@@ -92,6 +92,54 @@ func.func @arith_index(%arg0: i32, %arg1: i32) {
   return
 }
 
+// -----
+
+// CHECK-LABEL: arith_bitwise
+func.func @arith_bitwise(%arg0: i32, %arg1: i32) {
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui32
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui32
+  // CHECK: %[[AND:[^ ]*]] = emitc.bitwise_and %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[AND]] : ui32 to i32
+  %5 = arith.andi %arg0, %arg1 : i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui32
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui32
+  // CHECK: %[[OR:[^ ]*]] = emitc.bitwise_or %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[OR]] : ui32 to i32
+  %6 = arith.ori %arg0, %arg1 : i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui32
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui32
+  // CHECK: %[[XOR:[^ ]*]] = emitc.bitwise_xor %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[XOR]] : ui32 to i32
+  %7 = arith.xori %arg0, %arg1 : i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui32
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui32
+  // CHECK: %[[SHL:[^ ]*]] = emitc.bitwise_left_shift %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[SHL]] : ui32 to i32
+  %8 = arith.shli %arg0, %arg1 : i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui32
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui32
+  // CHECK: %[[SHR:[^ ]*]] = emitc.bitwise_right_shift %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[SHR]] : ui32 to i32
+  %9 = arith.shrsi %arg0, %arg1 : i32
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: arith_bitwise_bool
+func.func @arith_bitwise_bool(%arg0: i1, %arg1: i1) {
+  // CHECK: %[[AND:[^ ]*]] = emitc.bitwise_and %arg0, %arg1 : (i1, i1) -> i1
+  %5 = arith.andi %arg0, %arg1 : i1
+  // CHECK: %[[OR:[^ ]*]] = emitc.bitwise_or %arg0, %arg1 : (i1, i1) -> i1
+  %6 = arith.ori %arg0, %arg1 : i1
+  // CHECK: %[[xor:[^ ]*]] = emitc.bitwise_xor %arg0, %arg1 : (i1, i1) -> i1
+  %7 = arith.xori %arg0, %arg1 : i1
+
+  return
+}
+
+
 // -----
 
 func.func @arith_select(%arg0: i1, %arg1: tensor<8xi32>, %arg2: tensor<8xi32>) -> () {