[SPIR-V] Fix emission of debug and annotation instructions and add SPV_EXT_optnone SPIR-V extension (llvm#118402)

This PR fixes:
* emission of OpNames (added newly inserted internal intrinsics and
basic blocks)
* emission of function attributes (SRet is added)
* implementation of SPV_INTEL_optnone so that it emits OptNoneINTEL
Function Control flag, and add implementation of the SPV_EXT_optnone
SPIR-V extension.

Author: VyacheslavLevytskyy

llvm/docs/SPIRVUsage.rst

@@ -141,16 +141,18 @@ list of supported SPIR-V extensions, sorted alphabetically by their extension na
 
    * - Extension Name
      - Description
+   * - ``SPV_EXT_arithmetic_fence``
+     - Adds an instruction that prevents fast-math optimizations between its argument and the expression that contains it.
+   * - ``SPV_EXT_demote_to_helper_invocation``
+     - Adds an instruction that demotes a fragment shader invocation to a helper invocation.
+   * - ``SPV_EXT_optnone``
+     - Adds OptNoneEXT value for Function Control mask that indicates a request to not optimize the function.
    * - ``SPV_EXT_shader_atomic_float16_add``
      - Extends the SPV_EXT_shader_atomic_float_add extension to support atomically adding to 16-bit floating-point numbers in memory.
    * - ``SPV_EXT_shader_atomic_float_add``
      - Adds atomic add instruction on floating-point numbers.
    * - ``SPV_EXT_shader_atomic_float_min_max``
      - Adds atomic min and max instruction on floating-point numbers.
-   * - ``SPV_EXT_arithmetic_fence``
-     - Adds an instruction that prevents fast-math optimizations between its argument and the expression that contains it.
-   * - ``SPV_EXT_demote_to_helper_invocation``
-     - Adds an instruction that demotes a fragment shader invocation to a helper invocation.
    * - ``SPV_INTEL_arbitrary_precision_integers``
      - Allows generating arbitrary width integer types.
    * - ``SPV_INTEL_bfloat16_conversion``

llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp

@@ -45,6 +45,7 @@ using namespace llvm;
 namespace {
 class SPIRVAsmPrinter : public AsmPrinter {
   unsigned NLabels = 0;
+  SmallPtrSet<const MachineBasicBlock *, 8> LabeledMBB;
 
 public:
   explicit SPIRVAsmPrinter(TargetMachine &TM,
@@ -152,13 +153,9 @@ void SPIRVAsmPrinter::outputOpFunctionEnd() {
   outputMCInst(FunctionEndInst);
 }
 
-// Emit OpFunctionEnd at the end of MF and clear BBNumToRegMap.
 void SPIRVAsmPrinter::emitFunctionBodyEnd() {
-  // Do not emit anything if it's an internal service function.
-  if (isHidden())
-    return;
-  outputOpFunctionEnd();
-  MAI->BBNumToRegMap.clear();
+  if (!isHidden())
+    outputOpFunctionEnd();
 }
 
 void SPIRVAsmPrinter::emitOpLabel(const MachineBasicBlock &MBB) {
@@ -171,6 +168,7 @@ void SPIRVAsmPrinter::emitOpLabel(const MachineBasicBlock &MBB) {
   LabelInst.addOperand(MCOperand::createReg(MAI->getOrCreateMBBRegister(MBB)));
   outputMCInst(LabelInst);
   ++NLabels;
+  LabeledMBB.insert(&MBB);
 }
 
 void SPIRVAsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
@@ -267,7 +265,7 @@ void SPIRVAsmPrinter::emitInstruction(const MachineInstr *MI) {
 
   // Output OpLabel after OpFunction and OpFunctionParameter in the first MBB.
   const MachineInstr *NextMI = MI->getNextNode();
-  if (!MAI->hasMBBRegister(*MI->getParent()) && isFuncOrHeaderInstr(MI, TII) &&
+  if (!LabeledMBB.contains(MI->getParent()) && isFuncOrHeaderInstr(MI, TII) &&
       (!NextMI || !isFuncOrHeaderInstr(NextMI, TII))) {
     assert(MI->getParent()->getNumber() == MF->front().getNumber() &&
            "OpFunction is not in the front MBB of MF");

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -65,7 +65,8 @@ bool SPIRVCallLowering::lowerReturn(MachineIRBuilder &MIRBuilder,
 }
 
 // Based on the LLVM function attributes, get a SPIR-V FunctionControl.
-static uint32_t getFunctionControl(const Function &F) {
+static uint32_t getFunctionControl(const Function &F,
+                                   const SPIRVSubtarget *ST) {
   MemoryEffects MemEffects = F.getMemoryEffects();
 
   uint32_t FuncControl = static_cast<uint32_t>(SPIRV::FunctionControl::None);
@@ -80,6 +81,11 @@ static uint32_t getFunctionControl(const Function &F) {
   else if (MemEffects.onlyReadsMemory())
     FuncControl |= static_cast<uint32_t>(SPIRV::FunctionControl::Const);
 
+  if (ST->canUseExtension(SPIRV::Extension::SPV_INTEL_optnone) ||
+      ST->canUseExtension(SPIRV::Extension::SPV_EXT_optnone))
+    if (F.hasFnAttribute(Attribute::OptimizeNone))
+      FuncControl |= static_cast<uint32_t>(SPIRV::FunctionControl::OptNoneEXT);
+
   return FuncControl;
 }
 
@@ -346,6 +352,12 @@ bool SPIRVCallLowering::lowerFormalArguments(MachineIRBuilder &MIRBuilder,
         buildOpDecorate(VRegs[i][0], MIRBuilder,
                         SPIRV::Decoration::FuncParamAttr, {Attr});
       }
+      if (Arg.hasAttribute(Attribute::StructRet)) {
+        auto Attr =
+            static_cast<unsigned>(SPIRV::FunctionParameterAttribute::Sret);
+        buildOpDecorate(VRegs[i][0], MIRBuilder,
+                        SPIRV::Decoration::FuncParamAttr, {Attr});
+      }
 
       if (F.getCallingConv() == CallingConv::SPIR_KERNEL) {
         std::vector<SPIRV::Decoration::Decoration> ArgTypeQualDecs =
@@ -397,7 +409,7 @@ bool SPIRVCallLowering::lowerFormalArguments(MachineIRBuilder &MIRBuilder,
   FTy = fixFunctionTypeIfPtrArgs(GR, F, FTy, RetTy, ArgTypeVRegs);
   SPIRVType *FuncTy = GR->getOrCreateOpTypeFunctionWithArgs(
       FTy, RetTy, ArgTypeVRegs, MIRBuilder);
-  uint32_t FuncControl = getFunctionControl(F);
+  uint32_t FuncControl = getFunctionControl(F, ST);
 
   // Add OpFunction instruction
   MachineInstrBuilder MB = MIRBuilder.buildInstr(SPIRV::OpFunction)
@@ -427,10 +439,8 @@ bool SPIRVCallLowering::lowerFormalArguments(MachineIRBuilder &MIRBuilder,
 
   // Handle entry points and function linkage.
   if (isEntryPoint(F)) {
-    const auto &STI = MIRBuilder.getMF().getSubtarget<SPIRVSubtarget>();
-    auto executionModel = getExecutionModel(STI, F);
     auto MIB = MIRBuilder.buildInstr(SPIRV::OpEntryPoint)
-                   .addImm(static_cast<uint32_t>(executionModel))
+                   .addImm(static_cast<uint32_t>(getExecutionModel(*ST, F)))
                    .addUse(FuncVReg);
     addStringImm(F.getName(), MIB);
   } else if (F.getLinkage() != GlobalValue::InternalLinkage &&

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -42,6 +42,7 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
         {"SPV_INTEL_global_variable_host_access",
          SPIRV::Extension::Extension::SPV_INTEL_global_variable_host_access},
         {"SPV_INTEL_optnone", SPIRV::Extension::Extension::SPV_INTEL_optnone},
+        {"SPV_EXT_optnone", SPIRV::Extension::Extension::SPV_EXT_optnone},
         {"SPV_INTEL_usm_storage_classes",
          SPIRV::Extension::Extension::SPV_INTEL_usm_storage_classes},
         {"SPV_INTEL_split_barrier",

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -199,6 +199,8 @@ class SPIRVEmitIntrinsics
                             DenseMap<Function *, CallInst *> Ptrcasts);
 
   void replaceAllUsesWith(Value *Src, Value *Dest, bool DeleteOld = true);
+  void replaceAllUsesWithAndErase(IRBuilder<> &B, Instruction *Src,
+                                  Instruction *Dest, bool DeleteOld = true);
 
   bool runOnFunction(Function &F);
   bool postprocessTypes(Module &M);
@@ -322,6 +324,17 @@ static inline void reportFatalOnTokenType(const Instruction *I) {
                        false);
 }
 
+static void emitAssignName(Instruction *I, IRBuilder<> &B) {
+  if (!I->hasName() || I->getType()->isAggregateType() ||
+      expectIgnoredInIRTranslation(I))
+    return;
+  reportFatalOnTokenType(I);
+  setInsertPointAfterDef(B, I);
+  std::vector<Value *> Args = {I};
+  addStringImm(I->getName(), B, Args);
+  B.CreateIntrinsic(Intrinsic::spv_assign_name, {I->getType()}, Args);
+}
+
 void SPIRVEmitIntrinsics::replaceAllUsesWith(Value *Src, Value *Dest,
                                              bool DeleteOld) {
   Src->replaceAllUsesWith(Dest);
@@ -336,6 +349,19 @@ void SPIRVEmitIntrinsics::replaceAllUsesWith(Value *Src, Value *Dest,
   }
 }
 
+void SPIRVEmitIntrinsics::replaceAllUsesWithAndErase(IRBuilder<> &B,
+                                                     Instruction *Src,
+                                                     Instruction *Dest,
+                                                     bool DeleteOld) {
+  replaceAllUsesWith(Src, Dest, DeleteOld);
+  std::string Name = Src->hasName() ? Src->getName().str() : "";
+  Src->eraseFromParent();
+  if (!Name.empty()) {
+    Dest->setName(Name);
+    emitAssignName(Dest, B);
+  }
+}
+
 static bool IsKernelArgInt8(Function *F, StoreInst *SI) {
   return SI && F->getCallingConv() == CallingConv::SPIR_KERNEL &&
          isPointerTy(SI->getValueOperand()->getType()) &&
@@ -1308,8 +1334,7 @@ Instruction *SPIRVEmitIntrinsics::visitGetElementPtrInst(GetElementPtrInst &I) {
   for (auto &Op : I.operands())
     Args.push_back(Op);
   auto *NewI = B.CreateIntrinsic(Intrinsic::spv_gep, {Types}, {Args});
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1331,10 +1356,7 @@ Instruction *SPIRVEmitIntrinsics::visitBitCastInst(BitCastInst &I) {
   SmallVector<Type *, 2> Types = {I.getType(), Source->getType()};
   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() : "";
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
-  NewI->setName(InstName);
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1589,10 +1611,7 @@ Instruction *SPIRVEmitIntrinsics::visitInsertElementInst(InsertElementInst &I) {
   B.SetInsertPoint(&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() : "";
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
-  NewI->setName(InstName);
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1604,10 +1623,7 @@ SPIRVEmitIntrinsics::visitExtractElementInst(ExtractElementInst &I) {
                                   I.getIndexOperand()->getType()};
   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() : "";
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
-  NewI->setName(InstName);
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1641,8 +1657,7 @@ Instruction *SPIRVEmitIntrinsics::visitExtractValueInst(ExtractValueInst &I) {
     Args.push_back(B.getInt32(Op));
   auto *NewI =
       B.CreateIntrinsic(Intrinsic::spv_extractv, {I.getType()}, {Args});
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1701,10 +1716,7 @@ Instruction *SPIRVEmitIntrinsics::visitAllocaInst(AllocaInst &I) {
       ArraySize ? B.CreateIntrinsic(Intrinsic::spv_alloca_array,
                                     {PtrTy, ArraySize->getType()}, {ArraySize})
                 : B.CreateIntrinsic(Intrinsic::spv_alloca, {PtrTy}, {});
-  std::string InstName = I.hasName() ? I.getName().str() : "";
-  replaceAllUsesWith(&I, NewI);
-  I.eraseFromParent();
-  NewI->setName(InstName);
+  replaceAllUsesWithAndErase(B, &I, NewI);
   return NewI;
 }
 
@@ -1914,14 +1926,7 @@ void SPIRVEmitIntrinsics::processInstrAfterVisit(Instruction *I,
       I->setOperand(OpNo, NewOp);
     }
   }
-  if (I->hasName() && !I->getType()->isAggregateType() &&
-      !expectIgnoredInIRTranslation(I)) {
-    reportFatalOnTokenType(I);
-    setInsertPointAfterDef(B, I);
-    std::vector<Value *> Args = {I};
-    addStringImm(I->getName(), B, Args);
-    B.CreateIntrinsic(Intrinsic::spv_assign_name, {I->getType()}, Args);
-  }
+  emitAssignName(I, B);
 }
 
 Type *SPIRVEmitIntrinsics::deduceFunParamElementType(Function *F,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -421,6 +421,7 @@ void SPIRVModuleAnalysis::processOtherInstrs(const Module &M) {
       continue;
     MachineFunction *MF = MMI->getMachineFunction(*F);
     assert(MF);
+
     for (MachineBasicBlock &MBB : *MF)
       for (MachineInstr &MI : MBB) {
         if (MAI.getSkipEmission(&MI))
@@ -1548,11 +1549,14 @@ static void collectReqs(const Module &M, SPIRV::ModuleAnalysisInfo &MAI,
           SPIRV::OperandCategory::ExecutionModeOperand,
           SPIRV::ExecutionMode::VecTypeHint, ST);
 
-    if (F.hasOptNone() &&
-        ST.canUseExtension(SPIRV::Extension::SPV_INTEL_optnone)) {
-      // Output OpCapability OptNoneINTEL.
-      MAI.Reqs.addExtension(SPIRV::Extension::SPV_INTEL_optnone);
-      MAI.Reqs.addCapability(SPIRV::Capability::OptNoneINTEL);
+    if (F.hasOptNone()) {
+      if (ST.canUseExtension(SPIRV::Extension::SPV_EXT_optnone)) {
+        MAI.Reqs.addExtension(SPIRV::Extension::SPV_EXT_optnone);
+        MAI.Reqs.addCapability(SPIRV::Capability::OptNoneEXT);
+      } else if (ST.canUseExtension(SPIRV::Extension::SPV_INTEL_optnone)) {
+        MAI.Reqs.addExtension(SPIRV::Extension::SPV_INTEL_optnone);
+        MAI.Reqs.addCapability(SPIRV::Capability::OptNoneINTEL);
+      }
     }
   }
 }
@@ -1613,6 +1617,27 @@ static void addDecorations(const Module &M, const SPIRVInstrInfo &TII,
   }
 }
 
+static void addMBBNames(const Module &M, const SPIRVInstrInfo &TII,
+                        MachineModuleInfo *MMI, const SPIRVSubtarget &ST,
+                        SPIRV::ModuleAnalysisInfo &MAI) {
+  for (auto F = M.begin(), E = M.end(); F != E; ++F) {
+    MachineFunction *MF = MMI->getMachineFunction(*F);
+    if (!MF)
+      continue;
+    MachineRegisterInfo &MRI = MF->getRegInfo();
+    for (auto &MBB : *MF) {
+      if (!MBB.hasName() || MBB.empty())
+        continue;
+      // Emit basic block names.
+      Register Reg = MRI.createGenericVirtualRegister(LLT::scalar(64));
+      MRI.setRegClass(Reg, &SPIRV::IDRegClass);
+      buildOpName(Reg, MBB.getName(), *std::prev(MBB.end()), TII);
+      Register GlobalReg = MAI.getOrCreateMBBRegister(MBB);
+      MAI.setRegisterAlias(MF, Reg, GlobalReg);
+    }
+  }
+}
+
 struct SPIRV::ModuleAnalysisInfo SPIRVModuleAnalysis::MAI;
 
 void SPIRVModuleAnalysis::getAnalysisUsage(AnalysisUsage &AU) const {
@@ -1631,6 +1656,7 @@ bool SPIRVModuleAnalysis::runOnModule(Module &M) {
 
   setBaseInfo(M);
 
+  addMBBNames(M, *TII, MMI, *ST, MAI);
   addDecorations(M, *TII, MMI, *ST, MAI);
 
   collectReqs(M, MAI, MMI, *ST);

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h

@@ -157,7 +157,7 @@ struct ModuleAnalysisInfo {
   // The array contains lists of MIs for each module section.
   InstrList MS[NUM_MODULE_SECTIONS];
   // The table maps MBB number to SPIR-V unique ID register.
-  DenseMap<int, Register> BBNumToRegMap;
+  DenseMap<std::pair<const MachineFunction *, int>, Register> BBNumToRegMap;
 
   Register getFuncReg(const Function *F) {
     assert(F && "Function is null");
@@ -188,15 +188,17 @@ struct ModuleAnalysisInfo {
   }
   unsigned getNextID() { return MaxID++; }
   bool hasMBBRegister(const MachineBasicBlock &MBB) {
-    return BBNumToRegMap.contains(MBB.getNumber());
+    auto Key = std::make_pair(MBB.getParent(), MBB.getNumber());
+    return BBNumToRegMap.contains(Key);
   }
   // Convert MBB's number to corresponding ID register.
   Register getOrCreateMBBRegister(const MachineBasicBlock &MBB) {
-    auto f = BBNumToRegMap.find(MBB.getNumber());
-    if (f != BBNumToRegMap.end())
-      return f->second;
+    auto Key = std::make_pair(MBB.getParent(), MBB.getNumber());
+    auto It = BBNumToRegMap.find(Key);
+    if (It != BBNumToRegMap.end())
+      return It->second;
     Register NewReg = Register::index2VirtReg(getNextID());
-    BBNumToRegMap[MBB.getNumber()] = NewReg;
+    BBNumToRegMap[Key] = NewReg;
     return NewReg;
   }
 };

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -304,6 +304,7 @@ defm SPV_INTEL_global_variable_host_access : ExtensionOperand<109>;
 defm SPV_INTEL_global_variable_fpga_decorations : ExtensionOperand<110>;
 defm SPV_KHR_cooperative_matrix : ExtensionOperand<111>;
 defm SPV_EXT_arithmetic_fence : ExtensionOperand<112>;
+defm SPV_EXT_optnone : ExtensionOperand<113>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
@@ -463,6 +464,7 @@ defm PhysicalStorageBufferAddressesEXT : CapabilityOperand<5347, 0, 0, [], [Shad
 defm CooperativeMatrixNV : CapabilityOperand<5357, 0, 0, [], [Shader]>;
 defm ArbitraryPrecisionIntegersINTEL : CapabilityOperand<5844, 0, 0, [SPV_INTEL_arbitrary_precision_integers], [Int8, Int16]>;
 defm OptNoneINTEL : CapabilityOperand<6094, 0, 0, [SPV_INTEL_optnone], []>;
+defm OptNoneEXT : CapabilityOperand<6094, 0, 0, [SPV_EXT_optnone], []>;
 defm BitInstructions : CapabilityOperand<6025, 0, 0, [SPV_KHR_bit_instructions], []>;
 defm ExpectAssumeKHR : CapabilityOperand<5629, 0, 0, [SPV_KHR_expect_assume], []>;
 defm FunctionPointersINTEL : CapabilityOperand<5603, 0, 0, [SPV_INTEL_function_pointers], []>;
@@ -1433,6 +1435,7 @@ defm Inline : FunctionControlOperand<0x1>;
 defm DontInline : FunctionControlOperand<0x2>;
 defm Pure : FunctionControlOperand<0x4>;
 defm Const : FunctionControlOperand<0x8>;
+defm OptNoneEXT : FunctionControlOperand<0x10000>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define MemorySemantics enum values and at the same time

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -108,6 +108,16 @@ void buildOpName(Register Target, const StringRef &Name,
   }
 }
 
+void buildOpName(Register Target, const StringRef &Name, MachineInstr &I,
+                 const SPIRVInstrInfo &TII) {
+  if (!Name.empty()) {
+    auto MIB =
+        BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(SPIRV::OpName))
+            .addUse(Target);
+    addStringImm(Name, MIB);
+  }
+}
+
 static void finishBuildOpDecorate(MachineInstrBuilder &MIB,
                                   const std::vector<uint32_t> &DecArgs,
                                   StringRef StrImm) {