[SPIR-V] Improve type inference: fix types of return values in call lowering

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

@@ -469,12 +469,8 @@ static Register buildLoadInst(SPIRVType *BaseType, Register PtrRegister,
                               MachineIRBuilder &MIRBuilder,
                               SPIRVGlobalRegistry *GR, LLT LowLevelType,
                               Register DestinationReg = Register(0)) {
-  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
-  if (!DestinationReg.isValid()) {
-    DestinationReg = MRI->createVirtualRegister(&SPIRV::iIDRegClass);
-    MRI->setType(DestinationReg, LLT::scalar(64));
-    GR->assignSPIRVTypeToVReg(BaseType, DestinationReg, MIRBuilder.getMF());
-  }
+  if (!DestinationReg.isValid())
+    DestinationReg = createVirtualRegister(BaseType, GR, MIRBuilder);
   // TODO: consider using correct address space and alignment (p0 is canonical
   // type for selection though).
   MachinePointerInfo PtrInfo = MachinePointerInfo();
@@ -2151,7 +2147,7 @@ static bool buildEnqueueKernel(const SPIRV::IncomingCall *Call,
     const SPIRVType *PointerSizeTy = GR->getOrCreateSPIRVPointerType(
         Int32Ty, MIRBuilder, SPIRV::StorageClass::Function);
     for (unsigned I = 0; I < LocalSizeNum; ++I) {
-      Register Reg = MRI->createVirtualRegister(&SPIRV::iIDRegClass);
+      Register Reg = MRI->createVirtualRegister(&SPIRV::pIDRegClass);
       MRI->setType(Reg, LLType);
       GR->assignSPIRVTypeToVReg(PointerSizeTy, Reg, MIRBuilder.getMF());
       auto GEPInst = MIRBuilder.buildIntrinsic(
@@ -2539,23 +2535,11 @@ std::optional<bool> lowerBuiltin(const StringRef DemangledCall,
                                  SPIRVGlobalRegistry *GR) {
   LLVM_DEBUG(dbgs() << "Lowering builtin call: " << DemangledCall << "\n");
 
-  // SPIR-V type and return register.
-  Register ReturnRegister = OrigRet;
-  SPIRVType *ReturnType = nullptr;
-  if (OrigRetTy && !OrigRetTy->isVoidTy()) {
-    ReturnType = GR->assignTypeToVReg(OrigRetTy, ReturnRegister, MIRBuilder);
-    if (!MIRBuilder.getMRI()->getRegClassOrNull(ReturnRegister))
-      MIRBuilder.getMRI()->setRegClass(ReturnRegister,
-                                       GR->getRegClass(ReturnType));
-  } else if (OrigRetTy && OrigRetTy->isVoidTy()) {
-    ReturnRegister = MIRBuilder.getMRI()->createVirtualRegister(&IDRegClass);
-    MIRBuilder.getMRI()->setType(ReturnRegister, LLT::scalar(64));
-    ReturnType = GR->assignTypeToVReg(OrigRetTy, ReturnRegister, MIRBuilder);
-  }
-
   // Lookup the builtin in the TableGen records.
+  SPIRVType *SpvType = GR->getSPIRVTypeForVReg(OrigRet);
+  assert(SpvType && "Inconsistent return register: expected valid type info");
   std::unique_ptr<const IncomingCall> Call =
-      lookupBuiltin(DemangledCall, Set, ReturnRegister, ReturnType, Args);
+      lookupBuiltin(DemangledCall, Set, OrigRet, SpvType, Args);
 
   if (!Call) {
     LLVM_DEBUG(dbgs() << "Builtin record was not found!\n");

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -539,6 +539,23 @@ bool SPIRVCallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
 
   if (isFunctionDecl && !DemangledName.empty() &&
       (canUseGLSL || canUseOpenCL)) {
+    if (ResVReg.isValid()) {
+      if (!GR->getSPIRVTypeForVReg(ResVReg)) {
+        const Type *RetTy = OrigRetTy;
+        if (auto *PtrRetTy = dyn_cast<PointerType>(OrigRetTy)) {
+          const Value *OrigValue = Info.OrigRet.OrigValue;
+          if (!OrigValue)
+            OrigValue = Info.CB;
+          if (OrigValue)
+            if (Type *ElemTy = GR->findDeducedElementType(OrigValue))
+              RetTy =
+                  TypedPointerType::get(ElemTy, PtrRetTy->getAddressSpace());
+        }
+        setRegClassType(ResVReg, RetTy, GR, MIRBuilder);
+      }
+    } else {
+      ResVReg = createVirtualRegister(OrigRetTy, GR, MIRBuilder);
+    }
     SmallVector<Register, 8> ArgVRegs;
     for (auto Arg : Info.OrigArgs) {
       assert(Arg.Regs.size() == 1 && "Call arg has multiple VRegs");

llvm/lib/Target/SPIRV/SPIRVDuplicatesTracker.cpp

@@ -69,16 +69,31 @@ void SPIRVGeneralDuplicatesTracker::buildDepsGraph(
         MachineOperand *RegOp = &VRegDef->getOperand(0);
         if (Reg2Entry.count(RegOp) == 0 &&
             (MI->getOpcode() != SPIRV::OpVariable || i != 3)) {
-          std::string DiagMsg;
-          raw_string_ostream OS(DiagMsg);
-          OS << "Unexpected pattern while building a dependency "
-                "graph.\nInstruction: ";
-          MI->print(OS);
-          OS << "Operand: ";
-          Op.print(OS);
-          OS << "\nOperand definition: ";
-          VRegDef->print(OS);
-          report_fatal_error(DiagMsg.c_str());
+          // try to repair the unexpected code pattern
+          bool IsFixed = false;
+          if (VRegDef->getOpcode() == TargetOpcode::G_CONSTANT &&
+              RegOp->isReg() && MRI.getType(RegOp->getReg()).isScalar()) {
+            const Constant *C = VRegDef->getOperand(1).getCImm();
+            add(C, MI->getParent()->getParent(), RegOp->getReg());
+            auto Iter = CT.Storage.find(C);
+            if (Iter != CT.Storage.end()) {
+              SPIRV::DTSortableEntry &MissedEntry = Iter->second;
+              Reg2Entry[RegOp] = &MissedEntry;
+              IsFixed = true;
+            }
+          }
+          if (!IsFixed) {
+            std::string DiagMsg;
+            raw_string_ostream OS(DiagMsg);
+            OS << "Unexpected pattern while building a dependency "
+                  "graph.\nInstruction: ";
+            MI->print(OS);
+            OS << "Operand: ";
+            Op.print(OS);
+            OS << "\nOperand definition: ";
+            VRegDef->print(OS);
+            report_fatal_error(DiagMsg.c_str());
+          }
         }
         if (Reg2Entry.count(RegOp))
           E->addDep(Reg2Entry[RegOp]);