Enable tracing rays with OptiX backend (#1871)

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

* changing op to more descriptive name

* small tweak to allow for dumping out intermediate source for cuda targets

* initial trace ray call compiling

* hit attributes now work for float and int types, and vectors thereof

* Hitgroups using structs and arrays now work with optix

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

5 files changed, 129 insertions, 5 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 3d7217b99..bb573c2b2 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -4220,18 +4220,23 @@ static const RAY_FLAG RAY_FLAG_SKIP_PROCEDURAL_PRIMITIVES       = 0x200;
 // 10.1.2 - Ray Description Structure
 
 __target_intrinsic(hlsl, RayDesc)
+__target_intrinsic(cuda, RayDesc)
 struct RayDesc
 {
     __target_intrinsic(hlsl, Origin)
+    __target_intrinsic(cuda, Origin)
     float3 Origin;
 
     __target_intrinsic(hlsl, TMin)
+    __target_intrinsic(cuda, TMin)
     float  TMin;
 
     __target_intrinsic(hlsl, Direction)
+    __target_intrinsic(cuda, Direction)
     float3 Direction;
 
     __target_intrinsic(hlsl, TMax)
+    __target_intrinsic(cuda, TMax)
     float  TMax;
 };
 
@@ -4303,6 +4308,7 @@ void CallShader(uint shaderIndex, inout Payload payload)
 
 // 10.3.2
 __target_intrinsic(hlsl)
+__target_intrinsic(cuda, "traceOptiXRay")
 void TraceRay<payload_t>(
     RaytracingAccelerationStructure AccelerationStructure,
     uint                            RayFlags,
diff --git a/source/slang/slang-emit-cuda.cpp b/source/slang/slang-emit-cuda.cpp
index 4c430cfa6..0a375cb56 100644
--- a/source/slang/slang-emit-cuda.cpp
+++ b/source/slang/slang-emit-cuda.cpp
@@ -788,6 +788,29 @@ bool CUDASourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOu
             m_writer->emit(")getOptiXRayPayloadPtr()");
             return true;
         }
+        case kIROp_GetOptiXHitAttribute:
+        {
+            auto typeToFetch = inst->getOperand(0);
+            auto idxInst = as<IRIntLit>(inst->getOperand(1));
+            IRIntegerValue idx = idxInst->getValue();
+            if (typeToFetch->getOp() == kIROp_FloatType) {
+                m_writer->emit("__int_as_float(optixGetAttribute_");
+            }
+            else
+            {
+                m_writer->emit("optixGetAttribute_");
+            }
+            m_writer->emit(idx);
+            if (typeToFetch->getOp() == kIROp_FloatType)
+            {
+                m_writer->emit("())");
+            }
+            else
+            {
+                m_writer->emit("()");
+            }
+            return true;
+        }
         default: break;
     }
 
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 17ae652ae..847ce1d80 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -501,6 +501,9 @@ INST(GpuForeach, gpuForeach, 3, 0)
 // a pointer represented by two payload registers.
 INST(GetOptiXRayPayloadPtr, getOptiXRayPayloadPtr, 0, 0)
 
+// Wrapper for OptiX intrinsics used to load a single hit attribute
+INST(GetOptiXHitAttribute, getOptiXHitAttribute, 1, 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 f500b100f..8025bb088 100644
--- a/source/slang/slang-ir-legalize-varying-params.cpp
+++ b/source/slang/slang-ir-legalize-varying-params.cpp
@@ -1067,6 +1067,83 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
         return LayoutResourceKind::None;
     }
 
