[SPIR-V] Add SPV_INTEL_memory_access_aliasing extension (#129800)

Spec can be found here intel/llvm#15225
 TODO for future patches:
- During spec review need to decide whether only FunctionCall or Atomic
instructions can be decorated and if not - move the code around adding
handling for other instructions;
- Handle optional string metadata;
- Handle LLVM atomic instructions;
- Handle SPIR-V friendly atomic calls returning via sret argument.

Signed-off-by: Sidorov, Dmitry <[email protected]>

Author: MrSidims

llvm/docs/SPIRVUsage.rst

@@ -173,6 +173,8 @@ list of supported SPIR-V extensions, sorted alphabetically by their extension na
      - Adds decorations that can be applied to global (module scope) variables to help code generation for FPGA devices.
    * - ``SPV_INTEL_media_block_io``
      - Adds additional subgroup block read and write functionality that allow applications to flexibly specify the width and height of the block to read from or write to a 2D image.
+   * - ``SPV_INTEL_memory_access_aliasing``
+     - Adds instructions and decorations to specify memory access aliasing, similar to alias.scope and noalias LLVM metadata.
    * - ``SPV_INTEL_optnone``
      - Adds OptNoneINTEL value for Function Control mask that indicates a request to not optimize the function.
    * - ``SPV_INTEL_split_barrier``
@@ -301,6 +303,10 @@ SPIR-V backend, along with their descriptions and argument details.
      - None
      - `[Type, Metadata]`
      - Assigns decoration to values by associating them with metadatas. Not emitted directly but used to support SPIR-V representation in LLVM IR.
+   * - `int_spv_assign_aliasing_decoration`
+     - None
+     - `[Type, 32-bit Integer, Metadata]`
+     - Assigns one of two memory aliasing decorations (specified by the second argument) to instructions using original aliasing metadata node. Not emitted directly but used to support SPIR-V representation in LLVM IR.
    * - `int_spv_track_constant`
      - Type
      - `[Type, Metadata]`

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -138,4 +138,7 @@ let TargetPrefix = "spv" in {
   def int_spv_resource_store_typedbuffer
     : DefaultAttrsIntrinsic<[], [llvm_any_ty, llvm_i32_ty, llvm_anyvector_ty]>;
 
+  // Memory aliasing intrinsics
+  def int_spv_assign_aliasing_decoration : Intrinsic<[], [llvm_any_ty, llvm_i32_ty, llvm_metadata_ty], [ImmArg<ArgIndex<1>>]>;
+
 }

llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp

@@ -611,6 +611,9 @@ void SPIRVAsmPrinter::outputModuleSections() {
   outputModuleSection(SPIRV::MB_DebugNames);
   // 7c. Debug: all OpModuleProcessed instructions.
   outputModuleSection(SPIRV::MB_DebugModuleProcessed);
+  // xxx. SPV_INTEL_memory_access_aliasing instructions go before 8.
+  // "All annotation instructions"
+  outputModuleSection(SPIRV::MB_AliasingInsts);
   // 8. All annotation instructions (all decorations).
   outputAnnotations(*M);
   // 9. All type declarations (OpTypeXXX instructions), all constant

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -695,6 +695,20 @@ bool SPIRVCallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
       return false;
     MIB.addUse(Arg.Regs[0]);
   }
+
+  if (ST->canUseExtension(SPIRV::Extension::SPV_INTEL_memory_access_aliasing)) {
+    // Process aliasing metadata.
+    const CallBase *CI = Info.CB;
+    if (CI && CI->hasMetadata()) {
+      if (MDNode *MD = CI->getMetadata(LLVMContext::MD_alias_scope))
+        GR->buildMemAliasingOpDecorate(ResVReg, MIRBuilder,
+                                       SPIRV::Decoration::AliasScopeINTEL, MD);
+      if (MDNode *MD = CI->getMetadata(LLVMContext::MD_noalias))
+        GR->buildMemAliasingOpDecorate(ResVReg, MIRBuilder,
+                                       SPIRV::Decoration::NoAliasINTEL, MD);
+    }
+  }
+
   return MIB.constrainAllUses(MIRBuilder.getTII(), *ST->getRegisterInfo(),
                               *ST->getRegBankInfo());
 }

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -53,6 +53,8 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::SPV_INTEL_subgroups},
         {"SPV_INTEL_media_block_io",
          SPIRV::Extension::Extension::SPV_INTEL_media_block_io},
