Fix CUDA reflection for acceleration structure handle size. (#6055)

1 files changed, 20 insertions, 0 deletions

diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 74528c61a..20a18a5ac 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -791,6 +791,7 @@ static LayoutResourceKind _getHLSLLayoutResourceKind(ShaderParameterKind kind)
     case ShaderParameterKind::RawBuffer:
     case ShaderParameterKind::Buffer:
     case ShaderParameterKind::Texture:
+    case ShaderParameterKind::AccelerationStructure:
         return LayoutResourceKind::ShaderResource;
 
     case ShaderParameterKind::MutableStructuredBuffer:
@@ -798,6 +799,7 @@ static LayoutResourceKind _getHLSLLayoutResourceKind(ShaderParameterKind kind)
     case ShaderParameterKind::MutableBuffer:
     case ShaderParameterKind::MutableTexture:
     case ShaderParameterKind::AppendConsumeStructuredBuffer:
+    case ShaderParameterKind::ShaderStorageBuffer:
         return LayoutResourceKind::UnorderedAccess;
 
     case ShaderParameterKind::SamplerState:
@@ -900,8 +902,10 @@ struct HLSLObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         case ShaderParameterKind::RawBuffer:
         case ShaderParameterKind::Buffer:
         case ShaderParameterKind::Texture:
+        case ShaderParameterKind::AccelerationStructure:
             return SimpleLayoutInfo(LayoutResourceKind::ShaderResource, 1);
 
+        case ShaderParameterKind::ShaderStorageBuffer:
         case ShaderParameterKind::MutableStructuredBuffer:
         case ShaderParameterKind::MutableRawBuffer:
         case ShaderParameterKind::MutableBuffer:
@@ -1218,6 +1222,8 @@ struct CPUObjectLayoutRulesImpl : ObjectLayoutRulesImpl
                 sizeof(void*) * 2,
                 SLANG_ALIGN_OF(void*));
 
+        case ShaderParameterKind::ShaderStorageBuffer:
+        case ShaderParameterKind::AccelerationStructure:
         case ShaderParameterKind::SamplerState:
             // It's a pointer
             return SimpleLayoutInfo(
@@ -1308,6 +1314,15 @@ struct CUDAObjectLayoutRulesImpl : CPUObjectLayoutRulesImpl
                     _roundToAlignment(sizeof(CUDAPtr) + sizeof(CUDACount), sizeof(CUDAPtr));
                 return SimpleLayoutInfo(LayoutResourceKind::Uniform, size, sizeof(CUDAPtr));
             }
+        case ShaderParameterKind::ShaderStorageBuffer:
+        case ShaderParameterKind::AccelerationStructure:
+            {
+                // It's a pointer.
+                return SimpleLayoutInfo(
+                    LayoutResourceKind::Uniform,
+                    sizeof(CUDAPtr),
+                    sizeof(CUDAPtr));
+            }
         case ShaderParameterKind::SamplerState:
             {
                 // In CUDA it seems that sampler states are combined into texture objects.
@@ -1868,6 +1883,8 @@ struct MetalObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         case ShaderParameterKind::Buffer:
         case ShaderParameterKind::MutableRawBuffer:
         case ShaderParameterKind::MutableBuffer:
+        case ShaderParameterKind::ShaderStorageBuffer:
+        case ShaderParameterKind::AccelerationStructure:
             return SimpleLayoutInfo(LayoutResourceKind::MetalBuffer, 1);
         case ShaderParameterKind::AppendConsumeStructuredBuffer:
             return SimpleLayoutInfo(LayoutResourceKind::MetalBuffer, 2);
@@ -1934,6 +1951,8 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
         case ShaderParameterKind::TextureUniformBuffer:
         case ShaderParameterKind::Texture:
         case ShaderParameterKind::SamplerState:
+        case ShaderParameterKind::ShaderStorageBuffer:
+        case ShaderParameterKind::AccelerationStructure:
             {
                 return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 1);
             }
@@ -4657,6 +4676,7 @@ static TypeLayoutResult _createTypeLayout(TypeLayoutContext& context, Type* type
     CASE(GLSLInputAttachmentType, InputRenderTarget);
 
     // This case is mostly to allow users to add new resource types...
+    CASE(RaytracingAccelerationStructureType, AccelerationStructure);
     CASE(UntypedBufferResourceType, RawBuffer);
 
     CASE(GLSLShaderStorageBufferType, MutableRawBuffer);