[SPIR-V] Implement WaveMatch intrinsic function (#6546)

Adds support for the `WaveMatch()` intrinsic function from Shader Model
6.5 using the `OpGroupNonUniformPartitionNV` instruction from the
`SPV_NV_shader_subgroup_partitioned` extension.

SPIRV-Tools bumped to include:
KhronosGroup/SPIRV-Tools#5648

Fixes #6545

Author: sudonatalie

docs/SPIR-V.rst

@@ -3826,8 +3826,8 @@ RayQuery Mapping to SPIR-V
 |``.WorldRayOrigin`                                 | ``OpRayQueryGetWorldRayOriginKHR``                                      |
 +---------------------------------------------------+-------------------------------------------------------------------------+
 
-Shader Model 6.0 Wave Intrinsics
-================================
+Shader Model 6.0+ Wave Intrinsics
+=================================
 
 
 Note that Wave intrinsics requires SPIR-V 1.3, which is supported by Vulkan 1.1.
@@ -3865,6 +3865,7 @@ Quad          ``QuadReadAcrossX()``        ``OpGroupNonUniformQuadSwap``
 Quad          ``QuadReadAcrossY()``        ``OpGroupNonUniformQuadSwap``
 Quad          ``QuadReadAcrossDiagonal()`` ``OpGroupNonUniformQuadSwap``
 Quad          ``QuadReadLaneAt()``         ``OpGroupNonUniformQuadBroadcast``
+N/A           ``WaveMatch()``              ``OpGroupNonUniformPartitionNV``
 ============= ============================ =================================== ======================
 
 The Implicit ``vk`` Namespace

tools/clang/include/clang/SPIRV/FeatureManager.h

@@ -61,6 +61,7 @@ enum class Extension {
   KHR_fragment_shader_barycentric,
   KHR_maximal_reconvergence,
   KHR_float_controls,
+  NV_shader_subgroup_partitioned,
   Unknown,
 };
 

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -898,6 +898,10 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
       Extension::KHR_vulkan_memory_model,
       {spv::Capability::VulkanMemoryModelDeviceScope});
 
+  addExtensionAndCapabilitiesIfEnabled(
+      Extension::NV_shader_subgroup_partitioned,
+      {spv::Capability::GroupNonUniformPartitionedNV});
+
   return true;
 }
 

tools/clang/lib/SPIRV/FeatureManager.cpp

@@ -197,6 +197,8 @@ Extension FeatureManager::getExtensionSymbol(llvm::StringRef name) {
       .Case("SPV_KHR_maximal_reconvergence",
             Extension::KHR_maximal_reconvergence)
       .Case("SPV_KHR_float_controls", Extension::KHR_float_controls)
+      .Case("SPV_NV_shader_subgroup_partitioned",
+            Extension::NV_shader_subgroup_partitioned)
       .Default(Extension::Unknown);
 }
 
@@ -264,6 +266,8 @@ const char *FeatureManager::getExtensionName(Extension symbol) {
     return "SPV_KHR_maximal_reconvergence";
   case Extension::KHR_float_controls:
     return "SPV_KHR_float_controls";
+  case Extension::NV_shader_subgroup_partitioned:
+    return "SPV_NV_shader_subgroup_partitioned";
   default:
     break;
   }

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -8716,6 +8716,9 @@ SpirvEmitter::processIntrinsicCallExpr(const CallExpr *callExpr) {
   case hlsl::IntrinsicOp::IOP_WaveReadLaneFirst:
     retVal = processWaveBroadcast(callExpr);
     break;
+  case hlsl::IntrinsicOp::IOP_WaveMatch:
+    retVal = processWaveMatch(callExpr);
+    break;
   case hlsl::IntrinsicOp::IOP_QuadReadAcrossX:
   case hlsl::IntrinsicOp::IOP_QuadReadAcrossY:
   case hlsl::IntrinsicOp::IOP_QuadReadAcrossDiagonal:
@@ -9710,6 +9713,17 @@ SpirvEmitter::processWaveActiveAllEqualMatrix(SpirvInstruction *arg,
   return spvBuilder.createCompositeConstruct(booleanMatrixType, rows, srcLoc);
 }
 
+SpirvInstruction *SpirvEmitter::processWaveMatch(const CallExpr *callExpr) {
+  assert(callExpr->getNumArgs() == 1);
+  const auto loc = callExpr->getExprLoc();
+
+  // The SPV_NV_shader_subgroup_partitioned extension requires SPIR-V 1.3.
+  featureManager.requestTargetEnv(SPV_ENV_VULKAN_1_1, "Wave Operation", loc);
+
+  SpirvInstruction *arg = doExpr(callExpr->getArg(0));
+  return spvBuilder.createUnaryOp(spv::Op::OpGroupNonUniformPartitionNV,
+                                  callExpr->getType(), arg, loc);
+}
 SpirvInstruction *SpirvEmitter::processIntrinsicModf(const CallExpr *callExpr) {
   // Signature is: ret modf(x, ip)
   // [in]    x: the input floating-point value.

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -655,6 +655,9 @@ class SpirvEmitter : public ASTConsumer {
   processWaveActiveAllEqualMatrix(SpirvInstruction *arg, QualType,
                                   clang::SourceLocation srcLoc);
 
+  /// Processes SM6.5 WaveMatch function.
+  SpirvInstruction *processWaveMatch(const CallExpr *);
+
   /// Processes the NonUniformResourceIndex intrinsic function.
   SpirvInstruction *processIntrinsicNonUniformResourceIndex(const CallExpr *);
 

tools/clang/test/CodeGenSPIRV/sm6_5.wave-match.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -E main -T ps_6_5 -spirv -O0 -fspv-target-env=vulkan1.1 %s | FileCheck %s
+// RUN: not %dxc -E main -T ps_6_5 -spirv -O0 %s 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+
+// CHECK-ERROR: error: Vulkan 1.1 is required for Wave Operation but not permitted to use
+
+// CHECK: OpCapability GroupNonUniformPartitionedNV
+// CHECK: OpExtension "SPV_NV_shader_subgroup_partitioned"
+
+uint4 main(uint4 input : ATTR0) : SV_Target {
+// CHECK: [[input:%[0-9]+]] = OpLoad %v4uint %input
+// CHECK:       {{%[0-9]+}} = OpGroupNonUniformPartitionNV %v4uint [[input]]
+    uint4 res = WaveMatch(input);
+    return res;
+}