[SPIR-V] Generate SPIR-V instructions when 'enable-maximal-reconvergence' function attribute is set

[luciechoi]

Implement maximal reconvergence in SPIR-V codegen.

Addresses #136930

[llvmbot]

@llvm/pr-subscribers-hlsl
@llvm/pr-subscribers-clang

@llvm/pr-subscribers-clang-driver

Author: Lucie Choi (luciechoi)

Changes

Implement maximal reconvergence in SPIR-V codegen.

Addresses #136930

---

Full diff: https://github.com/llvm/llvm-project/pull/163682.diff

7 Files Affected:

• (modified) clang/lib/Driver/ToolChains/HLSL.cpp (+25-13)
• (modified) llvm/docs/SPIRVUsage.rst (+2)
• (modified) llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp (+20)
• (modified) llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp (+3-1)
• (modified) llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp (+3)
• (modified) llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td (+3-1)
• (added) llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_maximal_reconvergence/enable-maximal-reconvergence.ll (+12)

diff --git a/clang/lib/Driver/ToolChains/HLSL.cpp b/clang/lib/Driver/ToolChains/HLSL.cpp
index 5c8891fbfb98a..20a320ea233d4 100644
--- a/clang/lib/Driver/ToolChains/HLSL.cpp
+++ b/clang/lib/Driver/ToolChains/HLSL.cpp
@@ -191,23 +191,35 @@ void getSpirvExtOperand(StringRef SpvExtensionArg, raw_ostream &out) {
   // The extensions that are commented out are supported in DXC, but the SPIR-V
   // backend does not know about them yet.
   static const std::vector<StringRef> DxcSupportedExtensions = {
-      "SPV_KHR_16bit_storage", "SPV_KHR_device_group",
-      "SPV_KHR_fragment_shading_rate", "SPV_KHR_multiview",
-      "SPV_KHR_post_depth_coverage", "SPV_KHR_non_semantic_info",
-      "SPV_KHR_shader_draw_parameters", "SPV_KHR_ray_tracing",
-      "SPV_KHR_shader_clock", "SPV_EXT_demote_to_helper_invocation",
-      "SPV_EXT_descriptor_indexing", "SPV_EXT_fragment_fully_covered",
+      "SPV_KHR_16bit_storage",
+      "SPV_KHR_device_group",
+      "SPV_KHR_fragment_shading_rate",
+      "SPV_KHR_multiview",
+      "SPV_KHR_post_depth_coverage",
+      "SPV_KHR_non_semantic_info",
+      "SPV_KHR_shader_draw_parameters",
+      "SPV_KHR_ray_tracing",
+      "SPV_KHR_shader_clock",
+      "SPV_EXT_demote_to_helper_invocation",
+      "SPV_EXT_descriptor_indexing",
+      "SPV_EXT_fragment_fully_covered",
       "SPV_EXT_fragment_invocation_density",
-      "SPV_EXT_fragment_shader_interlock", "SPV_EXT_mesh_shader",
-      "SPV_EXT_shader_stencil_export", "SPV_EXT_shader_viewport_index_layer",
+      "SPV_EXT_fragment_shader_interlock",
+      "SPV_EXT_mesh_shader",
+      "SPV_EXT_shader_stencil_export",
+      "SPV_EXT_shader_viewport_index_layer",
       // "SPV_AMD_shader_early_and_late_fragment_tests",
-      "SPV_GOOGLE_hlsl_functionality1", "SPV_GOOGLE_user_type",
-      "SPV_KHR_ray_query", "SPV_EXT_shader_image_int64",
-      "SPV_KHR_fragment_shader_barycentric", "SPV_KHR_physical_storage_buffer",
+      "SPV_GOOGLE_hlsl_functionality1",
+      "SPV_GOOGLE_user_type",
+      "SPV_KHR_ray_query",
+      "SPV_EXT_shader_image_int64",
+      "SPV_KHR_fragment_shader_barycentric",
+      "SPV_KHR_physical_storage_buffer",
       "SPV_KHR_vulkan_memory_model",
       // "SPV_KHR_compute_shader_derivatives",
-      // "SPV_KHR_maximal_reconvergence",
-      "SPV_KHR_float_controls", "SPV_NV_shader_subgroup_partitioned",
+      "SPV_KHR_maximal_reconvergence",
+      "SPV_KHR_float_controls",
+      "SPV_NV_shader_subgroup_partitioned",
       // "SPV_KHR_quad_control"
   };
 
diff --git a/llvm/docs/SPIRVUsage.rst b/llvm/docs/SPIRVUsage.rst
index d2d66462b5df9..85eeabf10244a 100644
--- a/llvm/docs/SPIRVUsage.rst
+++ b/llvm/docs/SPIRVUsage.rst
@@ -235,6 +235,8 @@ Below is a list of supported SPIR-V extensions, sorted alphabetically by their e
      - Adds execution modes and decorations to control floating-point computations in both kernels and shaders. It can be used on whole modules and individual instructions.
    * - ``SPV_INTEL_predicated_io``
      - Adds predicated load and store instructions that conditionally read from or write to memory based on a boolean predicate.
+   * - ``SPV_KHR_maximal_reconvergence``
+     - Adds execution mode and capability to enable maximal reconvergence.
 
 SPIR-V representation in LLVM IR
 ================================
diff --git a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
index 56a6168775365..640b014646f36 100644
--- a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
@@ -78,6 +78,8 @@ class SPIRVAsmPrinter : public AsmPrinter {
   void outputExecutionModeFromNumthreadsAttribute(
       const MCRegister &Reg, const Attribute &Attr,
       SPIRV::ExecutionMode::ExecutionMode EM);
+  void outputExecutionModeFromEnableMaximalReconvergenceAttr(
+      const MCRegister &Reg, const SPIRVSubtarget &ST);
   void outputExecutionMode(const Module &M);
   void outputAnnotations(const Module &M);
   void outputModuleSections();
@@ -495,6 +497,20 @@ void SPIRVAsmPrinter::outputExecutionModeFromNumthreadsAttribute(
   outputMCInst(Inst);
 }
 
+void SPIRVAsmPrinter::outputExecutionModeFromEnableMaximalReconvergenceAttr(
+    const MCRegister &Reg, const SPIRVSubtarget &ST) {
+  assert(ST.canUseExtension(SPIRV::Extension::SPV_KHR_maximal_reconvergence) &&
+         "Function called when SPV_KHR_maximal_reconvergence is not enabled.");
+
+  MCInst Inst;
+  Inst.setOpcode(SPIRV::OpExecutionMode);
+  Inst.addOperand(MCOperand::createReg(Reg));
+  unsigned EM =
+      static_cast<unsigned>(SPIRV::ExecutionMode::MaximallyReconvergesKHR);
+  Inst.addOperand(MCOperand::createImm(EM));
+  outputMCInst(Inst);
+}
+
 void SPIRVAsmPrinter::outputExecutionMode(const Module &M) {
   NamedMDNode *Node = M.getNamedMetadata("spirv.ExecutionMode");
   if (Node) {
@@ -551,6 +567,10 @@ void SPIRVAsmPrinter::outputExecutionMode(const Module &M) {
     if (Attribute Attr = F.getFnAttribute("hlsl.numthreads"); Attr.isValid())
       outputExecutionModeFromNumthreadsAttribute(
           FReg, Attr, SPIRV::ExecutionMode::LocalSize);
+    if (Attribute Attr = F.getFnAttribute("enable-maximal-reconvergence");
+        Attr.getValueAsBool()) {
+      outputExecutionModeFromEnableMaximalReconvergenceAttr(FReg, *ST);
+    }
     if (MDNode *Node = F.getMetadata("work_group_size_hint"))
       outputExecutionModeFromMDNode(FReg, Node,
                                     SPIRV::ExecutionMode::LocalSizeHint, 3, 1);
diff --git a/llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp b/llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp
index 5f3ed862af893..96f5dee21bc2a 100644
--- a/llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp
@@ -153,7 +153,9 @@ static const std::map<std::string, SPIRV::Extension::Extension, std::less<>>
          SPIRV::Extension::Extension::
              SPV_EXT_relaxed_printf_string_address_space},
         {"SPV_INTEL_predicated_io",
-         SPIRV::Extension::Extension::SPV_INTEL_predicated_io}};
+         SPIRV::Extension::Extension::SPV_INTEL_predicated_io},
+        {"SPV_KHR_maximal_reconvergence",
+         SPIRV::Extension::Extension::SPV_KHR_maximal_reconvergence}};
 
 bool SPIRVExtensionsParser::parse(cl::Option &O, StringRef ArgName,
                                   StringRef ArgValue,
diff --git a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
index 5144fb14fa6a6..6760f5a2829d3 100644
--- a/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp
@@ -2155,6 +2155,9 @@ static void collectReqs(const Module &M, SPIRV::ModuleAnalysisInfo &MAI,
           SPIRV::OperandCategory::ExecutionModeOperand,
           SPIRV::ExecutionMode::LocalSize, ST);
     }
+    if (F.getFnAttribute("enable-maximal-reconvergence").getValueAsBool()) {
+      MAI.Reqs.addExtension(SPIRV::Extension::SPV_KHR_maximal_reconvergence);
+    }
     if (F.getMetadata("work_group_size_hint"))
       MAI.Reqs.getAndAddRequirements(
           SPIRV::OperandCategory::ExecutionModeOperand,
diff --git a/llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td b/llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td
index 26256429537e9..7d08b29a51a6e 100644
--- a/llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td
+++ b/llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td
@@ -386,6 +386,7 @@ defm SPV_KHR_float_controls2 : ExtensionOperand<124, [EnvVulkan, EnvOpenCL]>;
 defm SPV_INTEL_tensor_float32_conversion : ExtensionOperand<125, [EnvOpenCL]>;
 defm SPV_KHR_bfloat16 : ExtensionOperand<126, [EnvVulkan, EnvOpenCL]>;
 defm SPV_INTEL_predicated_io : ExtensionOperand<127, [EnvOpenCL]>;
+defm SPV_KHR_maximal_reconvergence : ExtensionOperand<128, [EnvVulkan]>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
@@ -698,7 +699,7 @@ defm IntersectionNV: ExecutionModelOperand<5314, [RayTracingNV]>;
 defm AnyHitNV: ExecutionModelOperand<5315, [RayTracingNV]>;
 defm ClosestHitNV: ExecutionModelOperand<5316, [RayTracingNV]>;
 defm MissNV: ExecutionModelOperand<5317, [RayTracingNV]>;
-defm CallableNV: ExecutionModelOperand<5318, [RayTracingNV]>;
+defm CallableNV : ExecutionModelOperand<5318, [RayTracingNV]>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define MemoryModel enum values and at the same time
@@ -805,6 +806,7 @@ defm RoundingModeRTNINTEL : ExecutionModeOperand<5621, [RoundToInfinityINTEL]>;
 defm FloatingPointModeALTINTEL : ExecutionModeOperand<5622, [FloatingPointModeINTEL]>;
 defm FloatingPointModeIEEEINTEL : ExecutionModeOperand<5623, [FloatingPointModeINTEL]>;
 defm FPFastMathDefault : ExecutionModeOperand<6028, [FloatControls2]>;
+defm MaximallyReconvergesKHR : ExecutionModeOperand<6023, [Shader]>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define StorageClass enum values and at the same time
diff --git a/llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_maximal_reconvergence/enable-maximal-reconvergence.ll b/llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_maximal_reconvergence/enable-maximal-reconvergence.ll
new file mode 100644
index 0000000000000..662190e039957
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_maximal_reconvergence/enable-maximal-reconvergence.ll
@@ -0,0 +1,12 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv1.6-unknown-vulkan1.3-compute --spirv-ext=+SPV_KHR_maximal_reconvergence %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-unknown-vulkan1.3-compute --spirv-ext=+SPV_KHR_maximal_reconvergence %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpCapability Shader
+; CHECK: OpExtension "SPV_KHR_maximal_reconvergence"
+; CHECK: OpExecutionMode {{.*}} MaximallyReconvergesKHR
+define void @main() local_unnamed_addr #0 {
+entry:
+  ret void
+}
+
+attributes #0 = { mustprogress nofree noinline norecurse nosync nounwind willreturn memory(none) "enable-maximal-reconvergence"="true" "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-infs-fp-math"="true" "no-nans-fp-math"="true" "no-signed-zeros-fp-math"="true" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }

[luciechoi]

cc: @s-perron , @Keenuts

[Keenuts]

Thanks for this! Code looks good to me, just a request regarding testing 😊

[s-perron]

This looks really good. Add the test that Nathan suggested, and I think this is ready to merge.