[mlir] Saturation arithmetic intrinsics

[omri123]

Summary

Add the following saturation arithmetic llvm intrinsic ops: sadd.sat, uadd.sat, ssub.sat, usub.sat

[gysit]

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

162–163

nit: There are quite a lot of changes in the diff for things that seemingly did not change (e.g. the overflow ops or the memove memset tests). Could it be your editor inserts tabs instead of spaces?

187

Nice!

I have just landed a revision that makes the type constraints for llvm intrinsics more precise (https://reviews.llvm.org/D136360). Can you update your newly added intrinsics to have more accurate type constraints as well (e.g. AnySignlessInteger instead of just LLVM_Type)? I think they may be best defined by deriving from LLVM_BinarySameArgsIntrOpI since they take two arguments and return one result of the same type if I am not mistaken:

def LLVM_SMaxOp : LLVM_BinarySameArgsIntrOpI<"sadd.sat">;
def LLVM_SMaxOp : LLVM_BinarySameArgsIntrOpI<"ssub.sat">;

[omri123]

Thanks for your comments! In this revision I fixed the diff and updated the intrinsics to use LLVM_BinarySameArgsIntrOpI.

[gysit]

Looks good from my side modulo the nit comment. Let's wait for @ftynse (the code owner) for the final approval.

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

778

nit: the #0 seems unnecessary.

[gysit]

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

179

Ah I missed that one. The Op names need to differ of course LLVM_SAddSat, LLVM_UAddSat, etc.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

179

Indeed, this is breaking the tests.

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

809–816

Nit: I'd rather remove the #0 matadata because the 0 is transient. I see some cases above also match for that, but it is a mistake that makes tests fragile.

[omri123]

Thanks again! In this revision I fixed both issues mentioned above (removed #0 from tests and renamed the tblgen definitions)

[ftynse]

Please wait for the tests to pass before submitting.

[omri123]

Thanks! Submitting.

[gysit]

@omri123 are you interested in rebasing and landing this?

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_LifetimeEndOp : LLVM_LifetimeBaseOp<"lifetime.end"> {
   // Custom builder to convert the size argument to an attribute.
   string mlirBuilder = [{
     $_builder.create<LLVM::LifetimeEndOp>(
       $_location, $_int_attr($size), $ptr);
   }];
 }
 
 // Intrinsics with multiple returns.
 
 class LLVM_ArithWithOverflowOp<string mnem>
     : LLVM_IntrOp<mnem, [0], [], [SameOperandsElementType], 2>,
       Arguments<(ins LLVM_ScalarOrVectorOf<AnySignlessInteger>,
                  LLVM_ScalarOrVectorOf<AnySignlessInteger>)>;
 
 def LLVM_SAddWithOverflowOp : LLVM_ArithWithOverflowOp<"sadd.with.overflow">;
 def LLVM_UAddWithOverflowOp : LLVM_ArithWithOverflowOp<"uadd.with.overflow">;
 def LLVM_SSubWithOverflowOp : LLVM_ArithWithOverflowOp<"ssub.with.overflow">;
 def LLVM_USubWithOverflowOp : LLVM_ArithWithOverflowOp<"usub.with.overflow">;
 def LLVM_SMulWithOverflowOp : LLVM_ArithWithOverflowOp<"smul.with.overflow">;
 def LLVM_UMulWithOverflowOp : LLVM_ArithWithOverflowOp<"umul.with.overflow">;
 
+// Saturation Arithmetic Intrinsics.
+
+def LLVM_SAddSat : LLVM_BinarySameArgsIntrOpI<"sadd.sat">;
+def LLVM_UAddSat : LLVM_BinarySameArgsIntrOpI<"uadd.sat">;
+def LLVM_SSubSat : LLVM_BinarySameArgsIntrOpI<"ssub.sat">;
+def LLVM_USubSat : LLVM_BinarySameArgsIntrOpI<"usub.sat">;
+
 def LLVM_AssumeOp
   : LLVM_ZeroResultIntrOp<"assume", []>, Arguments<(ins I1:$cond)>;
 
 //
 // Coroutine intrinsics.
 //
 
 def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
   let arguments = (ins I32:$align,
                        LLVM_i8Ptr:$promise,
                        LLVM_i8Ptr:$coroaddr,
                        LLVM_i8Ptr:$fnaddrs);
   let assemblyFormat = "$align `,` $promise `,` $coroaddr `,` $fnaddrs"

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

 define void @sadd_with_overflow_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
   ; CHECK: "llvm.intr.sadd.with.overflow"(%{{.*}}, %{{.*}}) : (i32, i32) -> !llvm.struct<(i32, i1)>
   %5 = call { i32, i1 } @llvm.sadd.with.overflow.i32(i32 %0, i32 %1)
   ; CHECK: "llvm.intr.sadd.with.overflow"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> !llvm.struct<(vector<8xi32>, vector<8xi1>)>
   %6 = call { <8 x i32>, <8 x i1> } @llvm.sadd.with.overflow.v8i32(<8 x i32> %2, <8 x i32> %3)
   ret void
 }
 
+; CHECK-LABEL:  llvm.func @sadd_sat_test
+define void @sadd_sat_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
+  ; CHECK: "llvm.intr.sadd.sat"(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
+  %5 = call i32 @llvm.sadd.sat.i32(i32 %0, i32 %1)
+  ; CHECK: "llvm.intr.sadd.sat"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  %6 = call <8 x i32> @llvm.sadd.sat.v8i32(<8 x i32> %2, <8 x i32> %3)
+  ret void
+}
+
+; CHECK-LABEL:  llvm.func @uadd_sat_test
+define void @uadd_sat_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
+  ; CHECK: "llvm.intr.uadd.sat"(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
+  %5 = call i32 @llvm.uadd.sat.i32(i32 %0, i32 %1)
+  ; CHECK: "llvm.intr.uadd.sat"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  %6 = call <8 x i32> @llvm.uadd.sat.v8i32(<8 x i32> %2, <8 x i32> %3)
+  ret void
+}
+
+; CHECK-LABEL:  llvm.func @ssub_sat_test
+define void @ssub_sat_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
+  ; CHECK: "llvm.intr.ssub.sat"(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
+  %5 = call i32 @llvm.ssub.sat.i32(i32 %0, i32 %1)
+  ; CHECK: "llvm.intr.ssub.sat"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  %6 = call <8 x i32> @llvm.ssub.sat.v8i32(<8 x i32> %2, <8 x i32> %3)
+  ret void
+}
+
+; CHECK-LABEL:  llvm.func @usub_sat_test
+define void @usub_sat_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
+  ; CHECK: "llvm.intr.usub.sat"(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
+  %5 = call i32 @llvm.usub.sat.i32(i32 %0, i32 %1)
+  ; CHECK: "llvm.intr.usub.sat"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  %6 = call <8 x i32> @llvm.usub.sat.v8i32(<8 x i32> %2, <8 x i32> %3)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @uadd_with_overflow_test
 define void @uadd_with_overflow_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
   ; CHECK: "llvm.intr.uadd.with.overflow"(%{{.*}}, %{{.*}}) : (i32, i32) -> !llvm.struct<(i32, i1)>
   %5 = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %0, i32 %1)
   ; CHECK: "llvm.intr.uadd.with.overflow"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> !llvm.struct<(vector<8xi32>, vector<8xi1>)>
   %6 = call { <8 x i32>, <8 x i1> } @llvm.uadd.with.overflow.v8i32(<8 x i32> %2, <8 x i32> %3)
   ret void
 }

 declare { i32, i1 } @llvm.ssub.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.ssub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.usub.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.usub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.smul.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.smul.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.umul.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32(<8 x i32>, <8 x i32>)
