Add raypayload decoration to ray payload structs (#6164)

* Add raypayload decoration to ray payload structs

Closes https://github.com/shader-slang/slang/issues/6104

* Disable PAQs when compiling with DXC

See https://github.com/shader-slang/slang/issues/3448

10 files changed, 131 insertions, 27 deletions

diff --git a/source/compiler-core/slang-dxc-compiler.cpp b/source/compiler-core/slang-dxc-compiler.cpp
index cfc4771e6..2539826ea 100644
--- a/source/compiler-core/slang-dxc-compiler.cpp
+++ b/source/compiler-core/slang-dxc-compiler.cpp
@@ -479,6 +479,13 @@ SlangResult DXCDownstreamCompiler::compile(const CompileOptions& inOptions, IArt
         args.add(compilerSpecific[i]);
     }
 
+    // This can be re-enabled when we add PAQs: https://github.com/shader-slang/slang/issues/3448
+    const bool enablePAQs = false;
+    if (!enablePAQs)
+        args.add(L"-disable-payload-qualifiers");
+    else
+        args.add(L"-enable-payload-qualifiers");
+
     // TODO: deal with
     bool treatWarningsAsErrors = false;
     if (treatWarningsAsErrors)
diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 884621960..a10e747c0 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -16464,6 +16464,13 @@ __generic<T>
 __intrinsic_op($(kIROp_ForceVarIntoStructTemporarily))
 Ref<T> __forceVarIntoStructTemporarily(inout T maybeStruct);
 
+// Some functions require a struct type which is decorated with a [raypayload]
+// attribute. This will do the same as __forceVarIntoStructTemporarily and also
+// ensure that the struct type in question is decorated appropriately.
+__generic<T>
+__intrinsic_op($(kIROp_ForceVarIntoRayPayloadStructTemporarily))
+Ref<T> __forceVarIntoRayPayloadStructTemporarily(inout T maybeStruct);
+
 __generic<payload_t>
 [require(hlsl, raytracing)]
 void __traceRayHLSL(
@@ -16548,7 +16555,7 @@ void TraceRay(
             MultiplierForGeometryContributionToHitGroupIndex,
             MissShaderIndex,
             Ray,
-            __forceVarIntoStructTemporarily(Payload));
+            __forceVarIntoRayPayloadStructTemporarily(Payload));
         return;
     case cuda: __intrinsic_asm "traceOptiXRay";
     case glsl:
@@ -16686,7 +16693,7 @@ void TraceMotionRay(
             MissShaderIndex,
             Ray,
             CurrentTime,
-            __forceVarIntoStructTemporarily(Payload));
+            __forceVarIntoRayPayloadStructTemporarily(Payload));
         return;
     case glsl:
     {
@@ -18830,7 +18837,7 @@ struct HitObject
                     MultiplierForGeometryContributionToHitGroupIndex, 
                     MissShaderIndex, 
                     Ray, 
-                    __forceVarIntoStructTemporarily(Payload),
+                    __forceVarIntoRayPayloadStructTemporarily(Payload),
                     hitObj);
                 return hitObj;
             }
@@ -18923,7 +18930,7 @@ struct HitObject
                 MissShaderIndex,
                 Ray,
                 CurrentTime,
-                __forceVarIntoStructTemporarily(Payload));
+                __forceVarIntoRayPayloadStructTemporarily(Payload));
         case glsl:
             {
                 [__vulkanRayPayload]
@@ -19441,7 +19448,7 @@ struct HitObject
             __InvokeHLSL(
                 AccelerationStructure,
                 HitOrMiss,
-                __forceVarIntoStructTemporarily(Payload));
+                __forceVarIntoRayPayloadStructTemporarily(Payload));
         case glsl:
             {
                 [__vulkanRayPayload]
diff --git a/source/slang/slang-emit-hlsl.cpp b/source/slang/slang-emit-hlsl.cpp
index ff4514d69..1cb9b769f 100644
--- a/source/slang/slang-emit-hlsl.cpp
+++ b/source/slang/slang-emit-hlsl.cpp
@@ -1669,6 +1669,15 @@ void HLSLSourceEmitter::emitPostKeywordTypeAttributesImpl(IRInst* inst)
     {
         m_writer->emit("[payload] ");
     }
+    // This can be re-enabled when we add PAQs: https://github.com/shader-slang/slang/issues/3448
+    const bool enablePAQs = false;
+    if (enablePAQs)
+    {
+        if (const auto payloadDecoration = inst->findDecoration<IRRayPayloadDecoration>())
+        {
+            m_writer->emit("[raypayload] ");
+        }
+    }
 }
 
 void HLSLSourceEmitter::_emitPrefixTypeAttr(IRAttr* attr)
