spirv: correctly choose ray_tracing/ray_query caps (microsoft#5662)

Both the ray_query and ray_tracing extensions brings new SPIRV types.
But some hardware only support one type.

This decision cannot be taken by the trimming pass as both choices could
be valid, instead, we should check what extension the user needed.

What can be removed however is the logic to emit the capability. This
commit also does that.

Fixes microsoft#4385

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -23,15 +23,18 @@ void CapabilityVisitor::addExtension(Extension ext, llvm::StringRef target,
     spvBuilder.requireExtension(featureManager.getExtensionName(ext), loc);
 }
 
-void CapabilityVisitor::addExtensionAndCapabilitiesIfEnabled(
+bool CapabilityVisitor::addExtensionAndCapabilitiesIfEnabled(
     Extension ext, llvm::ArrayRef<spv::Capability> capabilities) {
-  if (featureManager.isExtensionEnabled(ext)) {
-    addExtension(ext, "", {});
+  if (!featureManager.isExtensionEnabled(ext)) {
+    return false;
+  }
 
-    for (auto cap : capabilities) {
-      addCapability(cap);
-    }
+  addExtension(ext, "", {});
+
+  for (auto cap : capabilities) {
+    addCapability(cap);
   }
+  return true;
 }
 
 void CapabilityVisitor::addCapability(spv::Capability cap, SourceLocation loc) {
@@ -225,18 +228,6 @@ void CapabilityVisitor::addCapabilityForType(const SpirvType *type,
     for (auto field : structType->getFields())
       addCapabilityForType(field.type, loc, sc);
   }
-  // AccelerationStructureTypeNV and RayQueryTypeKHR type
-  // Note: Because AccelerationStructureType can be provided by both
-  // SPV_KHR_ray_query and SPV_{NV,KHR}_ray_tracing extensions, this logic will
-  // result in SPV_KHR_ray_query being unnecessarily required in some cases. If
-  // this is an issue in future (more devices are identified that support
-  // ray_tracing but not ray_query), then we should consider addressing this
-  // interaction with a spirv-opt pass instead.
-  else if (isa<AccelerationStructureTypeNV>(type) ||
-           isa<RayQueryTypeKHR>(type)) {
-    addCapability(spv::Capability::RayQueryKHR);
-    addExtension(Extension::KHR_ray_query, "SPV_KHR_ray_query", {});
-  }
 }
 
 bool CapabilityVisitor::visit(SpirvDecoration *decor) {
@@ -888,6 +879,18 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
           spv::Capability::FragmentShaderPixelInterlockEXT,
           spv::Capability::FragmentShaderShadingRateInterlockEXT,
       });
+
+  // AccelerationStructureType or RayQueryType can be provided by both
+  // ray_tracing and ray_query extension. By default, we select ray_query to
+  // provide it. This is an arbitrary decision. If the user wants avoid one
+  // extension (lack of support by ex), if can be done by providing the list
+  // of enabled extensions.
+  if (!addExtensionAndCapabilitiesIfEnabled(Extension::KHR_ray_query,
+                                            {spv::Capability::RayQueryKHR})) {
+    addExtensionAndCapabilitiesIfEnabled(Extension::KHR_ray_tracing,
+                                         {spv::Capability::RayTracingKHR});
+  }
+
   return true;
 }
 

tools/clang/lib/SPIRV/CapabilityVisitor.h

@@ -66,7 +66,8 @@ class CapabilityVisitor : public Visitor {
 
   /// Checks that the given extension is enabled based on command line arguments
   /// before calling addExtension and addCapability.
-  void addExtensionAndCapabilitiesIfEnabled(
+  /// Returns `true` if the extension was enabled, `false` otherwise.
+  bool addExtensionAndCapabilitiesIfEnabled(
       Extension ext, llvm::ArrayRef<spv::Capability> capabilities);
 
   /// Checks that the given capability is a valid capability. And if so,

tools/clang/test/CodeGenSPIRV_Lit/raytracing.acceleration-structure.ray_query.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -E main -T vs_6_5 -fspv-target-env=vulkan1.2 -fspv-extension=SPV_KHR_ray_query %s -spirv | FileCheck %s
+
+// CHECK-NOT: OpCapability RayTracingKHR
+// CHECK: OpCapability RayQueryKHR
+// CHECK: OpCapability Shader
+// CHECK-NEXT: OpExtension "SPV_KHR_ray_query"
+// CHECK-NOT: OpExtension "SPV_KHR_ray_tracing"
+
+RaytracingAccelerationStructure test_bvh;
+
+float main(RayDesc rayDesc : RAYDESC) : OUT {
+  RayQuery<RAY_FLAG_FORCE_OPAQUE|RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES> rayQuery1;
+  rayQuery1.TraceRayInline(test_bvh, 0, 1, rayDesc);
+  return 0;
+}

tools/clang/test/CodeGenSPIRV_Lit/raytracing.acceleration-structure.ray_tracing.hlsl

@@ -0,0 +1,24 @@
+// RUN: %dxc -E main -T lib_6_6 -fspv-target-env=vulkan1.2 -fspv-extension=SPV_KHR_ray_tracing  %s -spirv | FileCheck %s
+
+// CHECK-NOT: OpCapability RayQueryKHR
+// CHECK: OpCapability RayTracingKHR
+// CHECK-NEXT: OpExtension "SPV_KHR_ray_tracing"
+// CHECK-NOT: OpExtension "SPV_KHR_ray_query"
+
+RaytracingAccelerationStructure rs;
+
+struct Payload { float4 color; };
+struct Attribute { float2 bary; };
+
+[shader("closesthit")]
+void main(inout Payload MyPayload, in Attribute MyAttr) {
+
+  Payload myPayload = { float4(0.0f,0.0f,0.0f,0.0f) };
+  RayDesc rayDesc;
+  rayDesc.Origin = float3(0.0f, 0.0f, 0.0f);
+  rayDesc.Direction = float3(0.0f, 0.0f, -1.0f);
+  rayDesc.TMin = 0.0f;
+  rayDesc.TMax = 1000.0f;
+
+  TraceRay(rs, 0x0, 0xff, 0, 1, 0, rayDesc, myPayload);
+}