+declare i32 @llvm.sadd.sat.i32(i32, i32)
+declare <8 x i32> @llvm.sadd.sat.v8i32(<8 x i32>, <8 x i32>)
+declare i32 @llvm.uadd.sat.i32(i32, i32)
+declare <8 x i32> @llvm.uadd.sat.v8i32(<8 x i32>, <8 x i32>)
+declare i32 @llvm.ssub.sat.i32(i32, i32)
+declare <8 x i32> @llvm.ssub.sat.v8i32(<8 x i32>, <8 x i32>)
+declare i32 @llvm.usub.sat.i32(i32, i32)
+declare <8 x i32> @llvm.usub.sat.v8i32(<8 x i32>, <8 x i32>)
 declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
 declare i8* @llvm.coro.begin(token, i8* writeonly)
 declare i64 @llvm.coro.size.i64()
 declare i32 @llvm.coro.size.i32()
 declare i64 @llvm.coro.align.i64()
 declare i32 @llvm.coro.align.i32()
 declare token @llvm.coro.save(i8*)
 declare i8 @llvm.coro.suspend(token, i1)

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

 llvm.func @umul_with_overflow_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
   // CHECK: call { i32, i1 } @llvm.umul.with.overflow.i32
   "llvm.intr.umul.with.overflow"(%arg0, %arg1) : (i32, i32) -> !llvm.struct<(i32, i1)>
   // CHECK: call { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32
   "llvm.intr.umul.with.overflow"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> !llvm.struct<(vector<8xi32>, vector<8xi1>)>
   llvm.return
 }
 
