Support SPV_INTEL_maximum_registers extension (#2344)

vmaksimo · web-flow · commit 92ea64b77207 · 2024-02-26T18:37:25.000+01:00
Spec: KhronosGroup/SPIRV-Registry#235
diff --git a/include/LLVMSPIRVExtensions.inc b/include/LLVMSPIRVExtensions.inc
@@ -71,3 +71,4 @@ EXT(SPV_INTEL_fp_max_error)
 EXT(SPV_INTEL_cache_controls)
 EXT(SPV_INTEL_subgroup_requirements)
 EXT(SPV_INTEL_task_sequence)
+EXT(SPV_INTEL_maximum_registers)
diff --git a/lib/SPIRV/SPIRVReader.cpp b/lib/SPIRV/SPIRVReader.cpp
@@ -4259,6 +4259,50 @@ bool SPIRVToLLVM::transMetadata() {
       F->setMetadata(kSPIR2MD::IntelFPGAIPInterface,
                      MDNode::get(*Context, InterfaceMDVec));
     }
+    if (auto *EM = BF->getExecutionMode(
+            internal::ExecutionModeMaximumRegistersINTEL)) {
+      NamedMDNode *ExecModeMD =
+          M->getOrInsertNamedMetadata(kSPIRVMD::ExecutionMode);
+
+      SmallVector<Metadata *, 4> ValueVec;
+      ValueVec.push_back(ConstantAsMetadata::get(F));
+      ValueVec.push_back(
+          ConstantAsMetadata::get(getUInt32(M, EM->getExecutionMode())));
+      ValueVec.push_back(
+          ConstantAsMetadata::get(getUInt32(M, EM->getLiterals()[0])));
+      ExecModeMD->addOperand(MDNode::get(*Context, ValueVec));
+    }
+    if (auto *EM = BF->getExecutionMode(
+            internal::ExecutionModeMaximumRegistersIdINTEL)) {
+      NamedMDNode *ExecModeMD =
+          M->getOrInsertNamedMetadata(kSPIRVMD::ExecutionMode);
+
+      SmallVector<Metadata *, 4> ValueVec;
+      ValueVec.push_back(ConstantAsMetadata::get(F));
+      ValueVec.push_back(
+          ConstantAsMetadata::get(getUInt32(M, EM->getExecutionMode())));
+
+      auto *ExecOp = BF->getModule()->getValue(EM->getLiterals()[0]);
+      ValueVec.push_back(
+          MDNode::get(*Context, ConstantAsMetadata::get(cast<ConstantInt>(
+                                    transValue(ExecOp, nullptr, nullptr)))));
+      ExecModeMD->addOperand(MDNode::get(*Context, ValueVec));
+    }
+    if (auto *EM = BF->getExecutionMode(
+            internal::ExecutionModeNamedMaximumRegistersINTEL)) {
+      NamedMDNode *ExecModeMD =
+          M->getOrInsertNamedMetadata(kSPIRVMD::ExecutionMode);
+
+      SmallVector<Metadata *, 4> ValueVec;
+      ValueVec.push_back(ConstantAsMetadata::get(F));
+      ValueVec.push_back(
+          ConstantAsMetadata::get(getUInt32(M, EM->getExecutionMode())));
+
+      assert(EM->getLiterals()[0] == 0 &&
+             "Invalid named maximum number of registers");
+      ValueVec.push_back(MDString::get(*Context, "AutoINTEL"));
+      ExecModeMD->addOperand(MDNode::get(*Context, ValueVec));
+    }
   }
   NamedMDNode *MemoryModelMD =
       M->getOrInsertNamedMetadata(kSPIRVMD::MemoryModel);
