Implement the UAddc DXIL op to lower AddUint64 to DXIL

Author: Icohedron

llvm/lib/Target/DirectX/DXIL.td

@@ -50,6 +50,7 @@ def HandleTy : DXILOpParamType;
 def ResBindTy : DXILOpParamType;
 def ResPropsTy : DXILOpParamType;
 def SplitDoubleTy : DXILOpParamType;
+def BinaryWithCarryTy : DXILOpParamType;
 
 class DXILOpClass;
 
@@ -738,6 +739,18 @@ def UMin : DXILOp<40, binary> {
   let attributes = [Attributes<DXIL1_0, [ReadNone]>];
 }
 
+def UAddc : DXILOp<44, binaryWithCarryOrBorrow > {
+  let Doc = "Unsigned 32-bit integer arithmetic add with carry. uaddc(a,b) = (a+b, a+b overflowed ? 1 : 0)";
+  // TODO: This `let intrinsics = ...` line may be uncommented when 
+  // https://github.com/llvm/llvm-project/issues/113192 is fixed
+  // let intrinsics = [IntrinSelect<int_uadd_with_overflow>];
+  let arguments = [OverloadTy, OverloadTy];
+  let result = BinaryWithCarryTy;
+  let overloads = [Overloads<DXIL1_0, [Int32Ty]>];
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+  let attributes = [Attributes<DXIL1_0, [ReadNone]>];
+}
+
 def FMad : DXILOp<46, tertiary> {
   let Doc = "Floating point arithmetic multiply/add operation. fmad(m,a,b) = m "
             "* a + b.";

llvm/lib/Target/DirectX/DXILOpBuilder.cpp

@@ -230,6 +230,14 @@ static StructType *getSplitDoubleType(LLVMContext &Context) {
   return StructType::create({Int32Ty, Int32Ty}, "dx.types.splitdouble");
 }
 
+static StructType *getBinaryWithCarryType(LLVMContext &Context) {
+  if (auto *ST = StructType::getTypeByName(Context, "dx.types.i32c"))
+    return ST;
+  Type *Int32Ty = Type::getInt32Ty(Context);
+  Type *Int1Ty= Type::getInt1Ty(Context);
+  return StructType::create({Int32Ty, Int1Ty}, "dx.types.i32c");
+}
+
 static Type *getTypeFromOpParamType(OpParamType Kind, LLVMContext &Ctx,
                                     Type *OverloadTy) {
   switch (Kind) {
@@ -273,6 +281,8 @@ static Type *getTypeFromOpParamType(OpParamType Kind, LLVMContext &Ctx,
     return getResPropsType(Ctx);
   case OpParamType::SplitDoubleTy:
     return getSplitDoubleType(Ctx);
+  case OpParamType::BinaryWithCarryTy:
+    return getBinaryWithCarryType(Ctx);
   }
   llvm_unreachable("Invalid parameter kind");
   return nullptr;
@@ -539,6 +549,10 @@ StructType *DXILOpBuilder::getSplitDoubleType(LLVMContext &Context) {
   return ::getSplitDoubleType(Context);
 }
 
+StructType *DXILOpBuilder::getBinaryWithCarryType(LLVMContext &Context) {
+  return ::getBinaryWithCarryType(Context);
+}
+
 StructType *DXILOpBuilder::getHandleType() {
   return ::getHandleType(IRB.getContext());
 }

llvm/lib/Target/DirectX/DXILOpBuilder.h

@@ -53,6 +53,9 @@ class DXILOpBuilder {
   /// Get the `%dx.types.splitdouble` type.
   StructType *getSplitDoubleType(LLVMContext &Context);
 
+  /// Get the `%dx.types.i32c` type.
+  StructType *getBinaryWithCarryType(LLVMContext &Context);
+
   /// Get the `%dx.types.Handle` type.
   StructType *getHandleType();
 

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -359,17 +359,16 @@ class OpLowerer {
     return lowerToBindAndAnnotateHandle(F);
   }
 
-  Error replaceSplitDoubleCallUsages(CallInst *Intrin, CallInst *Op) {
+  Error replaceExtractElementTypeOfCallUsages(CallInst *Intrin, CallInst *Op) {
     for (Use &U : make_early_inc_range(Intrin->uses())) {
       if (auto *EVI = dyn_cast<ExtractValueInst>(U.getUser())) {
 
         if (EVI->getNumIndices() != 1)
-          return createStringError(std::errc::invalid_argument,
-                                   "Splitdouble has only 2 elements");
+          return createStringError(std::errc::invalid_argument, (std::string(Intrin->getOpcodeName()) + " has only 2 elements").c_str());
         EVI->setOperand(0, Op);
       } else {
         return make_error<StringError>(
-            "Splitdouble use is not ExtractValueInst",
+            (std::string(Intrin->getOpcodeName()) + " use is not ExtractValueInst").c_str(),
             inconvertibleErrorCode());
       }
     }
@@ -821,7 +820,17 @@ class OpLowerer {
             F, OpCode::SplitDouble,
             OpBuilder.getSplitDoubleType(M.getContext()),
             [&](CallInst *CI, CallInst *Op) {
-              return replaceSplitDoubleCallUsages(CI, Op);
+              return replaceExtractElementTypeOfCallUsages(CI, Op);
+            });
+        break;
+      // TODO: this can be removed when
+      // https://github.com/llvm/llvm-project/issues/113192 is fixed
+      case Intrinsic::uadd_with_overflow:
+        HasErrors |= replaceFunctionWithNamedStructOp(
+            F, OpCode::UAddc,
+            OpBuilder.getBinaryWithCarryType(M.getContext()),
+            [&](CallInst *CI, CallInst *Op) {
+              return replaceExtractElementTypeOfCallUsages(CI, Op);
             });
         break;
       case Intrinsic::ctpop:

llvm/test/CodeGen/DirectX/UAddc.ll

@@ -0,0 +1,25 @@
+; RUN: opt -S -dxil-op-lower -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s
+
+; CHECK: %dx.types.i32c = type { i32, i1 }
+
+define noundef i32 @test_UAddc(i32 noundef %a, i32 noundef %b) {
+; CHECK-LABEL: define noundef i32 @test_UAddc(
+; CHECK-SAME: i32 noundef [[A:%.*]], i32 noundef [[B:%.*]]) {
+; CHECK-NEXT:    [[UAddc:%.*]] = call %dx.types.i32c @dx.op.binaryWithCarryOrBorrow.i32(i32 44, i32 [[A]], i32 [[B]])
+; CHECK-NEXT:    [[Carry:%.*]] = extractvalue %dx.types.i32c [[UAddc]], 1
+; CHECK-NEXT:    [[Sum:%.*]] = extractvalue %dx.types.i32c [[UAddc]], 0
+; CHECK-NEXT:    [[CarryZExt:%.*]] = zext i1 [[Carry]] to i32
+; CHECK-NEXT:    [[Result:%.*]] = add i32 [[Sum]], [[CarryZExt]]
+; CHECK-NEXT:    ret i32 [[Result]]
+; 
+  %uaddc = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %a, i32 %b)
+  %carry = extractvalue { i32, i1 } %uaddc, 1
+  %sum = extractvalue { i32, i1 } %uaddc, 0
+  %carry_zext = zext i1 %carry to i32
+  %result = add i32 %sum, %carry_zext
+  ret i32 %result
+}
+
+declare { i32, i1 } @llvm.uadd.with.overflow.i32(i32, i32)
+; CHECK: declare %dx.types.i32c @dx.op.binaryWithCarryOrBorrow.i32(i32, i32, i32)
+

llvm/test/CodeGen/DirectX/UAddc_errors.ll

@@ -0,0 +1,17 @@
+; RUN: not opt -S -dxil-op-lower -mtriple=dxil-pc-shadermodel6.3-library %s 2>&1 | FileCheck %s
+
+; DXIL operation UAddc only supports i32. Other integer types are unsupported.
+; CHECK: in function uaddc_i16
+; CHECK-SAME: Cannot create UAddc operation: Invalid overload type
+
+define noundef i16 @uaddc_i16(i16 noundef %a, i16 noundef %b) {
+  %uaddc = call { i16, i1 } @llvm.uadd.with.overflow.i16(i16 %a, i16 %b)
+  %carry = extractvalue { i16, i1 } %uaddc, 1
+  %sum = extractvalue { i16, i1 } %uaddc, 0
+  %carry_zext = zext i1 %carry to i16
+  %result = add i16 %sum, %carry_zext
+  ret i16 %result
+}
+
+declare { i16, i1 } @llvm.uadd.with.overflow.i16(i16, i16)
+