[mlir][LLVM] Add nsw and nuw flags to trunc

This implementation is based on the one already existing for the binary operations.

If the nuw keyword is present, and any of the truncated bits are non-zero, the result is a poison value. If the nsw keyword is present, and any of the truncated bits are not the same as the top bit of the truncation result, the result is a poison value.

Author: lfrenot

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -525,6 +525,23 @@ class LLVM_CastOpWithNNegFlag<string mnemonic, string instName, Type type,
   }];
 }
 
+class LLVM_CastOpWithOverflowFlag<string mnemonic, string instName, Type type,
+                  Type resultType, list<Trait> traits = []> :
+    LLVM_Op<mnemonic, !listconcat([Pure], [DeclareOpInterfaceMethods<IntegerOverflowFlagsInterface>], traits)>,
+    LLVM_Builder<"$res = builder.Create" # instName # "($arg, $_resultType, /*Name=*/\"\", op.hasNoUnsignedWrap(), op.hasNoSignedWrap());"> {
+  let arguments = (ins type:$arg, EnumProperty<"IntegerOverflowFlags", "", "IntegerOverflowFlags::none">:$overflowFlags);
+  let results = (outs resultType:$res);
+  let builders = [LLVM_OneResultOpBuilder];
+  let assemblyFormat = "$arg attr-dict `` custom<OverflowFlags>($overflowFlags) `:` type($arg) `to` type($res)";
+  string llvmInstName = instName;
+  string mlirBuilder = [{
+    auto op = $_builder.create<$_qualCppClassName>(
+      $_location, $_resultType, $arg);
+    moduleImport.setIntegerOverflowFlags(inst, op);
+    $res = op;
+  }];
+}
+
 def LLVM_BitcastOp : LLVM_CastOp<"bitcast", "BitCast", LLVM_AnyNonAggregate,
     LLVM_AnyNonAggregate, [DeclareOpInterfaceMethods<PromotableOpInterface>]> {
   let hasFolder = 1;
@@ -554,7 +571,7 @@ def LLVM_ZExtOp : LLVM_CastOpWithNNegFlag<"zext", "ZExt",
   let hasFolder = 1;
   let hasVerifier = 1;
 }
-def LLVM_TruncOp : LLVM_CastOp<"trunc", "Trunc",
+def LLVM_TruncOp : LLVM_CastOpWithOverflowFlag<"trunc", "Trunc",
                                LLVM_ScalarOrVectorOf<AnySignlessInteger>,
                                LLVM_ScalarOrVectorOf<AnySignlessInteger>>;
 def LLVM_SIToFPOp : LLVM_CastOp<"sitofp", "SIToFP",

mlir/test/Dialect/LLVMIR/roundtrip.mlir

@@ -338,6 +338,23 @@ func.func @nneg_casts(%arg0: i32, %arg1: i64, %arg2: vector<4xi32>,
   llvm.return
 }
 
+// CHECK-LABEL: @casts_overflow
+// CHECK-SAME: (%[[I32:.*]]: i32, %[[I64:.*]]: i64, %[[V4I32:.*]]: vector<4xi32>, %[[V4I64:.*]]: vector<4xi64>, %[[PTR:.*]]: !llvm.ptr)
+func.func @casts_overflow(%arg0: i32, %arg1: i64, %arg2: vector<4xi32>,
+            %arg3: vector<4xi64>, %arg4: !llvm.ptr) {
+// CHECK:  = llvm.trunc %[[I64]] overflow<nsw> : i64 to i56
+  %0 = llvm.trunc %arg1 overflow<nsw> : i64 to i56
+// CHECK:  = llvm.trunc %[[I64]] overflow<nuw> : i64 to i56
+  %1 = llvm.trunc %arg1 overflow<nuw> : i64 to i56
+// CHECK:  = llvm.trunc %[[I64]] overflow<nsw, nuw> : i64 to i56
+  %2 = llvm.trunc %arg1 overflow<nsw, nuw> : i64 to i56
+// CHECK:  = llvm.trunc %[[I64]] overflow<nsw, nuw> : i64 to i56
+  %3 = llvm.trunc %arg1 overflow<nuw, nsw> : i64 to i56
+// CHECK:  = llvm.trunc %[[V4I64]] overflow<nsw> : vector<4xi64> to vector<4xi56>
+  %4 = llvm.trunc %arg3 overflow<nsw> : vector<4xi64> to vector<4xi56>
+  llvm.return
+}
+
 // CHECK-LABEL: @vect
 func.func @vect(%arg0: vector<4xf32>, %arg1: i32, %arg2: f32, %arg3: !llvm.vec<2 x ptr>) {
 // CHECK:  = llvm.extractelement {{.*}} : vector<4xf32>

mlir/test/Target/LLVMIR/Import/nsw_nuw.ll

@@ -10,5 +10,7 @@ define void @intflag_inst(i64 %arg1, i64 %arg2) {
   %3 = mul nsw nuw i64 %arg1, %arg2
   ; CHECK: llvm.shl %{{.*}}, %{{.*}} overflow<nsw, nuw> : i64
   %4 = shl nuw nsw i64 %arg1, %arg2
+  ; CHECK: llvm.trunc %{{.*}} overflow<nsw> : i64 to i32
+  %5 = trunc nsw i64 %arg1 to i32
   ret void
 }

mlir/test/Target/LLVMIR/nsw_nuw.mlir

@@ -10,5 +10,7 @@ llvm.func @intflags_func(%arg0: i64, %arg1: i64) {
   %2 = llvm.mul %arg0, %arg1 overflow <nsw, nuw> : i64
   // CHECK: %{{.*}} = shl nuw nsw i64 %{{.*}}, %{{.*}}
   %3 = llvm.shl %arg0, %arg1 overflow <nsw, nuw> : i64
+  // CHECK: %{{.*}} = trunc nuw i64 %{{.*}} to i32
+  %4 = llvm.trunc %arg1 overflow<nuw> : i64 to i32
   llvm.return
 }