Avoid UB due to signed wrap around

Author: mgehre-amd

mlir/lib/Conversion/ArithToEmitC/ArithToEmitC.cpp

@@ -55,6 +55,50 @@ class ArithOpConversion final : public OpConversionPattern<ArithOp> {
   }
 };
 
+template <typename ArithOp, typename EmitCOp>
+class IntegerOpConversion 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, IndexType>(type)) {
+      return rewriter.notifyMatchFailure(op, "expected integer type");
+    }
+
+    Value lhs = adaptor.getLhs();
+    Value rhs = adaptor.getRhs();
+    Type arithmeticType = type;
+    if ((type.isSignlessInteger() || type.isSignedInteger()) &&
+        !bitEnumContainsAll(op.getOverflowFlags(),
+                            arith::IntegerOverflowFlags::nsw)) {
+      // If the C type is signed and the op doesn't guarantee "No Signed Wrap",
+      // we compute in unsigned integers to avoid UB.
+      arithmeticType = rewriter.getIntegerType(type.getIntOrFloatBitWidth(),
+                                               /*isSigned=*/false);
+    }
+    if (arithmeticType != type) {
+      lhs = rewriter.template create<emitc::CastOp>(op.getLoc(), arithmeticType,
+                                                    lhs);
+      rhs = rewriter.template create<emitc::CastOp>(op.getLoc(), arithmeticType,
+                                                    rhs);
+    }
+
+    Value result = rewriter.template create<EmitCOp>(op.getLoc(),
+                                                     arithmeticType, lhs, rhs);
+
+    if (arithmeticType != type) {
+      result =
+          rewriter.template create<emitc::CastOp>(op.getLoc(), type, result);
+    }
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
 class SelectOpConversion : public OpConversionPattern<arith::SelectOp> {
 public:
   using OpConversionPattern<arith::SelectOp>::OpConversionPattern;
@@ -96,9 +140,9 @@ void mlir::populateArithToEmitCPatterns(TypeConverter &typeConverter,
     ArithOpConversion<arith::DivFOp, emitc::DivOp>,
     ArithOpConversion<arith::MulFOp, emitc::MulOp>,
     ArithOpConversion<arith::SubFOp, emitc::SubOp>,
-    ArithOpConversion<arith::AddIOp, emitc::AddOp>,
-    ArithOpConversion<arith::MulIOp, emitc::MulOp>,
-    ArithOpConversion<arith::SubIOp, emitc::SubOp>,
+    IntegerOpConversion<arith::AddIOp, emitc::AddOp>,
+    IntegerOpConversion<arith::MulIOp, emitc::MulOp>,
+    IntegerOpConversion<arith::SubIOp, emitc::SubOp>,
     SelectOpConversion
   >(typeConverter, ctx);
   // clang-format on

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

@@ -37,19 +37,57 @@ func.func @arith_ops(%arg0: f32, %arg1: f32) {
 
 // -----
 
+// CHECK-LABEL: arith_integer_ops
 func.func @arith_integer_ops(%arg0: i32, %arg1: i32) {
-  // CHECK: emitc.add %arg0, %arg1 : (i32, i32) -> i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui3
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui3
+  // CHECK: %[[ADD:[^ ]*]] = emitc.add %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[ADD]] : ui32 to i3
   %0 = arith.addi %arg0, %arg1 : i32
-  // CHECK: emitc.sub %arg0, %arg1 : (i32, i32) -> i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui3
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui3
+  // CHECK: %[[SUB:[^ ]*]] = emitc.sub %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[SUB]] : ui32 to i3
   %1 = arith.subi %arg0, %arg1 : i32
-  // CHECK: emitc.mul %arg0, %arg1 : (i32, i32) -> i32
+  // CHECK: %[[C1:[^ ]*]] = emitc.cast %arg0 : i32 to ui3
+  // CHECK: %[[C2:[^ ]*]] = emitc.cast %arg1 : i32 to ui3
+  // CHECK: %[[MUL:[^ ]*]] = emitc.mul %[[C1]], %[[C2]] : (ui32, ui32) -> ui32
+  // CHECK: %[[C3:[^ ]*]] = emitc.cast %[[MUL]] : ui32 to i3
   %2 = arith.muli %arg0, %arg1 : i32
 
   return
 }
 
 // -----
 
+// CHECK-LABEL: arith_integer_ops_signed_nsw
+func.func @arith_integer_ops_signed_nsw(%arg0: i32, %arg1: i32) {
+  // CHECK: emitc.add %arg0, %arg1 : (i32, i32) -> i32
+  %0 = arith.addi %arg0, %arg1 overflow<nsw> : i32
+  // CHECK: emitc.sub %arg0, %arg1 : (i32, i32) -> i32
+  %1 = arith.subi %arg0, %arg1 overflow<nsw>  : i32
+  // CHECK: emitc.mul %arg0, %arg1 : (i32, i32) -> i32
+  %2 = arith.muli %arg0, %arg1 overflow<nsw> : i32
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: arith_index
+func.func @arith_index(%arg0: index, %arg1: index) {
+  // CHECK: emitc.add %arg0, %arg1 : (index, index) -> index
+  %0 = arith.addi %arg0, %arg1 : index
+  // CHECK: emitc.sub %arg0, %arg1 : (index, index) -> index
+  %1 = arith.subi %arg0, %arg1 : index
+  // CHECK: emitc.mul %arg0, %arg1 : (index, index) -> index
+  %2 = arith.muli %arg0, %arg1 : index
+
+  return
+}
+
+// -----
+
 func.func @arith_select(%arg0: i1, %arg1: tensor<8xi32>, %arg2: tensor<8xi32>) -> () {
   // CHECK: [[V0:[^ ]*]] = emitc.conditional %arg0, %arg1, %arg2 : tensor<8xi32>
   %0 = arith.select %arg0, %arg1, %arg2 : i1, tensor<8xi32>