[SPIR-V] Implement builtins for OpIAddCarry/OpISubBorrow and improve/fix type inference (#115192)

This PR is to solve several intertwined issues with type inference while
adding support for builtins for OpIAddCarry and OpISubBorrow:
* OpIAddCarry and OpISubBorrow generation in a way of supporting SPIR-V
friendly builtins `__spirv_...` -- introduces a new element to account
for, namely, `ptr sret (%struct) %0` argument that is a place to put a
result of the instruction;
* fix early definition of SPIR-V types during call lowering -- namely,
the goal of the PR is to ensure that correct types are applied to
virtual registers which were used as arguments in call lowering and so
caused early definition of SPIR-V types; reproducers are attached as a
new test cases;
* improve parsing of builtin names (e.g., understand a name of a kind
`"anon<int, int> __spirv_IAddCarry<int, int>(int, int)"` that was
incorrectly parsed as `anon` before the PR);
* improve type inference and fix access to erased from parent after
visit instructions -- before the PR visiting of instructions in
emitintrinsics pass replaced old alloca's, bitcast's, etc. instructions
with a newly generated internal SPIR-V intrinsics and after erasing old
instructions there were still references to them in a postprocessing
working list, while records for newly deduced pointee types were lost;
this PR fixes the issue by adding as consistent wrt. internal data
structures action `SPIRVEmitIntrinsics::replaceAllUsesWith()` that fixes
above mentioned problems;
* LLVM IR add/sub instructions result in logical SPIR-V instructions
when applied to bool type;
* fix validation of pointer types for frexp and lgamma_r,
* fix hardcoded reference to AS0 as a Function storage class in
lib/Target/SPIRV/SPIRVBuiltins.cpp -- now it's
`storageClassToAddressSpace(SPIRV::StorageClass::Function)`,
* re-use the same OpTypeStruct for two identical references to struct's
in arithmetic with overflow instructions.

Author: VyacheslavLevytskyy

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

@@ -190,8 +190,14 @@ std::string lookupBuiltinNameHelper(StringRef DemangledCall) {
   // Check if the extracted name contains type information between angle
   // brackets. If so, the builtin is an instantiated template - needs to have
   // the information after angle brackets and return type removed.
-  if (BuiltinName.find('<') && BuiltinName.back() == '>') {
-    BuiltinName = BuiltinName.substr(0, BuiltinName.find('<'));
+  std::size_t Pos1 = BuiltinName.rfind('<');
+  if (Pos1 != std::string::npos && BuiltinName.back() == '>') {
+    std::size_t Pos2 = BuiltinName.rfind(' ', Pos1);
+    if (Pos2 == std::string::npos)
+      Pos2 = 0;
+    else
+      ++Pos2;
+    BuiltinName = BuiltinName.substr(Pos2, Pos1 - Pos2);
     BuiltinName = BuiltinName.substr(BuiltinName.find_last_of(' ') + 1);
   }
 
@@ -461,9 +467,11 @@ static Register buildBuiltinVariableLoad(
     SPIRVGlobalRegistry *GR, SPIRV::BuiltIn::BuiltIn BuiltinValue, LLT LLType,
     Register Reg = Register(0), bool isConst = true, bool hasLinkageTy = true) {
   Register NewRegister =
-      MIRBuilder.getMRI()->createVirtualRegister(&SPIRV::iIDRegClass);
-  MIRBuilder.getMRI()->setType(NewRegister,
-                               LLT::pointer(0, GR->getPointerSize()));
+      MIRBuilder.getMRI()->createVirtualRegister(&SPIRV::pIDRegClass);
+  MIRBuilder.getMRI()->setType(
+      NewRegister,
+      LLT::pointer(storageClassToAddressSpace(SPIRV::StorageClass::Function),
+                   GR->getPointerSize()));
   SPIRVType *PtrType = GR->getOrCreateSPIRVPointerType(
       VariableType, MIRBuilder, SPIRV::StorageClass::Input);
   GR->assignSPIRVTypeToVReg(PtrType, NewRegister, MIRBuilder.getMF());
@@ -1556,6 +1564,55 @@ static bool generateWaveInst(const SPIRV::IncomingCall *Call,
       /* isConst= */ false, /* hasLinkageTy= */ false);
 }
 
+// We expect a builtin
+//     Name(ptr sret([RetType]) %result, Type %operand1, Type %operand1)
+// where %result is a pointer to where the result of the builtin execution
+// is to be stored, and generate the following instructions:
+//     Res = Opcode RetType Operand1 Operand1
+//     OpStore RetVariable Res
+static bool generateICarryBorrowInst(const SPIRV::IncomingCall *Call,
+                                     MachineIRBuilder &MIRBuilder,
+                                     SPIRVGlobalRegistry *GR) {
+  const SPIRV::DemangledBuiltin *Builtin = Call->Builtin;
+  unsigned Opcode =
+      SPIRV::lookupNativeBuiltin(Builtin->Name, Builtin->Set)->Opcode;
+
+  Register SRetReg = Call->Arguments[0];
+  SPIRVType *PtrRetType = GR->getSPIRVTypeForVReg(SRetReg);
+  SPIRVType *RetType = GR->getPointeeType(PtrRetType);
+  if (!RetType)
+    report_fatal_error("The first parameter must be a pointer");
+  if (RetType->getOpcode() != SPIRV::OpTypeStruct)
+    report_fatal_error("Expected struct type result for the arithmetic with "
+                       "overflow builtins");
+
+  SPIRVType *OpType1 = GR->getSPIRVTypeForVReg(Call->Arguments[1]);
+  SPIRVType *OpType2 = GR->getSPIRVTypeForVReg(Call->Arguments[2]);
+  if (!OpType1 || !OpType2 || OpType1 != OpType2)
+    report_fatal_error("Operands must have the same type");
+  if (OpType1->getOpcode() == SPIRV::OpTypeVector)
+    switch (Opcode) {
+    case SPIRV::OpIAddCarryS:
+      Opcode = SPIRV::OpIAddCarryV;
+      break;
+    case SPIRV::OpISubBorrowS:
+      Opcode = SPIRV::OpISubBorrowV;
+      break;
+    }
+
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  Register ResReg = MRI->createGenericVirtualRegister(LLT::scalar(64));
+  MRI->setRegClass(ResReg, &SPIRV::iIDRegClass);
+  GR->assignSPIRVTypeToVReg(RetType, ResReg, MIRBuilder.getMF());
+  MIRBuilder.buildInstr(Opcode)
+      .addDef(ResReg)
+      .addUse(GR->getSPIRVTypeID(RetType))
+      .addUse(Call->Arguments[1])
+      .addUse(Call->Arguments[2]);
+  MIRBuilder.buildInstr(SPIRV::OpStore).addUse(SRetReg).addUse(ResReg);
+  return true;
+}
+
 static bool generateGetQueryInst(const SPIRV::IncomingCall *Call,
                                  MachineIRBuilder &MIRBuilder,
                                  SPIRVGlobalRegistry *GR) {
@@ -2511,6 +2568,8 @@ std::optional<bool> lowerBuiltin(const StringRef DemangledCall,
     return generateDotOrFMulInst(Call.get(), MIRBuilder, GR);
   case SPIRV::Wave:
     return generateWaveInst(Call.get(), MIRBuilder, GR);
+  case SPIRV::ICarryBorrow:
+    return generateICarryBorrowInst(Call.get(), MIRBuilder, GR);
   case SPIRV::GetQuery:
     return generateGetQueryInst(Call.get(), MIRBuilder, GR);
   case SPIRV::ImageSizeQuery:

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

@@ -63,6 +63,7 @@ def KernelClock : BuiltinGroup;
 def CastToPtr : BuiltinGroup;
 def Construct : BuiltinGroup;
 def CoopMatr : BuiltinGroup;
+def ICarryBorrow : BuiltinGroup;
 
 //===----------------------------------------------------------------------===//
 // Class defining a demangled builtin record. The information in the record
@@ -628,6 +629,10 @@ defm : DemangledNativeBuiltin<"barrier", OpenCL_std, Barrier, 1, 3, OpControlBar
 defm : DemangledNativeBuiltin<"work_group_barrier", OpenCL_std, Barrier, 1, 3, OpControlBarrier>;
 defm : DemangledNativeBuiltin<"__spirv_ControlBarrier", OpenCL_std, Barrier, 3, 3, OpControlBarrier>;
 
+// ICarryBorrow builtin record:
+defm : DemangledNativeBuiltin<"__spirv_IAddCarry", OpenCL_std, ICarryBorrow, 3, 3, OpIAddCarryS>;
+defm : DemangledNativeBuiltin<"__spirv_ISubBorrow", OpenCL_std, ICarryBorrow, 3, 3, OpISubBorrowS>;
+
 // cl_intel_split_work_group_barrier
 defm : DemangledNativeBuiltin<"intel_work_group_barrier_arrive", OpenCL_std, Barrier, 1, 2, OpControlBarrierArriveINTEL>;
 defm : DemangledNativeBuiltin<"__spirv_ControlBarrierArriveINTEL", OpenCL_std, Barrier, 3, 3, OpControlBarrierArriveINTEL>;

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -546,12 +546,36 @@ bool SPIRVCallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
       ArgVRegs.push_back(ArgReg);
       SPIRVType *SpvType = GR->getSPIRVTypeForVReg(ArgReg);
       if (!SpvType) {
-        SpvType = GR->getOrCreateSPIRVType(Arg.Ty, MIRBuilder);
-        GR->assignSPIRVTypeToVReg(SpvType, ArgReg, MF);
+        Type *ArgTy = nullptr;
+        if (auto *PtrArgTy = dyn_cast<PointerType>(Arg.Ty)) {
+          // If Arg.Ty is an untyped pointer (i.e., ptr [addrspace(...)]) and we
+          // don't have access to original value in LLVM IR or info about
+          // deduced pointee type, then we should wait with setting the type for
+          // the virtual register until pre-legalizer step when we access
+          // @llvm.spv.assign.ptr.type.p...(...)'s info.
+          if (Arg.OrigValue)
+            if (Type *ElemTy = GR->findDeducedElementType(Arg.OrigValue))
+              ArgTy =
+                  TypedPointerType::get(ElemTy, PtrArgTy->getAddressSpace());
+        } else {
+          ArgTy = Arg.Ty;
+        }
+        if (ArgTy) {
+          SpvType = GR->getOrCreateSPIRVType(ArgTy, MIRBuilder);
+          GR->assignSPIRVTypeToVReg(SpvType, ArgReg, MF);
+        }
       }
       if (!MRI->getRegClassOrNull(ArgReg)) {
-        MRI->setRegClass(ArgReg, GR->getRegClass(SpvType));
-        MRI->setType(ArgReg, GR->getRegType(SpvType));
+        // Either we have SpvType created, or Arg.Ty is an untyped pointer and
+        // we know its virtual register's class and type even if we don't know
+        // pointee type.
+        MRI->setRegClass(ArgReg, SpvType ? GR->getRegClass(SpvType)
+                                         : &SPIRV::pIDRegClass);
+        MRI->setType(
+            ArgReg,
+            SpvType ? GR->getRegType(SpvType)
+                    : LLT::pointer(cast<PointerType>(Arg.Ty)->getAddressSpace(),
+                                   GR->getPointerSize()));
       }
     }
     auto instructionSet = canUseOpenCL ? SPIRV::InstructionSet::OpenCL_std

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -76,8 +76,8 @@ class SPIRVEmitIntrinsics
   SPIRV::InstructionSet::InstructionSet InstrSet;
 
   // a register of Instructions that don't have a complete type definition
-  SmallPtrSet<Value *, 8> UncompleteTypeInfo;
-  SmallVector<Instruction *> PostprocessWorklist;
+  DenseMap<Value *, unsigned> UncompleteTypeInfo;
+  SmallVector<Value *> PostprocessWorklist;
 
   // well known result types of builtins
   enum WellKnownTypes { Event };
@@ -147,6 +147,7 @@ class SPIRVEmitIntrinsics
                                   std::unordered_set<Function *> &FVisited);
   void replaceWithPtrcasted(Instruction *CI, Type *NewElemTy, Type *KnownElemTy,
                             CallInst *AssignCI);
+  void replaceAllUsesWith(Value *Src, Value *Dest, bool DeleteOld = true);
 
   bool runOnFunction(Function &F);
   bool postprocessTypes();
@@ -272,6 +273,27 @@ static inline void reportFatalOnTokenType(const Instruction *I) {
                        false);
 }
 
+void SPIRVEmitIntrinsics::replaceAllUsesWith(Value *Src, Value *Dest,
+                                             bool DeleteOld) {
+  Src->replaceAllUsesWith(Dest);
+  // Update deduced type records
+  GR->updateIfExistDeducedElementType(Src, Dest, DeleteOld);
+  GR->updateIfExistAssignPtrTypeInstr(Src, Dest, DeleteOld);
+  // Update uncomplete type records if any
+  auto It = UncompleteTypeInfo.find(Src);
+  if (It == UncompleteTypeInfo.end())
+    return;
+  if (DeleteOld) {
+    unsigned Pos = It->second;
+    UncompleteTypeInfo.erase(Src);
+    UncompleteTypeInfo[Dest] = Pos;
+    PostprocessWorklist[Pos] = Dest;
+  } else {
+    UncompleteTypeInfo[Dest] = PostprocessWorklist.size();
+    PostprocessWorklist.push_back(Dest);
+  }
+}
+
 static bool IsKernelArgInt8(Function *F, StoreInst *SI) {
   return SI && F->getCallingConv() == CallingConv::SPIR_KERNEL &&
          isPointerTy(SI->getValueOperand()->getType()) &&
@@ -434,7 +456,7 @@ void SPIRVEmitIntrinsics::maybeAssignPtrType(Type *&Ty, Value *Op, Type *RefTy,
     if (!UnknownElemTypeI8)
       return;
     if (auto *I = dyn_cast<Instruction>(Op)) {
-      UncompleteTypeInfo.insert(I);
+      UncompleteTypeInfo[I] = PostprocessWorklist.size();
       PostprocessWorklist.push_back(I);
     }
   }
@@ -640,7 +662,7 @@ Type *SPIRVEmitIntrinsics::deduceElementType(Value *I, bool UnknownElemTypeI8) {
   if (!UnknownElemTypeI8)
     return nullptr;
   if (auto *Instr = dyn_cast<Instruction>(I)) {
-    UncompleteTypeInfo.insert(Instr);
+    UncompleteTypeInfo[Instr] = PostprocessWorklist.size();
     PostprocessWorklist.push_back(Instr);
   }
   return IntegerType::getInt8Ty(I->getContext());
@@ -1062,7 +1084,7 @@ Instruction *SPIRVEmitIntrinsics::visitSwitchInst(SwitchInst &I) {
                                      {I.getOperand(0)->getType()}, {Args});
   // remove switch to avoid its unneeded and undesirable unwrap into branches
   // and conditions
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   // insert artificial and temporary instruction to preserve valid CFG,
   // it will be removed after IR translation pass
@@ -1084,7 +1106,7 @@ Instruction *SPIRVEmitIntrinsics::visitGetElementPtrInst(GetElementPtrInst &I) {
   for (auto &Op : I.operands())
     Args.push_back(Op);
   auto *NewI = B.CreateIntrinsic(Intrinsic::spv_gep, {Types}, {Args});
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   return NewI;
 }
@@ -1099,7 +1121,7 @@ Instruction *SPIRVEmitIntrinsics::visitBitCastInst(BitCastInst &I) {
   // such bitcasts do not provide sufficient information, should be just skipped
   // here, and handled in insertPtrCastOrAssignTypeInstr.
   if (isPointerTy(I.getType())) {
-    I.replaceAllUsesWith(Source);
+    replaceAllUsesWith(&I, Source);
     I.eraseFromParent();
     return nullptr;
   }
@@ -1108,7 +1130,7 @@ Instruction *SPIRVEmitIntrinsics::visitBitCastInst(BitCastInst &I) {
   SmallVector<Value *> Args(I.op_begin(), I.op_end());
   auto *NewI = B.CreateIntrinsic(Intrinsic::spv_bitcast, {Types}, {Args});
   std::string InstName = I.hasName() ? I.getName().str() : "";
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   NewI->setName(InstName);
   return NewI;
@@ -1219,6 +1241,8 @@ void SPIRVEmitIntrinsics::replacePointerOperandWithPtrCast(
   SmallVector<Value *, 2> Args = {Pointer, VMD, B.getInt32(AddressSpace)};
   auto *PtrCastI = B.CreateIntrinsic(Intrinsic::spv_ptrcast, {Types}, Args);
   I->setOperand(OperandToReplace, PtrCastI);
+  // We need to set up a pointee type for the newly created spv_ptrcast.
+  buildAssignPtr(B, ExpectedElementType, PtrCastI);
 }
 
 void SPIRVEmitIntrinsics::insertPtrCastOrAssignTypeInstr(Instruction *I,
@@ -1331,7 +1355,7 @@ Instruction *SPIRVEmitIntrinsics::visitInsertElementInst(InsertElementInst &I) {
   SmallVector<Value *> Args(I.op_begin(), I.op_end());
   auto *NewI = B.CreateIntrinsic(Intrinsic::spv_insertelt, {Types}, {Args});
   std::string InstName = I.hasName() ? I.getName().str() : "";
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   NewI->setName(InstName);
   return NewI;
@@ -1346,7 +1370,7 @@ SPIRVEmitIntrinsics::visitExtractElementInst(ExtractElementInst &I) {
   SmallVector<Value *, 2> Args = {I.getVectorOperand(), I.getIndexOperand()};
   auto *NewI = B.CreateIntrinsic(Intrinsic::spv_extractelt, {Types}, {Args});
   std::string InstName = I.hasName() ? I.getName().str() : "";
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   NewI->setName(InstName);
   return NewI;
@@ -1382,7 +1406,7 @@ Instruction *SPIRVEmitIntrinsics::visitExtractValueInst(ExtractValueInst &I) {
     Args.push_back(B.getInt32(Op));
   auto *NewI =
       B.CreateIntrinsic(Intrinsic::spv_extractv, {I.getType()}, {Args});
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   return NewI;
 }
@@ -1443,7 +1467,7 @@ Instruction *SPIRVEmitIntrinsics::visitAllocaInst(AllocaInst &I) {
                                     {PtrTy, ArraySize->getType()}, {ArraySize})
                 : B.CreateIntrinsic(Intrinsic::spv_alloca, {PtrTy}, {});
   std::string InstName = I.hasName() ? I.getName().str() : "";
-  I.replaceAllUsesWith(NewI);
+  replaceAllUsesWith(&I, NewI);
   I.eraseFromParent();
   NewI->setName(InstName);
   return NewI;
@@ -1613,7 +1637,7 @@ void SPIRVEmitIntrinsics::processInstrAfterVisit(Instruction *I,
     auto *NewOp =
         buildIntrWithMD(Intrinsic::spv_track_constant,
                         {II->getType(), II->getType()}, t->second, I, {}, B);
-    I->replaceAllUsesWith(NewOp);
+    replaceAllUsesWith(I, NewOp, false);
     NewOp->setArgOperand(0, I);
   }
   bool IsPhi = isa<PHINode>(I), BPrepared = false;

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -180,6 +180,15 @@ class SPIRVGlobalRegistry {
     auto It = AssignPtrTypeInstr.find(Val);
     return It == AssignPtrTypeInstr.end() ? nullptr : It->second;
   }
+  // - Find a record and update its key or add a new record, if found.
+  void updateIfExistAssignPtrTypeInstr(Value *OldVal, Value *NewVal,
+                                       bool DeleteOld) {
+    if (CallInst *CI = findAssignPtrTypeInstr(OldVal)) {
+      if (DeleteOld)
+        AssignPtrTypeInstr.erase(OldVal);
+      AssignPtrTypeInstr[NewVal] = CI;
+    }
+  }
 
   // A registry of mutated values
   // (see `SPIRVPrepareFunctions::removeAggregateTypesFromSignature()`):
@@ -214,6 +223,15 @@ class SPIRVGlobalRegistry {
     auto It = DeducedElTys.find(Val);
     return It == DeducedElTys.end() ? nullptr : It->second;
   }
+  // - Find a record and update its key or add a new record, if found.
+  void updateIfExistDeducedElementType(Value *OldVal, Value *NewVal,
+                                       bool DeleteOld) {
+    if (Type *Ty = findDeducedElementType(OldVal)) {
+      if (DeleteOld)
+        DeducedElTys.erase(OldVal);
+      DeducedElTys[NewVal] = Ty;
+    }
+  }
   // - Add a record to the map of deduced composite types.
   void addDeducedCompositeType(Value *Val, Type *Ty) {
     DeducedNestedTys[Val] = Ty;

llvm/lib/Target/SPIRV/SPIRVISelLowering.cpp

@@ -414,6 +414,17 @@ void SPIRVTargetLowering::finalizeLowering(MachineFunction &MF) const {
         validateForwardCalls(STI, MRI, GR, MI);
         break;
 
+      // ensure that LLVM IR add/sub instructions result in logical SPIR-V
+      // instructions when applied to bool type
+      case SPIRV::OpIAddS:
+      case SPIRV::OpIAddV:
+      case SPIRV::OpISubS:
+      case SPIRV::OpISubV:
+        if (GR.isScalarOrVectorOfType(MI.getOperand(1).getReg(),
+                                      SPIRV::OpTypeBool))
+          MI.setDesc(STI.getInstrInfo()->get(SPIRV::OpLogicalNotEqual));
+        break;
+
       // ensure that LLVM IR bitwise instructions result in logical SPIR-V
       // instructions when applied to bool type
       case SPIRV::OpBitwiseOrS:
@@ -473,8 +484,11 @@ void SPIRVTargetLowering::finalizeLowering(MachineFunction &MF) const {
             MI.getOperand(2).getImm() != SPIRV::InstructionSet::OpenCL_std)
           continue;
         switch (MI.getOperand(3).getImm()) {
+        case SPIRV::OpenCLExtInst::frexp:
+        case SPIRV::OpenCLExtInst::lgamma_r:
         case SPIRV::OpenCLExtInst::remquo: {
-          // The last operand must be of a pointer to the return type.
+          // The last operand must be of a pointer to i32 or vector of i32
+          // values.
           MachineIRBuilder MIB(MI);
           SPIRVType *Int32Type = GR.getOrCreateSPIRVIntegerType(32, MIB);
           SPIRVType *RetType = MRI->getVRegDef(MI.getOperand(1).getReg());
@@ -487,8 +501,6 @@ void SPIRVTargetLowering::finalizeLowering(MachineFunction &MF) const {
                         Int32Type, RetType->getOperand(2).getImm(), MIB));
         } break;
         case SPIRV::OpenCLExtInst::fract:
-        case SPIRV::OpenCLExtInst::frexp:
-        case SPIRV::OpenCLExtInst::lgamma_r:
         case SPIRV::OpenCLExtInst::modf:
         case SPIRV::OpenCLExtInst::sincos:
           // The last operand must be of a pointer to the base type represented