diff --git a/lib/SPIRV/SPIRVWriter.cpp b/lib/SPIRV/SPIRVWriter.cpp
@@ -979,7 +979,10 @@ SPIRVFunction *LLVMToSPIRVBase::transFunctionDecl(Function *F) {
 
   transFPGAFunctionMetadata(BF, F);
 
-  transFunctionMetadataAsUserSemanticDecoration(BF, F);
+  if (BM->isAllowedToUseExtension(ExtensionID::SPV_INTEL_maximum_registers))
+    transFunctionMetadataAsExecutionMode(BF, F);
+  else
+    transFunctionMetadataAsUserSemanticDecoration(BF, F);
 
   transAuxDataInst(BF, F);
 
@@ -1120,6 +1123,38 @@ void LLVMToSPIRVBase::transFPGAFunctionMetadata(SPIRVFunction *BF,
     transMetadataDecorations(FDecoMD, BF);
 }
 
+void LLVMToSPIRVBase::transFunctionMetadataAsExecutionMode(SPIRVFunction *BF,
+                                                           Function *F) {
+  SmallVector<MDNode *, 1> RegisterAllocModeMDs;
+  F->getMetadata("RegisterAllocMode", RegisterAllocModeMDs);
+
+  for (unsigned I = 0; I < RegisterAllocModeMDs.size(); I++) {
+    auto *RegisterAllocMode = RegisterAllocModeMDs[I]->getOperand(0).get();
+    if (isa<MDString>(RegisterAllocMode)) {
+      StringRef Str = getMDOperandAsString(RegisterAllocModeMDs[I], 0);
+      internal::InternalNamedMaximumNumberOfRegisters NamedValue =
+          SPIRVNamedMaximumNumberOfRegistersNameMap::rmap(Str.str());
+      BF->addExecutionMode(BM->add(new SPIRVExecutionMode(
+          OpExecutionMode, BF,
+          internal::ExecutionModeNamedMaximumRegistersINTEL, NamedValue)));
+    } else if (isa<MDNode>(RegisterAllocMode)) {
+      auto *RegisterAllocNodeMDOp =
+          getMDOperandAsMDNode(RegisterAllocModeMDs[I], 0);
+      int Num = getMDOperandAsInt(RegisterAllocNodeMDOp, 0);
+      auto *Const =
+          BM->addConstant(transType(Type::getInt32Ty(F->getContext())), Num);
+      BF->addExecutionMode(BM->add(new SPIRVExecutionModeId(
+          BF, internal::ExecutionModeMaximumRegistersIdINTEL, Const->getId())));
+    } else {
+      int64_t RegisterAllocVal =
+          mdconst::dyn_extract<ConstantInt>(RegisterAllocMode)->getZExtValue();
+      BF->addExecutionMode(BM->add(new SPIRVExecutionMode(
+          OpExecutionMode, BF, internal::ExecutionModeMaximumRegistersINTEL,
+          RegisterAllocVal)));
+    }
+  }
+}
+
 void LLVMToSPIRVBase::transFunctionMetadataAsUserSemanticDecoration(
     SPIRVFunction *BF, Function *F) {
   if (auto *RegisterAllocModeMD = F->getMetadata("RegisterAllocMode")) {
diff --git a/lib/SPIRV/SPIRVWriter.h b/lib/SPIRV/SPIRVWriter.h
@@ -131,6 +131,7 @@ class LLVMToSPIRVBase : protected BuiltinCallHelper {
   SPIRVFunction *transFunctionDecl(Function *F);
   void transVectorComputeMetadata(Function *F);
   void transFPGAFunctionMetadata(SPIRVFunction *BF, Function *F);
+  void transFunctionMetadataAsExecutionMode(SPIRVFunction *BF, Function *F);
   void transFunctionMetadataAsUserSemanticDecoration(SPIRVFunction *BF,
                                                      Function *F);
   void transAuxDataInst(SPIRVFunction *BF, Function *F);
diff --git a/lib/SPIRV/libSPIRV/SPIRVEntry.cpp b/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
@@ -661,6 +661,9 @@ void SPIRVExecutionMode::decode(std::istream &I) {
   case ExecutionModeRegisterMapInterfaceINTEL:
   case ExecutionModeStreamingInterfaceINTEL:
   case spv::internal::ExecutionModeNamedSubgroupSizeINTEL:
+  case internal::ExecutionModeMaximumRegistersINTEL:
+  case internal::ExecutionModeMaximumRegistersIdINTEL:
+  case internal::ExecutionModeNamedMaximumRegistersINTEL:
     WordLiterals.resize(1);
     break;
   default:
diff --git a/lib/SPIRV/libSPIRV/SPIRVEntry.h b/lib/SPIRV/libSPIRV/SPIRVEntry.h
@@ -696,6 +696,17 @@ class SPIRVExecutionMode : public SPIRVAnnotation {
     }
   }
 
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    switch (static_cast<unsigned>(ExecMode)) {
+    case internal::ExecutionModeMaximumRegistersINTEL:
+    case internal::ExecutionModeMaximumRegistersIdINTEL:
+    case internal::ExecutionModeNamedMaximumRegistersINTEL:
+      return ExtensionID::SPV_INTEL_maximum_registers;
+    default:
+      return {};
+    }
+  }
+
 protected:
   _SPIRV_DCL_ENCDEC
   SPIRVExecutionModeKind ExecMode;
@@ -757,6 +768,11 @@ class SPIRVComponentExecutionModes {
       return IsDenorm(EMK) || IsRoundingMode(EMK) || IsFPMode(EMK) ||
              IsOtherFP(EMK);
     };
+    auto IsMaxRegisters = [&](auto EMK) {
+      return EMK == internal::ExecutionModeMaximumRegistersINTEL ||
+             EMK == internal::ExecutionModeMaximumRegistersIdINTEL ||
+             EMK == internal::ExecutionModeNamedMaximumRegistersINTEL;
+    };
     auto IsCompatible = [&](SPIRVExecutionMode *EM0, SPIRVExecutionMode *EM1) {
       if (EM0->getTargetId() != EM1->getTargetId())
         return true;
@@ -770,7 +786,8 @@ class SPIRVComponentExecutionModes {
         return true;
       return !(IsDenorm(EMK0) && IsDenorm(EMK1)) &&
              !(IsRoundingMode(EMK0) && IsRoundingMode(EMK1)) &&
-             !(IsFPMode(EMK0) && IsFPMode(EMK1));
+             !(IsFPMode(EMK0) && IsFPMode(EMK1)) &&
+             !(IsMaxRegisters(EMK0) && IsMaxRegisters(EMK1));
     };
     for (auto I = ExecModes.begin(); I != ExecModes.end(); ++I) {
       assert(IsCompatible(ExecMode, (*I).second) &&
diff --git a/lib/SPIRV/libSPIRV/SPIRVEnum.h b/lib/SPIRV/libSPIRV/SPIRVEnum.h
@@ -293,6 +293,12 @@ template <> inline void SPIRVMap<SPIRVExecutionModeKind, SPIRVCapVec>::init() {
                {CapabilityVectorComputeINTEL});
   ADD_VEC_INIT(internal::ExecutionModeNamedSubgroupSizeINTEL,
                {internal::CapabilitySubgroupRequirementsINTEL});
+  ADD_VEC_INIT(internal::ExecutionModeMaximumRegistersINTEL,
+               {internal::CapabilityRegisterLimitsINTEL});
+  ADD_VEC_INIT(internal::ExecutionModeMaximumRegistersIdINTEL,
+               {internal::CapabilityRegisterLimitsINTEL});
+  ADD_VEC_INIT(internal::ExecutionModeNamedMaximumRegistersINTEL,
+               {internal::CapabilityRegisterLimitsINTEL});
 }
 
 template <> inline void SPIRVMap<SPIRVMemoryModelKind, SPIRVCapVec>::init() {
diff --git a/lib/SPIRV/libSPIRV/SPIRVIsValidEnum.h b/lib/SPIRV/libSPIRV/SPIRVIsValidEnum.h
@@ -73,6 +73,9 @@ inline bool isValid(spv::ExecutionModel V) {
   case ExecutionModelCallableKHR:
   case ExecutionModeRegisterMapInterfaceINTEL:
   case ExecutionModeStreamingInterfaceINTEL:
+  case internal::ExecutionModeMaximumRegistersINTEL:
+  case internal::ExecutionModeMaximumRegistersIdINTEL:
+  case internal::ExecutionModeNamedMaximumRegistersINTEL:
     return true;
   default:
     return false;
diff --git a/lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h b/lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h
@@ -672,6 +672,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(internal::CapabilitySubgroupRequirementsINTEL,
       "SubgroupRequirementsINTEL");
   add(internal::CapabilityTaskSequenceINTEL, "TaskSequenceINTEL");
+  add(internal::CapabilityRegisterLimitsINTEL, "RegisterLimitsINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 
@@ -695,6 +696,14 @@ template <> inline void SPIRVMap<HostAccessQualifier, std::string>::init() {
 }
 SPIRV_DEF_NAMEMAP(HostAccessQualifier, SPIRVHostAccessQualifierNameMap)
 
+template <>
+inline void
+SPIRVMap<internal::InternalNamedMaximumNumberOfRegisters, std::string>::init() {
+  add(internal::NamedMaximumNumberOfRegistersAutoINTEL, "AutoINTEL");
+}
+SPIRV_DEF_NAMEMAP(internal::InternalNamedMaximumNumberOfRegisters,
+                  SPIRVNamedMaximumNumberOfRegistersNameMap);
+
 } /* namespace SPIRV */
 
 #endif // SPIRV_LIBSPIRV_SPIRVNAMEMAPENUM_H
diff --git a/lib/SPIRV/libSPIRV/SPIRVStream.cpp b/lib/SPIRV/libSPIRV/SPIRVStream.cpp
@@ -147,6 +147,7 @@ SPIRV_DEF_ENCDEC(SPIRVDebugExtOpKind)
 SPIRV_DEF_ENCDEC(NonSemanticAuxDataOpKind)
 SPIRV_DEF_ENCDEC(InitializationModeQualifier)
 SPIRV_DEF_ENCDEC(HostAccessQualifier)
+SPIRV_DEF_ENCDEC(internal::InternalNamedMaximumNumberOfRegisters)
 SPIRV_DEF_ENCDEC(LinkageType)
 
 // Read a string with padded 0's at the end so that they form a stream of
diff --git a/lib/SPIRV/libSPIRV/SPIRVStream.h b/lib/SPIRV/libSPIRV/SPIRVStream.h
@@ -231,6 +231,7 @@ SPIRV_DEC_ENCDEC(SPIRVDebugExtOpKind)
 SPIRV_DEC_ENCDEC(NonSemanticAuxDataOpKind)
 SPIRV_DEC_ENCDEC(InitializationModeQualifier)
 SPIRV_DEC_ENCDEC(HostAccessQualifier)
+SPIRV_DEC_ENCDEC(internal::InternalNamedMaximumNumberOfRegisters)
 SPIRV_DEC_ENCDEC(LinkageType)
 
 const SPIRVEncoder &operator<<(const SPIRVEncoder &O, const std::string &Str);
diff --git a/lib/SPIRV/libSPIRV/spirv_internal.hpp b/lib/SPIRV/libSPIRV/spirv_internal.hpp
@@ -122,14 +122,18 @@ enum InternalCapability {
   ICapabilityJointMatrixPackedInt2ComponentTypeINTEL = 6438,
   ICapabilityJointMatrixPackedInt4ComponentTypeINTEL = 6439,
   ICapabilityCacheControlsINTEL = 6441,
-  ICapabilitySubgroupRequirementsINTEL = 6445
+  ICapabilitySubgroupRequirementsINTEL = 6445,
+  ICapRegisterLimitsINTEL = 6460
 };
 
 enum InternalFunctionControlMask { IFunctionControlOptNoneINTELMask = 0x10000 };
 
 enum InternalExecutionMode {
   IExecModeFastCompositeKernelINTEL = 6088,
-  IExecModeNamedSubgroupSizeINTEL = 6446
+  IExecModeNamedSubgroupSizeINTEL = 6446,
+  IExecModeMaximumRegistersINTEL = 6461,
+  IExecModeMaximumRegistersIdINTEL = 6462,
+  IExecModeNamedMaximumRegistersINTEL = 6463
 };
 
 constexpr LinkageType LinkageTypeInternal =
@@ -157,6 +161,10 @@ enum InternalBuiltIn {
   IBuiltInGlobalHWThreadIDINTEL = 6136,
 };
 
+enum InternalNamedMaximumNumberOfRegisters {
+  NamedMaximumNumberOfRegistersAutoINTEL = 0,
+};
+
 #define _SPIRV_OP(x, y) constexpr x x##y = static_cast<x>(I##x##y);
 _SPIRV_OP(Capability, JointMatrixINTEL)
 _SPIRV_OP(Capability, JointMatrixWIInstructionsINTEL)
@@ -286,12 +294,20 @@ constexpr Capability CapabilityBfloat16ConversionINTEL =
     static_cast<Capability>(ICapBfloat16ConversionINTEL);
 constexpr Capability CapabilityGlobalVariableDecorationsINTEL =
     static_cast<Capability>(ICapGlobalVariableDecorationsINTEL);
+constexpr Capability CapabilityRegisterLimitsINTEL =
+    static_cast<Capability>(ICapRegisterLimitsINTEL);
 
 constexpr FunctionControlMask FunctionControlOptNoneINTELMask =
     static_cast<FunctionControlMask>(IFunctionControlOptNoneINTELMask);
 
 constexpr ExecutionMode ExecutionModeFastCompositeKernelINTEL =
     static_cast<ExecutionMode>(IExecModeFastCompositeKernelINTEL);
+constexpr ExecutionMode ExecutionModeMaximumRegistersINTEL =
+    static_cast<ExecutionMode>(IExecModeMaximumRegistersINTEL);
+constexpr ExecutionMode ExecutionModeMaximumRegistersIdINTEL =
+    static_cast<ExecutionMode>(IExecModeMaximumRegistersIdINTEL);
+constexpr ExecutionMode ExecutionModeNamedMaximumRegistersINTEL =
+    static_cast<ExecutionMode>(IExecModeNamedMaximumRegistersINTEL);
 
 constexpr ExecutionMode ExecutionModeNamedSubgroupSizeINTEL =
     static_cast<ExecutionMode>(IExecModeNamedSubgroupSizeINTEL);
diff --git a/test/extensions/INTEL/SPV_INTEL_maximum_registers/registerallocmode_maxreg_extension.ll b/test/extensions/INTEL/SPV_INTEL_maximum_registers/registerallocmode_maxreg_extension.ll
@@ -0,0 +1,85 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv -spirv-text --spirv-ext=+SPV_INTEL_maximum_registers %t.bc
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_maximum_registers -o %t.spv
+; RUN: llvm-spirv -r %t.spv -spirv-target-env=SPV-IR -o - | llvm-dis -o %t.rev.ll
+; RUN: FileCheck < %t.rev.ll %s --check-prefix=CHECK-LLVM
+
+; CHECK-SPIRV: EntryPoint [[#]] [[#FUNC0:]] "main_l3"
+; CHECK-SPIRV: EntryPoint [[#]] [[#FUNC1:]] "main_l6"
+; CHECK-SPIRV: EntryPoint [[#]] [[#FUNC2:]] "main_l9"
+; CHECK-SPIRV: EntryPoint [[#]] [[#FUNC3:]] "main_l13"
+; CHECK-SPIRV: EntryPoint [[#]] [[#FUNC4:]] "main_l19"
+
+; CHECK-SPIRV: ExecutionMode [[#FUNC0]] 6461 2
+; CHECK-SPIRV: ExecutionMode [[#FUNC1]] 6461 1
+; CHECK-SPIRV: ExecutionMode [[#FUNC2]] 6463 0
+; CHECK-SPIRV: ExecutionModeId [[#FUNC3]] 6462 [[#Const3:]]
+; CHECK-SPIRV: TypeInt [[#TypeInt:]] 32 0
+; CHECK-SPIRV: Constant [[#TypeInt]] [[#Const3]] 3
+
+; CHECK-SPIRV-NOT: ExecutionMode [[#FUNC4]]
+
+; CHECK-LLVM: !spirv.ExecutionMode = !{![[#FLAG0:]], ![[#FLAG1:]], ![[#FLAG2:]], ![[#FLAG3:]]}
+; CHECK-LLVM: ![[#FLAG0]] = !{ptr @main_l3, i32 6461, i32 2}
+; CHECK-LLVM: ![[#FLAG1]] = !{ptr @main_l6, i32 6461, i32 1}
+; CHECK-LLVM: ![[#FLAG2]] = !{ptr @main_l9, i32 6463, !"AutoINTEL"}
+; CHECK-LLVM: ![[#FLAG3]] = !{ptr @main_l13, i32 6462, ![[#VAL:]]}
+; CHECK-LLVM: ![[#VAL]] = !{i32 3}
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64"
+
+; Function Attrs: noinline nounwind optnone
+define weak dso_local spir_kernel void @main_l3() #0 !RegisterAllocMode !10 {
+newFuncRoot:
+  ret void
+}
+
+; Function Attrs: noinline nounwind optnone
+define weak dso_local spir_kernel void @main_l6() #0 !RegisterAllocMode !11 {
+newFuncRoot:
+  ret void
+}
+
+; Function Attrs: noinline nounwind optnone
+define weak dso_local spir_kernel void @main_l9() #0 !RegisterAllocMode !12 {
+newFuncRoot:
+  ret void
+}
+
+; Function Attrs: noinline nounwind optnone
+define weak dso_local spir_kernel void @main_l13() #0 !RegisterAllocMode !13 {
+newFuncRoot:
+  ret void
+}
+
+; Function Attrs: noinline nounwind optnone
+define weak dso_local spir_kernel void @main_l19() #0 {
+newFuncRoot:
+  ret void
+}
+
+attributes #0 = { noinline nounwind optnone }
+
+
+!opencl.compiler.options = !{!0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0}
+!spirv.Source = !{!2, !3, !3, !3, !3, !3, !2, !3, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2}
+!llvm.module.flags = !{!4, !5, !6, !7, !8}
+!spirv.MemoryModel = !{!9, !9, !9, !9, !9, !9}
+!spirv.ExecutionMode = !{}
+
+!0 = !{}
+!2 = !{i32 4, i32 200000}
+!3 = !{i32 3, i32 200000}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 7, !"openmp", i32 50}
+!6 = !{i32 7, !"openmp-device", i32 50}
+!7 = !{i32 8, !"PIC Level", i32 2}
+!8 = !{i32 7, !"frame-pointer", i32 2}
+!9 = !{i32 2, i32 2}
+!10 = !{i32 2}
+!11 = !{i32 1}
+!12 = !{!"AutoINTEL"}
+!13 = !{!14}
+!14 = !{i32 3}