diff --git a/source/slang/slang-ir-hlsl-legalize.cpp b/source/slang/slang-ir-hlsl-legalize.cpp
index 0670babdc..ec2419985 100644
--- a/source/slang/slang-ir-hlsl-legalize.cpp
+++ b/source/slang/slang-ir-hlsl-legalize.cpp
@@ -29,14 +29,20 @@ void searchChildrenForForceVarIntoStructTemporarily(IRModule* module, IRInst* in
                 for (UInt i = 0; i < call->getArgCount(); i++)
                 {
                     auto arg = call->getArg(i);
-                    if (arg->getOp() != kIROp_ForceVarIntoStructTemporarily)
+                    const bool isForcedStruct = arg->getOp() == kIROp_ForceVarIntoStructTemporarily;
+                    const bool isForcedRayPayloadStruct =
+                        arg->getOp() == kIROp_ForceVarIntoRayPayloadStructTemporarily;
+                    if (!(isForcedStruct || isForcedRayPayloadStruct))
                         continue;
                     auto forceStructArg = arg->getOperand(0);
                     auto forceStructBaseType =
                         as<IRType>(forceStructArg->getDataType()->getOperand(0));
+                    IRBuilder builder(call);
                     if (forceStructBaseType->getOp() == kIROp_StructType)
                     {
                         call->setArg(i, arg->getOperand(0));
+                        if (isForcedRayPayloadStruct)
+                            builder.addRayPayloadDecoration(forceStructBaseType);
                         continue;
                     }
 
@@ -47,14 +53,19 @@ void searchChildrenForForceVarIntoStructTemporarily(IRModule* module, IRInst* in
                     // `__forceVarIntoStructTemporarily` is a parameter to a side effect type
                     // (`ref`, `out`, `inout`) we copy the struct back into our original non-struct
                     // parameter.
-                    IRBuilder builder(call);
+
+                    const auto typeNameHint = isForcedRayPayloadStruct
+                                                  ? "RayPayload_t"
+                                                  : "ForceVarIntoStructTemporarily_t";
+                    const auto varNameHint =
+                        isForcedRayPayloadStruct ? "rayPayload" : "forceVarIntoStructTemporarily";
 
                     builder.setInsertBefore(call->getCallee());
                     auto structType = builder.createStructType();
                     StringBuilder structName;
-                    builder.addNameHintDecoration(
-                        structType,
-                        UnownedStringSlice("ForceVarIntoStructTemporarily_t"));
+                    builder.addNameHintDecoration(structType, UnownedStringSlice(typeNameHint));
+                    if (isForcedRayPayloadStruct)
+                        builder.addRayPayloadDecoration(structType);
 
                     auto elementBufferKey = builder.createStructKey();
                     builder.addNameHintDecoration(elementBufferKey, UnownedStringSlice("data"));
@@ -65,9 +76,7 @@ void searchChildrenForForceVarIntoStructTemporarily(IRModule* module, IRInst* in
 
                     builder.setInsertBefore(call);
                     auto structVar = builder.emitVar(structType);
-                    builder.addNameHintDecoration(
-                        structVar,
-                        UnownedStringSlice("forceVarIntoStructTemporarily"));
+                    builder.addNameHintDecoration(structVar, UnownedStringSlice(varNameHint));
                     builder.emitStore(
                         builder.emitFieldAddress(
                             builder.getPtrType(_dataField->getFieldType()),
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 4de7457a3..55880eab5 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -759,6 +759,9 @@ INST(GetPerVertexInputArray, GetPerVertexInputArray, 1, HOISTABLE)
 INST(ResolveVaryingInputRef, ResolveVaryingInputRef, 1, HOISTABLE)
 
 INST(ForceVarIntoStructTemporarily, ForceVarIntoStructTemporarily, 1, 0)
+INST(ForceVarIntoRayPayloadStructTemporarily, ForceVarIntoRayPayloadStructTemporarily, 1, 0)
+INST_RANGE(ForceVarIntoStructTemporarily, ForceVarIntoStructTemporarily, ForceVarIntoRayPayloadStructTemporarily)
+
 INST(MetalAtomicCast, MetalAtomicCast, 1, 0)
 
 INST(IsTextureAccess, IsTextureAccess, 1, 0)
@@ -992,6 +995,7 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
     INST(GLSLLocationDecoration, glslLocation, 1, 0)
     INST(GLSLOffsetDecoration, glslOffset, 1, 0)
     INST(PayloadDecoration, payload, 0, 0)
+    INST(RayPayloadDecoration, raypayload, 0, 0)
 
     /* Mesh Shader outputs */
         INST(VerticesDecoration, vertices, 1, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 2bc5bca1e..9c3892c0e 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -1605,6 +1605,11 @@ struct IRPayloadDecoration : public IRDecoration
     IR_LEAF_ISA(PayloadDecoration)
 };
 
+struct IRRayPayloadDecoration : public IRDecoration
+{
+    IR_LEAF_ISA(RayPayloadDecoration)
+};
+
 // Mesh shader decorations
 
 struct IRMeshOutputDecoration : public IRDecoration
@@ -5289,6 +5294,8 @@ public:
     {
         addDecoration(inst, kIROp_EntryPointParamDecoration, entryPointFunc);
     }
+
+    void addRayPayloadDecoration(IRType* inst) { addDecoration(inst, kIROp_RayPayloadDecoration); }
 };
 
 // Helper to establish the source location that will be used
diff --git a/tests/expected-failure-github.txt b/tests/expected-failure-github.txt
index 6bf4f041d..60d632785 100644
--- a/tests/expected-failure-github.txt
+++ b/tests/expected-failure-github.txt
@@ -13,6 +13,4 @@ tests/bugs/buffer-swizzle-store.slang.3 syn (wgpu)
 tests/compute/interface-shader-param-in-struct.slang.4 syn (wgpu)
 tests/compute/interface-shader-param.slang.5 syn (wgpu)
 tests/language-feature/shader-params/interface-shader-param-ordinary.slang.4 syn (wgpu)
-gfx-unit-test-tool/RayTracingTestAD3D12.internal
-gfx-unit-test-tool/RayTracingTestBD3D12.internal
 gfx-unit-test-tool/precompiledTargetModule2Vulkan.internal
diff --git a/tests/hlsl/raypayload-attribute-no-struct.slang b/tests/hlsl/raypayload-attribute-no-struct.slang
new file mode 100644
index 000000000..c7ad94593
--- /dev/null
+++ b/tests/hlsl/raypayload-attribute-no-struct.slang
@@ -0,0 +1,29 @@
+//enable when https://github.com/shader-slang/slang/issues/3448 is implemented
+//DISABLE_TEST:SIMPLE(filecheck=CHECK): -target hlsl -stage raygeneration -entry rayGenShaderA
+
+// CHECK: struct [raypayload]
+
+uniform RWTexture2D resultTexture;
+uniform RaytracingAccelerationStructure sceneBVH;
+
+[shader("raygeneration")]
+void rayGenShaderA()
+{
+    int2 threadIdx = DispatchRaysIndex().xy;
+
+    float3 rayDir = float3(0, 0, 1);
+    float3 rayOrigin = 0;
+    rayOrigin.x = (threadIdx.x * 2) - 1;
+    rayOrigin.y = (threadIdx.y * 2) - 1;
+
+    // Trace the ray.
+    RayDesc ray;
+    ray.Origin = rayOrigin;
+    ray.Direction = rayDir;
+    ray.TMin = 0.001;
+    ray.TMax = 10000.0;
+    float4 payload = float4(0, 0, 0, 0);
+    TraceRay(sceneBVH, RAY_FLAG_NONE, ~0, 0, 0, 0, ray, payload);
+
+    resultTexture[threadIdx.xy] = payload;
+}
diff --git a/tests/hlsl/raypayload-attribute.slang b/tests/hlsl/raypayload-attribute.slang
new file mode 100644
index 000000000..b981589ac
--- /dev/null
+++ b/tests/hlsl/raypayload-attribute.slang
@@ -0,0 +1,34 @@
+//enable when https://github.com/shader-slang/slang/issues/3448 is implemented
+//DISABLE_TEST:SIMPLE(filecheck=CHECK): -target hlsl -stage raygeneration -entry rayGenShaderA
+
+// CHECK: struct [raypayload]
+
+struct RayPayload
+{
+    float4 color;
+};
+
+uniform RWTexture2D resultTexture;
+uniform RaytracingAccelerationStructure sceneBVH;
+
+[shader("raygeneration")]
+void rayGenShaderA()
+{
+    int2 threadIdx = DispatchRaysIndex().xy;
+
+    float3 rayDir = float3(0, 0, 1);
+    float3 rayOrigin = 0;
+    rayOrigin.x = (threadIdx.x * 2) - 1;
+    rayOrigin.y = (threadIdx.y * 2) - 1;
+
+    // Trace the ray.
+    RayDesc ray;
+    ray.Origin = rayOrigin;
+    ray.Direction = rayDir;
+    ray.TMin = 0.001;
+    ray.TMax = 10000.0;
+    RayPayload payload = { float4(0, 0, 0, 0) };
+    TraceRay(sceneBVH, RAY_FLAG_NONE, ~0, 0, 0, 0, ray, payload);
+
+    resultTexture[threadIdx.xy] = payload.color;
+}
diff --git a/tests/vkray/raygen-trace-ray-param-non-struct.slang b/tests/vkray/raygen-trace-ray-param-non-struct.slang
index c4451d941..b0a129761 100644
--- a/tests/vkray/raygen-trace-ray-param-non-struct.slang
+++ b/tests/vkray/raygen-trace-ray-param-non-struct.slang
@@ -22,13 +22,13 @@ void main()
     ray.Direction = float3(0,0,1);
     ray.TMax = 100.0f;
 
-    // CHECK: ForceVarIntoStructTemporarily_t{{_[0-9]}} forceVarIntoStructTemporarily{{_[0-9]}};
+    // CHECK: RayPayload_t{{_[0-9]}} rayPayload{{_[0-9]}};
     float someInData1 = 5.0f;
     addComplexity1(someInData1);
 
-    // CHECK: forceVarIntoStructTemporarily{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
+    // CHECK: rayPayload{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
     // CHECK: TraceRay(
-    // CHECK: {{.*}} = forceVarIntoStructTemporarily{{.*}}.data{{.*}};
+    // CHECK: {{.*}} = rayPayload{{.*}}.data{{.*}};
     TraceRay(as,
              1,
              0xff,
@@ -39,9 +39,9 @@ void main()
              someInData1);
     outputBuffer1[0] = outputBuffer1[0]+someInData1;
 
-    // CHECK: forceVarIntoStructTemporarily{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
+    // CHECK: rayPayload{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
     // CHECK: TraceMotionRay(
-    // CHECK: {{.*}} = forceVarIntoStructTemporarily{{.*}}.data{{.*}};
+    // CHECK: {{.*}} = rayPayload{{.*}}.data{{.*}};
     TraceMotionRay(as,
             1,
             0xff,
@@ -53,9 +53,9 @@ void main()
             someInData1);
     outputBuffer1[0] = outputBuffer1[0]+someInData1;
 
-    // CHECK: forceVarIntoStructTemporarily{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
+    // CHECK: rayPayload{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
     // CHECK: NvTraceRayHitObject(
-    // CHECK: {{.*}} = forceVarIntoStructTemporarily{{.*}}.data{{.*}};
+    // CHECK: {{.*}} = rayPayload{{.*}}.data{{.*}};
     HitObject::TraceRay(as,
              1,
              0xff,
@@ -66,9 +66,9 @@ void main()
              someInData1);
     outputBuffer1[0] = outputBuffer1[0]+someInData1;
 
-    // CHECK: forceVarIntoStructTemporarily{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
+    // CHECK: rayPayload{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
     // CHECK: TraceMotionRay(
-    // CHECK: {{.*}} = forceVarIntoStructTemporarily{{.*}}.data{{.*}};
+    // CHECK: {{.*}} = rayPayload{{.*}}.data{{.*}};
     HitObject::TraceMotionRay(as,
             1,
             0xff,
@@ -80,9 +80,9 @@ void main()
             someInData1);
     outputBuffer1[0] = outputBuffer1[0]+someInData1;
 
-    // CHECK: forceVarIntoStructTemporarily{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
+    // CHECK: rayPayload{{_[0-9]}}.data{{_[0-9]}} = {{.*}}
     // CHECK: NvInvokeHitObject(
-    // CHECK: {{.*}} = forceVarIntoStructTemporarily{{.*}}.data{{.*}};
+    // CHECK: {{.*}} = rayPayload{{.*}}.data{{.*}};
     HitObject hitObject_HitOrMiss;
     HitObject::Invoke(
         as,
@@ -91,4 +91,4 @@ void main()
     outputBuffer1[0] = outputBuffer1[0]+someInData1;
 
     addComplexity2(someInData1);
-}
\ No newline at end of file
+}