OptiX ray payload read/write support in raytracing pipeline shaders (#1853)

* OptiX ray payload can now be read from and written to using the two payload register pointer method

* changing op to more descriptive name

* fixup: comment change to re-trigger CI

Co-authored-by: T. Foley <tfoleyNV@users.noreply.github.com>

5 files changed, 104 insertions, 9 deletions

diff --git a/prelude/slang-cuda-prelude.h b/prelude/slang-cuda-prelude.h
index 4df60e965..a6c4f70dc 100644
--- a/prelude/slang-cuda-prelude.h
+++ b/prelude/slang-cuda-prelude.h
@@ -1612,3 +1612,29 @@ found via reflection or defined such that it matches the shader appropriately.
 */
 struct UniformEntryPointParams;
 struct UniformState;
+
+// ---------------------- OptiX Ray Payload --------------------------------------
+#ifdef SLANG_CUDA_ENABLE_OPTIX
+static __forceinline__ __device__
+void *unpackOptiXRayPayloadPointer(uint32_t i0, uint32_t i1)
+{
+	const uint64_t uptr = static_cast<uint64_t>(i0) << 32 | i1;
+	void*           ptr = reinterpret_cast<void*>(uptr);
+	return ptr;
+}
+
+static __forceinline__ __device__
+void  packOptiXRayPayloadPointer(void* ptr, uint32_t& i0, uint32_t& i1)
+{
+	const uint64_t uptr = reinterpret_cast<uint64_t>(ptr);
+	i0 = uptr >> 32;
+	i1 = uptr & 0x00000000ffffffff;
+}
+
+static __forceinline__ __device__ void *getOptiXRayPayloadPtr()
+{
+	const uint32_t u0 = optixGetPayload_0();
+	const uint32_t u1 = optixGetPayload_1();
+	return unpackOptiXRayPayloadPointer(u0, u1);
+}
+#endif
diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index 09ea7ef9e..4c430cfa6 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -781,6 +781,13 @@ bool CUDASourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
             m_writer->emit(")");
             return true;
         }
+        case kIROp_GetOptiXRayPayloadPtr:
+        {
+            m_writer->emit("(");
+            emitType(inst->getDataType());
+            m_writer->emit(")getOptiXRayPayloadPtr()");
+            return true;
+        }
         default: break;
     }
 
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index f73dbd006..17ae652ae 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -497,6 +497,10 @@ INST(GroupMemoryBarrierWithGroupSync, GroupMemoryBarrierWithGroupSync, 0, 0)
 // GPU_FOREACH loop of the form 
 INST(GpuForeach, gpuForeach, 3, 0)
 
+// Wrapper for OptiX intrinsics used to load and store ray payload data using
+// a pointer represented by two payload registers.
+INST(GetOptiXRayPayloadPtr, getOptiXRayPayloadPtr, 0, 0)
+
 /* Decoration */
 
 INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
diff --git a/source/slang/slang-ir-legalize-varying-params.cpp b/source/slang/slang-ir-legalize-varying-params.cpp
index c802513e8..f500b100f 100644
--- a/source/slang/slang-ir-legalize-varying-params.cpp
+++ b/source/slang/slang-ir-legalize-varying-params.cpp
@@ -895,11 +895,7 @@ protected:
 
     virtual LegalizedVaryingVal createLegalUserVaryingValImpl(VaryingParamInfo const& info)
     {
-        SLANG_UNUSED(info);
-
-        m_sink->diagnose(m_param, Diagnostics::unimplemented, "this target doesn't support user-defined varying parameters");
-
-        return LegalizedVaryingVal();
+        return diagnoseUnsupportedUserVal(info);
     }
 
     virtual LegalizedVaryingVal createLegalSystemVaryingValImpl(VaryingParamInfo const& info)
@@ -915,7 +911,16 @@ protected:
     {
         SLANG_UNUSED(info);
 
-        m_sink->diagnose(m_param, Diagnostics::unimplemented, "this target doesn't support this system-defined varying parameters");
+        m_sink->diagnose(m_param, Diagnostics::unimplemented, "this target doesn't support this system-defined varying parameter");
+
+        return LegalizedVaryingVal();
+    }
+
+    LegalizedVaryingVal diagnoseUnsupportedUserVal(VaryingParamInfo const& info)
+    {
+        SLANG_UNUSED(info);
+
+        m_sink->diagnose(m_param, Diagnostics::unimplemented, "this target doesn't support this user-defined varying parameter");
 
         return LegalizedVaryingVal();
     }
@@ -1054,6 +1059,14 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
     //
     IRType* uint3Type = nullptr;
 
