Add support for SPV_INTEL_subgroup_requirements (#2317)

Spec: intel/llvm#11301

More accurately, this PR adds support for the named subgroup related features of SPV_INTEL_subgroup_requirements to support implementation of sycl_ext_named_sub_group_sizes (also see intel/llvm#12335). The features related to subgroup lane mapping are not added yet.

Author: jzc

include/LLVMSPIRVExtensions.inc

@@ -69,3 +69,4 @@ EXT(SPV_INTEL_fpga_argument_interfaces)
 EXT(SPV_INTEL_fpga_latency_control)
 EXT(SPV_INTEL_fp_max_error)
 EXT(SPV_INTEL_cache_controls)
+EXT(SPV_INTEL_subgroup_requirements)

lib/SPIRV/PreprocessMetadata.cpp

@@ -167,10 +167,19 @@ void PreprocessMetadataBase::visit(Module *M) {
 
     // !{void (i32 addrspace(1)*)* @kernel, i32 35, i32 size}
     if (MDNode *ReqdSubgroupSize = Kernel.getMetadata(kSPIR2MD::SubgroupSize)) {
+      // A primary named subgroup size is encoded as
+      // the metadata intel_reqd_sub_group_size with value 0.
+      auto Val = getMDOperandAsInt(ReqdSubgroupSize, 0);
+      if (Val == 0)
+        EM.addOp()
+            .add(&Kernel)
+            .add(spv::internal::ExecutionModeNamedSubgroupSizeINTEL)
+            .add(/* PrimarySubgroupSizeINTEL = */ 0U)
+            .done();
       EM.addOp()
           .add(&Kernel)
           .add(spv::ExecutionModeSubgroupSize)
-          .add(getMDOperandAsInt(ReqdSubgroupSize, 0))
+          .add(Val)
           .done();
     }
 

lib/SPIRV/SPIRVReader.cpp

@@ -4240,6 +4240,18 @@ bool SPIRVToLLVM::transMetadata() {
           ConstantAsMetadata::get(getUInt32(M, EM->getLiterals()[0]));
       F->setMetadata(kSPIR2MD::SubgroupSize, MDNode::get(*Context, SizeMD));
     }
+    // Generate metadata for intel_reqd_sub_group_size
+    if (BF->getExecutionMode(internal::ExecutionModeNamedSubgroupSizeINTEL)) {
+      // For now, there is only one named sub group size: primary, which is
+      // represented as a value of 0 as the argument of the OpExecutionMode.
+      assert(BF->getExecutionMode(internal::ExecutionModeNamedSubgroupSizeINTEL)
+                     ->getLiterals()[0] == 0 &&
+             "Invalid named sub group size");
+      // On the LLVM IR side, this is represented as the metadata
+      // intel_reqd_sub_group_size with value 0.
+      auto *SizeMD = ConstantAsMetadata::get(getUInt32(M, 0));
+      F->setMetadata(kSPIR2MD::SubgroupSize, MDNode::get(*Context, SizeMD));
+    }
     // Generate metadata for max_work_group_size
     if (auto *EM = BF->getExecutionMode(ExecutionModeMaxWorkgroupSizeINTEL)) {
       F->setMetadata(kSPIR2MD::MaxWGSize,

lib/SPIRV/SPIRVWriter.cpp

@@ -5835,6 +5835,12 @@ bool LLVMToSPIRVBase::transExecutionMode() {
           BF->addExecutionMode(BM->add(new SPIRVExecutionMode(
               OpExecutionMode, BF, static_cast<ExecutionMode>(EMode))));
       } break;
+      case spv::internal::ExecutionModeNamedSubgroupSizeINTEL: {
+        if (!BM->isAllowedToUseExtension(
+                ExtensionID::SPV_INTEL_subgroup_requirements))
+          break;
+        AddSingleArgExecutionMode(static_cast<ExecutionMode>(EMode));
+      } break;
       default:
         llvm_unreachable("invalid execution mode");
       }

lib/SPIRV/libSPIRV/SPIRVEntry.cpp

@@ -660,6 +660,7 @@ void SPIRVExecutionMode::decode(std::istream &I) {
   case ExecutionModeSchedulerTargetFmaxMhzINTEL:
   case ExecutionModeRegisterMapInterfaceINTEL:
   case ExecutionModeStreamingInterfaceINTEL:
+  case spv::internal::ExecutionModeNamedSubgroupSizeINTEL:
     WordLiterals.resize(1);
     break;
   default:

lib/SPIRV/libSPIRV/SPIRVEntry.h

@@ -885,6 +885,8 @@ class SPIRVCapability : public SPIRVEntryNoId<OpCapability> {
       return ExtensionID::SPV_INTEL_vector_compute;
     case internal::CapabilityFastCompositeINTEL:
       return ExtensionID::SPV_INTEL_fast_composite;
+    case internal::CapabilitySubgroupRequirementsINTEL:
+      return ExtensionID::SPV_INTEL_subgroup_requirements;
     default:
       return {};
     }

lib/SPIRV/libSPIRV/SPIRVEnum.h

@@ -291,6 +291,8 @@ template <> inline void SPIRVMap<SPIRVExecutionModeKind, SPIRVCapVec>::init() {
                {CapabilityFPGAKernelAttributesINTEL});
   ADD_VEC_INIT(ExecutionModeNamedBarrierCountINTEL,
                {CapabilityVectorComputeINTEL});
+  ADD_VEC_INIT(internal::ExecutionModeNamedSubgroupSizeINTEL,
+               {internal::CapabilitySubgroupRequirementsINTEL});
 }
 
 template <> inline void SPIRVMap<SPIRVMemoryModelKind, SPIRVCapVec>::init() {

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -668,6 +668,8 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
       "CooperativeMatrixInvocationInstructionsINTEL");
   add(internal::CapabilityCooperativeMatrixCheckedInstructionsINTEL,
       "CooperativeMatrixCheckedInstructionsINTEL");
+  add(internal::CapabilitySubgroupRequirementsINTEL,
+      "SubgroupRequirementsINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -117,13 +117,15 @@ enum InternalCapability {
   ICapabilityJointMatrixBF16ComponentTypeINTEL = 6437,
   ICapabilityJointMatrixPackedInt2ComponentTypeINTEL = 6438,
   ICapabilityJointMatrixPackedInt4ComponentTypeINTEL = 6439,
-  ICapabilityCacheControlsINTEL = 6441
+  ICapabilityCacheControlsINTEL = 6441,
+  ICapabilitySubgroupRequirementsINTEL = 6445
 };
 
 enum InternalFunctionControlMask { IFunctionControlOptNoneINTELMask = 0x10000 };
 
 enum InternalExecutionMode {
   IExecModeFastCompositeKernelINTEL = 6088,
+  IExecModeNamedSubgroupSizeINTEL = 6446
 };
 
 constexpr LinkageType LinkageTypeInternal =
@@ -211,6 +213,8 @@ _SPIRV_OP(Capability, TensorFloat32RoundingINTEL)
 _SPIRV_OP(Op, RoundFToTF32INTEL)
 
 _SPIRV_OP(Capability, CacheControlsINTEL)
+
+_SPIRV_OP(Capability, SubgroupRequirementsINTEL)
 #undef _SPIRV_OP
 
 constexpr SourceLanguage SourceLanguagePython =
@@ -296,6 +300,9 @@ constexpr FunctionControlMask FunctionControlOptNoneINTELMask =
 constexpr ExecutionMode ExecutionModeFastCompositeKernelINTEL =
     static_cast<ExecutionMode>(IExecModeFastCompositeKernelINTEL);
 
+constexpr ExecutionMode ExecutionModeNamedSubgroupSizeINTEL =
+    static_cast<ExecutionMode>(IExecModeNamedSubgroupSizeINTEL);
+
 } // namespace internal
 } // namespace spv
 

test/extensions/INTEL/SPV_INTEL_subgroup_requirements/named_sub_group_sizes.ll

@@ -0,0 +1,54 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_subgroup_requirements -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; RUN: llvm-spirv %t.bc -o %t2.spv
+; RUN: llvm-spirv %t2.spv -to-text -o %t2.spt
+; RUN: FileCheck < %t2.spt %s --check-prefix=CHECK-SPIRV-2
+
+; RUN: llvm-spirv -r %t2.spv -o %t2.rev.bc
+; RUN: llvm-dis < %t2.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; CHECK-SPIRV: Capability SubgroupRequirementsINTEL
+; CHECK-SPIRV: Extension "SPV_INTEL_subgroup_requirements"
+; CHECK-SPIRV: EntryPoint 6 [[kernel:[0-9]+]] "_ZTSZ4mainE7Kernel1"
+; CHECK-SPIRV: ExecutionMode [[kernel]] 6446 0
+
+; CHECK-LLVM: spir_kernel void @_ZTSZ4mainE7Kernel1() {{.*}} !intel_reqd_sub_group_size ![[MD:[0-9]+]]
+; CHECK-LLVM: ![[MD]] = !{i32 0}
+
+; CHECK-SPIRV-2-NOT: Capability SubgroupRequirementsINTEL
+; CHECK-SPIRV-2-NOT: Extension "SPV_INTEL_subgroup_requirements"
+; CHECK-SPIRV-2: EntryPoint 6 [[kernel:[0-9]+]] "_ZTSZ4mainE7Kernel1"
+; CHECK-SPIRV-2: ExecutionMode [[kernel]] 35 0
+
+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-unknown-unknown"
+
+$_ZTSZ4mainE7Kernel1 = comdat any
+
+; Function Attrs: mustprogress norecurse nounwind
+define weak_odr dso_local spir_kernel void @_ZTSZ4mainE7Kernel1() local_unnamed_addr #0 comdat !srcloc !5 !kernel_arg_buffer_location !6 !sycl_fixed_targets !6 !sycl_kernel_omit_args !6 !intel_reqd_sub_group_size !7 {
+entry:
+  ret void
+}
+
+attributes #0 = { mustprogress norecurse nounwind "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "sycl-module-id"="ex.cpp" "sycl-optlevel"="2" "sycl-sub-group-size"="0" "uniform-work-group-size"="true" }
+
+!llvm.module.flags = !{!0, !1}
+!opencl.spir.version = !{!2}
+!spirv.Source = !{!3}
+!llvm.ident = !{!4}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"frame-pointer", i32 2}
+!2 = !{i32 1, i32 2}
+!3 = !{i32 4, i32 100000}
+!4 = !{!"clang version 18.0.0git (/ws/llvm/clang 8fd29b3c2aa9f9ce163be557b51de39c95aaf230)"}
+!5 = !{i32 358}
+!6 = !{}
+!7 = !{i32 0}