+// CHECK-LABEL: @sadd_sat_test
+llvm.func @sadd_sat_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
+  // CHECK: call i32 @llvm.sadd.sat.i32
+  "llvm.intr.sadd.sat"(%arg0, %arg1) : (i32, i32) -> i32
+  // CHECK: call <8 x i32> @llvm.sadd.sat.v8i32
+  "llvm.intr.sadd.sat"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  llvm.return
+}
+
+// CHECK-LABEL: @uadd_sat_test
+llvm.func @uadd_sat_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
+  // CHECK: call i32 @llvm.uadd.sat.i32
+  "llvm.intr.uadd.sat"(%arg0, %arg1) : (i32, i32) -> i32
+  // CHECK: call <8 x i32> @llvm.uadd.sat.v8i32
+  "llvm.intr.uadd.sat"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  llvm.return
+}
+
+// CHECK-LABEL: @ssub_sat_test
+llvm.func @ssub_sat_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
+  // CHECK: call i32 @llvm.ssub.sat.i32
+  "llvm.intr.ssub.sat"(%arg0, %arg1) : (i32, i32) -> i32
+  // CHECK: call <8 x i32> @llvm.ssub.sat.v8i32
+  "llvm.intr.ssub.sat"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  llvm.return
+}
+
+// CHECK-LABEL: @usub_sat_test
+llvm.func @usub_sat_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
+  // CHECK: call i32 @llvm.usub.sat.i32
+  "llvm.intr.usub.sat"(%arg0, %arg1) : (i32, i32) -> i32
+  // CHECK: call <8 x i32> @llvm.usub.sat.v8i32
+  "llvm.intr.usub.sat"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
+  llvm.return
+}
+
 // CHECK-LABEL: @coro_id
 llvm.func @coro_id(%arg0: i32, %arg1: !llvm.ptr<i8>) {
   // CHECK: call token @llvm.coro.id
   %null = llvm.mlir.null : !llvm.ptr<i8>
   llvm.intr.coro.id %arg0, %arg1, %arg1, %null : !llvm.token
   llvm.return
 }
 

 // CHECK-DAG: declare { i32, i1 } @llvm.uadd.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.uadd.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.ssub.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.ssub.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.usub.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.usub.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
 // CHECK-DAG: declare { i32, i1 } @llvm.umul.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32(<8 x i32>, <8 x i32>) #0
+// CHECK-DAG: declare i32 @llvm.sadd.sat.i32(i32, i32)
+// CHECK-DAG: declare <8 x i32> @llvm.sadd.sat.v8i32(<8 x i32>, <8 x i32>)
+// CHECK-DAG: declare i32 @llvm.uadd.sat.i32(i32, i32)
+// CHECK-DAG: declare <8 x i32> @llvm.uadd.sat.v8i32(<8 x i32>, <8 x i32>)
+// CHECK-DAG: declare i32 @llvm.ssub.sat.i32(i32, i32)
+// CHECK-DAG: declare <8 x i32> @llvm.ssub.sat.v8i32(<8 x i32>, <8 x i32>)
+// CHECK-DAG: declare i32 @llvm.usub.sat.i32(i32, i32)
+// CHECK-DAG: declare <8 x i32> @llvm.usub.sat.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
 // CHECK-DAG: declare ptr @llvm.coro.begin(token, ptr writeonly)
 // CHECK-DAG: declare i64 @llvm.coro.size.i64()
 // CHECK-DAG: declare i32 @llvm.coro.size.i32()
 // CHECK-DAG: declare token @llvm.coro.save(ptr)
 // CHECK-DAG: declare i8 @llvm.coro.suspend(token, i1)
 // CHECK-DAG: declare i1 @llvm.coro.end(ptr, i1)
 // CHECK-DAG: declare ptr @llvm.coro.free(token, ptr nocapture readonly)