+    // Scans through and returns the first typeLayout attribute of non-zero size.
+    static LayoutResourceKind getLayoutResourceKind(IRTypeLayout* typeLayout) {
+        for (auto attr : typeLayout->getSizeAttrs()) {
+            if (attr->getSize() != 0) return attr->getResourceKind();
+        }
+        return LayoutResourceKind::None;
+    }
+
     void beginModuleImpl() SLANG_OVERRIDE
     {
         // Because many of the varying parameters are defined
@@ -1114,7 +1127,6 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
 
     IRInst* groupThreadIndex = nullptr;
     IRInst* dispatchThreadID = nullptr;
-
     void beginEntryPointImpl() SLANG_OVERRIDE
     {
         IRBuilder builder(m_sharedBuilder);
@@ -1174,11 +1186,38 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
         case SystemValueSemanticName::GroupThreadID:    return LegalizedVaryingVal::makeValue(threadIdxGlobalParam);
         case SystemValueSemanticName::GroupThreadIndex: return LegalizedVaryingVal::makeValue(groupThreadIndex);
         case SystemValueSemanticName::DispatchThreadID: return LegalizedVaryingVal::makeValue(dispatchThreadID);
-
         default:
             return diagnoseUnsupportedSystemVal(info);
         }
     }
+
+    LegalizedVaryingVal createLegalUserVaryingValImpl(VaryingParamInfo const& info) SLANG_OVERRIDE
+    {
+        auto layoutResourceKind = getLayoutResourceKind(info.typeLayout);
+        switch (layoutResourceKind)
+        {
+        case LayoutResourceKind::RayPayload: {
+            IRBuilder builder(m_sharedBuilder);
+			builder.setInsertBefore(m_firstOrdinaryInst);
+			IRPtrType* ptrType = builder.getPtrType(info.type);
+			IRInst* getRayPayload = builder.emitIntrinsicInst(ptrType, kIROp_GetOptiXRayPayloadPtr, 0, nullptr);
+            return LegalizedVaryingVal::makeAddress(getRayPayload);
+            // Todo: compute how many registers are required for the current payload. 
+            // If more than 32, use the above logic. 
+            // Otherwise, either use the optix_get_payload or optix_set_payload 
+            // intrinsics depending on input/output
+            /*if (info.kind == LayoutResourceKind::VaryingInput) {
+            }
+            else if (info.kind == LayoutResourceKind::VaryingOutput) {
+            }
+            else {
+                return diagnoseUnsupportedUserVal(info);
+            }*/ 
+        }
+        default:
+            return diagnoseUnsupportedUserVal(info);
+        }
+    }
 };
 
 
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index f21722bd7..f2869886e 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -686,6 +686,12 @@ struct HLSLRayTracingLayoutRulesImpl : DefaultVaryingLayoutRulesImpl
         : DefaultVaryingLayoutRulesImpl(kind)
     {}
 };
+struct CUDARayTracingLayoutRulesImpl : DefaultVaryingLayoutRulesImpl
+{
+    CUDARayTracingLayoutRulesImpl(LayoutResourceKind kind)
+        : DefaultVaryingLayoutRulesImpl(kind)
+    {}
+};
 
 DefaultLayoutRulesImpl kDefaultLayoutRulesImpl;
 Std140LayoutRulesImpl kStd140LayoutRulesImpl;
@@ -707,7 +713,13 @@ HLSLRayTracingLayoutRulesImpl kHLSLRayPayloadParameterLayoutRulesImpl(LayoutReso
 HLSLRayTracingLayoutRulesImpl kHLSLCallablePayloadParameterLayoutRulesImpl(LayoutResourceKind::CallablePayload);
 HLSLRayTracingLayoutRulesImpl kHLSLHitAttributesParameterLayoutRulesImpl(LayoutResourceKind::HitAttributes);
 
+// Just copying what was done above for now, but for CUDA...
+//CUDAVaryingLayoutRulesImpl kCUDAVaryingInputLayoutRulesImpl(LayoutResourceKind::VertexInput);
+//CUDAVaryingLayoutRulesImpl kCUDAVaryingOutputLayoutRulesImpl(LayoutResourceKind::FragmentOutput);
 //
+CUDARayTracingLayoutRulesImpl kCUDARayPayloadParameterLayoutRulesImpl(LayoutResourceKind::RayPayload);
+//CUDARayTracingLayoutRulesImpl kCUDACallablePayloadParameterLayoutRulesImpl(LayoutResourceKind::CallablePayload);
+//CUDARayTracingLayoutRulesImpl kCUDAHitAttributesParameterLayoutRulesImpl(LayoutResourceKind::HitAttributes);
 
 struct GLSLLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
 {
@@ -929,6 +941,12 @@ LayoutRulesImpl kCUDAAnyValueLayoutRulesImpl_ = {
     &kCUDAObjectLayoutRulesImpl,
 };
 
+// We want a custom layout for ray payloads to handle the logic of
+// copying payload registers vs reading / writing to and from memory
+LayoutRulesImpl kCUDARayPayloadParameterLayoutRulesImpl_ = {
+    &kCUDALayoutRulesFamilyImpl, &kCUDARayPayloadParameterLayoutRulesImpl, &kCUDAObjectLayoutRulesImpl,
+};
+
 // GLSL cases
 
 LayoutRulesImpl kStd140LayoutRulesImpl_ = {
@@ -1276,7 +1294,8 @@ LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getParameterBlockRules()
 }
 LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getRayPayloadParameterRules()
 {
-    return nullptr;
+    // Mimicking HLSL 
+    return &kCUDARayPayloadParameterLayoutRulesImpl_;
 }
 LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getCallablePayloadParameterRules()
 {