+    IRInst* emitOptiXAttributeFetch(int& ioBaseAttributeIndex, IRType* typeToFetch, IRBuilder* builder) {
+        if (auto structType = as<IRStructType>(typeToFetch))
+        {
+            List<IRInst*> fieldVals;
+            for (auto field : structType->getFields())
+            {
+                auto fieldType = field->getFieldType();
+                auto fieldVal = emitOptiXAttributeFetch(ioBaseAttributeIndex, fieldType, builder);
+                if (!fieldVal)
+                    return nullptr;
+
+                fieldVals.add(fieldVal);
+            }
+            return builder->emitMakeStruct(typeToFetch, fieldVals);
+        }
+        else if (auto arrayType = as<IRArrayTypeBase>(typeToFetch))
+        {
+            auto elementCountInst = as<IRIntLit>(arrayType->getElementCount());
+            IRIntegerValue elementCount = elementCountInst->getValue();
+            auto elementType = arrayType->getElementType();
+            List<IRInst*> elementVals;
+            for (IRIntegerValue ii = 0; ii < elementCount; ++ii)
+            {
+                auto elementVal = emitOptiXAttributeFetch(ioBaseAttributeIndex, elementType, builder);
+                if (!elementVal)
+                    return nullptr;
+                elementVals.add(elementVal);
+            }
+            return builder->emitMakeArray(typeToFetch, elementVals.getCount(), elementVals.getBuffer());
+        }
+        else if (auto matType = as<IRMatrixType>(typeToFetch))
+        {
+            auto rowCountInst = as<IRIntLit>(matType->getRowCount());
+            if (rowCountInst)
+            {
+                auto rowType = builder->getVectorType(matType->getElementType(), matType->getColumnCount());
+                IRType* elementType = rowType;
+                IRIntegerValue elementCount = rowCountInst->getValue();
+                List<IRInst*> elementVals;
+                for (IRIntegerValue ii = 0; ii < elementCount; ++ii)
+                {
+                    auto elementVal = emitOptiXAttributeFetch(ioBaseAttributeIndex, elementType, builder);
+                    if (!elementVal)
+                        return nullptr;
+                    elementVals.add(elementVal);
+                }
+                return builder->emitIntrinsicInst(typeToFetch, kIROp_MakeMatrix, elementVals.getCount(), elementVals.getBuffer());
+            }
+        }
+        else if (auto vecType = as<IRVectorType>(typeToFetch))
+        {
+            auto elementCountInst = as<IRIntLit>(vecType->getElementCount());
+            IRIntegerValue elementCount = elementCountInst->getValue();
+            IRType* elementType = vecType->getElementType();
+            List<IRInst*> elementVals;
+            for (IRIntegerValue ii = 0; ii < elementCount; ++ii)
+            {
+                auto elementVal = emitOptiXAttributeFetch(ioBaseAttributeIndex, elementType, builder);
+                if (!elementVal)
+                    return nullptr;
+                elementVals.add(elementVal);
+            }
+            return builder->emitMakeVector(typeToFetch, elementVals.getCount(), elementVals.getBuffer());
+        }
+        else if (auto basicType = as<IRBasicType>(typeToFetch))
+        {
+            IRIntegerValue idx = ioBaseAttributeIndex;
+            auto idxInst = builder->getIntValue(builder->getIntType(), idx);
+            ioBaseAttributeIndex++;
+            IRInst* args[] = { typeToFetch, idxInst };
+            IRInst* getAttr = builder->emitIntrinsicInst(typeToFetch, kIROp_GetOptiXHitAttribute, 2, args);
+            return getAttr;
+        }
+
+        return nullptr;
+    }
+
     void beginModuleImpl() SLANG_OVERRIDE
     {
         // Because many of the varying parameters are defined
@@ -1198,9 +1275,9 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
         {
         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);
+            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. 
@@ -1214,6 +1291,17 @@ struct CUDAEntryPointVaryingParamLegalizeContext : EntryPointVaryingParamLegaliz
                 return diagnoseUnsupportedUserVal(info);
             }*/ 
         }
+        case LayoutResourceKind::HitAttributes: {
+            IRBuilder builder(m_sharedBuilder);
+            builder.setInsertBefore(m_firstOrdinaryInst);
+            int ioBaseAttributeIndex = 0;
+            IRInst* getHitAttributes = emitOptiXAttributeFetch(/*ioBaseAttributeIndex*/ ioBaseAttributeIndex, /* type to fetch */info.type, /*the builder in use*/ &builder);
+            if (ioBaseAttributeIndex > 8) {
+                m_sink->diagnose(m_param, Diagnostics::unexpected, "the supplied hit attribute exceeds the maximum hit attribute structure size (32 bytes)");
+                return LegalizedVaryingVal();
+            }
+            return LegalizedVaryingVal::makeValue(getHitAttributes);
+        }
         default:
             return diagnoseUnsupportedUserVal(info);
         }
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 42a93f3f0..2568547a4 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -719,7 +719,7 @@ HLSLRayTracingLayoutRulesImpl kHLSLHitAttributesParameterLayoutRulesImpl(LayoutR
 //
 CUDARayTracingLayoutRulesImpl kCUDARayPayloadParameterLayoutRulesImpl(LayoutResourceKind::RayPayload);
 //CUDARayTracingLayoutRulesImpl kCUDACallablePayloadParameterLayoutRulesImpl(LayoutResourceKind::CallablePayload);
-//CUDARayTracingLayoutRulesImpl kCUDAHitAttributesParameterLayoutRulesImpl(LayoutResourceKind::HitAttributes);
+CUDARayTracingLayoutRulesImpl kCUDAHitAttributesParameterLayoutRulesImpl(LayoutResourceKind::HitAttributes);
 
 struct GLSLLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
 {
@@ -947,6 +947,10 @@ LayoutRulesImpl kCUDARayPayloadParameterLayoutRulesImpl_ = {
     &kCUDALayoutRulesFamilyImpl, &kCUDARayPayloadParameterLayoutRulesImpl, &kCUDAObjectLayoutRulesImpl,
 };
 
+LayoutRulesImpl kCUDAHitAttributesParameterLayoutRulesImpl_ = {
+    &kCUDALayoutRulesFamilyImpl, &kCUDAHitAttributesParameterLayoutRulesImpl, &kCUDAObjectLayoutRulesImpl,
+};
+
 // GLSL cases
 
 LayoutRulesImpl kStd140LayoutRulesImpl_ = {
@@ -1303,7 +1307,7 @@ LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getCallablePayloadParameterRules()
 }
 LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getHitAttributesParameterRules()
 {
-    return nullptr;
+    return &kCUDAHitAttributesParameterLayoutRulesImpl_;
 }
 LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getShaderRecordConstantBufferRules()
 {