+        {"SPV_INTEL_memory_access_aliasing",
+         SPIRV::Extension::Extension::SPV_INTEL_memory_access_aliasing},
         {"SPV_INTEL_joint_matrix",
          SPIRV::Extension::Extension::SPV_INTEL_joint_matrix},
         {"SPV_KHR_uniform_group_instructions",

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -151,6 +151,7 @@ class SPIRVEmitIntrinsics
                                         unsigned OperandToReplace,
                                         IRBuilder<> &B);
   void insertPtrCastOrAssignTypeInstr(Instruction *I, IRBuilder<> &B);
+  bool shouldTryToAddMemAliasingDecoration(Instruction *Inst);
   void insertSpirvDecorations(Instruction *I, IRBuilder<> &B);
   void processGlobalValue(GlobalVariable &GV, IRBuilder<> &B);
   void processParamTypes(Function *F, IRBuilder<> &B);
@@ -1169,6 +1170,7 @@ void SPIRVEmitIntrinsics::replaceMemInstrUses(Instruction *Old,
       llvm_unreachable("illegal aggregate intrinsic user");
     }
   }
+  New->copyMetadata(*Old);
   Old->eraseFromParent();
 }
 
@@ -1752,6 +1754,7 @@ Instruction *SPIRVEmitIntrinsics::visitStoreInst(StoreInst &I) {
       Intrinsic::spv_store, {I.getValueOperand()->getType(), PtrOp->getType()},
       {I.getValueOperand(), PtrOp, B.getInt16(Flags),
        B.getInt8(I.getAlign().value())});
+  NewI->copyMetadata(I);
   I.eraseFromParent();
   return NewI;
 }
@@ -1955,13 +1958,64 @@ void SPIRVEmitIntrinsics::insertAssignTypeIntrs(Instruction *I,
   }
 }
 
+bool SPIRVEmitIntrinsics::shouldTryToAddMemAliasingDecoration(
+    Instruction *Inst) {
+  const SPIRVSubtarget *STI = TM->getSubtargetImpl(*Inst->getFunction());
+  if (!STI->canUseExtension(SPIRV::Extension::SPV_INTEL_memory_access_aliasing))
+    return false;
+  // Add aliasing decorations to internal load and store intrinsics
+  // and atomic instructions, skipping atomic store as it won't have ID to
+  // attach the decoration.
+  CallInst *CI = dyn_cast<CallInst>(Inst);
+  if (!CI)
+    return false;
+  if (Function *Fun = CI->getCalledFunction()) {
+    if (Fun->isIntrinsic()) {
+      switch (Fun->getIntrinsicID()) {
+      case Intrinsic::spv_load:
+      case Intrinsic::spv_store:
+        return true;
+      default:
+        return false;
+      }
+    }
+    std::string Name = getOclOrSpirvBuiltinDemangledName(Fun->getName());
+    const std::string Prefix = "__spirv_Atomic";
+    const bool IsAtomic = Name.find(Prefix) == 0;
+
+    if (!Fun->getReturnType()->isVoidTy() && IsAtomic)
+      return true;
+  }
+  return false;
+}
+
 void SPIRVEmitIntrinsics::insertSpirvDecorations(Instruction *I,
                                                  IRBuilder<> &B) {
   if (MDNode *MD = I->getMetadata("spirv.Decorations")) {
     setInsertPointAfterDef(B, I);
     B.CreateIntrinsic(Intrinsic::spv_assign_decoration, {I->getType()},
                       {I, MetadataAsValue::get(I->getContext(), MD)});
   }
+  // Lower alias.scope/noalias metadata
+  {
+    auto processMemAliasingDecoration = [&](unsigned Kind) {
+      if (MDNode *AliasListMD = I->getMetadata(Kind)) {
+        if (shouldTryToAddMemAliasingDecoration(I)) {
+          uint32_t Dec = Kind == LLVMContext::MD_alias_scope
+                             ? SPIRV::Decoration::AliasScopeINTEL
+                             : SPIRV::Decoration::NoAliasINTEL;
+          SmallVector<Value *, 3> Args = {
+              I, ConstantInt::get(B.getInt32Ty(), Dec),
+              MetadataAsValue::get(I->getContext(), AliasListMD)};
+          setInsertPointAfterDef(B, I);
+          B.CreateIntrinsic(Intrinsic::spv_assign_aliasing_decoration,
+                            {I->getType()}, {Args});
+        }
+      }
+    };
+    processMemAliasingDecoration(LLVMContext::MD_alias_scope);
+    processMemAliasingDecoration(LLVMContext::MD_noalias);
+  }
 }
 
 void SPIRVEmitIntrinsics::processInstrAfterVisit(Instruction *I,

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -1759,6 +1759,75 @@ LLT SPIRVGlobalRegistry::getRegType(SPIRVType *SpvType) const {
   return LLT::scalar(64);
 }
 
+// Aliasing list MD contains several scope MD nodes whithin it. Each scope MD
+// has a selfreference and an extra MD node for aliasing domain and also it
+// can contain an optional string operand. Domain MD contains a self-reference
+// with an optional string operand. Here we unfold the list, creating SPIR-V
+// aliasing instructions.
+// TODO: add support for an optional string operand.
+MachineInstr *SPIRVGlobalRegistry::getOrAddMemAliasingINTELInst(
+    MachineIRBuilder &MIRBuilder, const MDNode *AliasingListMD) {
+  if (AliasingListMD->getNumOperands() == 0)
+    return nullptr;
+  if (auto L = AliasInstMDMap.find(AliasingListMD); L != AliasInstMDMap.end())
+    return L->second;
+
+  SmallVector<MachineInstr *> ScopeList;
+  MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+  for (const MDOperand &MDListOp : AliasingListMD->operands()) {
+    if (MDNode *ScopeMD = dyn_cast<MDNode>(MDListOp)) {
+      if (ScopeMD->getNumOperands() < 2)
+        return nullptr;
+      MDNode *DomainMD = dyn_cast<MDNode>(ScopeMD->getOperand(1));
+      if (!DomainMD)
+        return nullptr;
+      auto *Domain = [&] {
+        auto D = AliasInstMDMap.find(DomainMD);
+        if (D != AliasInstMDMap.end())
+          return D->second;
+        const Register Ret = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+        auto MIB =
+            MIRBuilder.buildInstr(SPIRV::OpAliasDomainDeclINTEL).addDef(Ret);
+        return MIB.getInstr();
+      }();
+      AliasInstMDMap.insert(std::make_pair(DomainMD, Domain));
+      auto *Scope = [&] {
+        auto S = AliasInstMDMap.find(ScopeMD);
+        if (S != AliasInstMDMap.end())
+          return S->second;
+        const Register Ret = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+        auto MIB = MIRBuilder.buildInstr(SPIRV::OpAliasScopeDeclINTEL)
+                       .addDef(Ret)
+                       .addUse(Domain->getOperand(0).getReg());
+        return MIB.getInstr();
+      }();
+      AliasInstMDMap.insert(std::make_pair(ScopeMD, Scope));
+      ScopeList.push_back(Scope);
+    }
+  }
+
+  const Register Ret = MRI->createVirtualRegister(&SPIRV::IDRegClass);
+  auto MIB =
+      MIRBuilder.buildInstr(SPIRV::OpAliasScopeListDeclINTEL).addDef(Ret);
+  for (auto *Scope : ScopeList)
+    MIB.addUse(Scope->getOperand(0).getReg());
+  auto List = MIB.getInstr();
+  AliasInstMDMap.insert(std::make_pair(AliasingListMD, List));
+  return List;
+}
+
+void SPIRVGlobalRegistry::buildMemAliasingOpDecorate(
+    Register Reg, MachineIRBuilder &MIRBuilder, uint32_t Dec,
+    const MDNode *AliasingListMD) {
+  MachineInstr *AliasList =
+      getOrAddMemAliasingINTELInst(MIRBuilder, AliasingListMD);
+  if (!AliasList)
+    return;
+  MIRBuilder.buildInstr(SPIRV::OpDecorate)
+      .addUse(Reg)
+      .addImm(Dec)
+      .addUse(AliasList->getOperand(0).getReg());
+}
 void SPIRVGlobalRegistry::replaceAllUsesWith(Value *Old, Value *New,
                                              bool DeleteOld) {
   Old->replaceAllUsesWith(New);

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -92,6 +92,9 @@ class SPIRVGlobalRegistry {
   // Maps OpVariable and OpFunction-related v-regs to its LLVM IR definition.
   DenseMap<std::pair<const MachineFunction *, Register>, const Value *> Reg2GO;
 
+  // map of aliasing decorations to aliasing metadata
+  std::unordered_map<const MDNode *, MachineInstr *> AliasInstMDMap;
+
   // Add a new OpTypeXXX instruction without checking for duplicates.
   SPIRVType *createSPIRVType(const Type *Type, MachineIRBuilder &MIRBuilder,
                              SPIRV::AccessQualifier::AccessQualifier AQ,
@@ -621,6 +624,10 @@ class SPIRVGlobalRegistry {
   const TargetRegisterClass *getRegClass(SPIRVType *SpvType) const;
   LLT getRegType(SPIRVType *SpvType) const;
 
+  MachineInstr *getOrAddMemAliasingINTELInst(MachineIRBuilder &MIRBuilder,
+                                             const MDNode *AliasingListMD);
+  void buildMemAliasingOpDecorate(Register Reg, MachineIRBuilder &MIRBuilder,
+                                  uint32_t Dec, const MDNode *GVarMD);
   // Replace all uses of a |Old| with |New| updates the global registry type
   // mappings.
   void replaceAllUsesWith(Value *Old, Value *New, bool DeleteOld = true);

llvm/lib/Target/SPIRV/SPIRVInstrInfo.cpp

@@ -97,6 +97,17 @@ bool SPIRVInstrInfo::isDecorationInstr(const MachineInstr &MI) const {
   }
 }
 
+bool SPIRVInstrInfo::isAliasingInstr(const MachineInstr &MI) const {
+  switch (MI.getOpcode()) {
+  case SPIRV::OpAliasDomainDeclINTEL:
+  case SPIRV::OpAliasScopeDeclINTEL:
+  case SPIRV::OpAliasScopeListDeclINTEL:
+    return true;
+  default:
+    return false;
+  }
+}
+
 bool SPIRVInstrInfo::isHeaderInstr(const MachineInstr &MI) const {
   switch (MI.getOpcode()) {
   case SPIRV::OpCapability:
@@ -115,7 +126,8 @@ bool SPIRVInstrInfo::isHeaderInstr(const MachineInstr &MI) const {
   case SPIRV::OpModuleProcessed:
     return true;
   default:
-    return isTypeDeclInstr(MI) || isConstantInstr(MI) || isDecorationInstr(MI);
+    return isTypeDeclInstr(MI) || isConstantInstr(MI) ||
+           isDecorationInstr(MI) || isAliasingInstr(MI);
   }
 }
 

llvm/lib/Target/SPIRV/SPIRVInstrInfo.h

@@ -34,6 +34,7 @@ class SPIRVInstrInfo : public SPIRVGenInstrInfo {
   bool isInlineAsmDefInstr(const MachineInstr &MI) const;
   bool isTypeDeclInstr(const MachineInstr &MI) const;
   bool isDecorationInstr(const MachineInstr &MI) const;
+  bool isAliasingInstr(const MachineInstr &MI) const;
   bool canUseFastMathFlags(const MachineInstr &MI) const;
   bool canUseNSW(const MachineInstr &MI) const;
   bool canUseNUW(const MachineInstr &MI) const;