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

4 files changed, 30 insertions, 1 deletions

diff --git a/source/slang/slang-reflection-json.cpp b/source/slang/slang-reflection-json.cpp
index 57183f54d..be6539246 100644
--- a/source/slang/slang-reflection-json.cpp
+++ b/source/slang/slang-reflection-json.cpp
@@ -416,6 +416,7 @@ static void emitReflectionResourceTypeBaseInfoJSON(
         CASE(TEXTURE_BUFFER, textureBuffer);
         CASE(STRUCTURED_BUFFER, structuredBuffer);
         CASE(BYTE_ADDRESS_BUFFER, byteAddressBuffer);
+        CASE(ACCELERATION_STRUCTURE, accelerationStructure);
 #undef CASE
     }
     writer << "\"";
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);
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index 91132c76a..7d3dd8369 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -975,7 +975,7 @@ enum class ShaderParameterKind
     AtomicUint,
 
     SubpassInput,
-
+    AccelerationStructure,
     ParameterBlock,
 };
 
diff --git a/tests/reflection/acceleration-structure.slang b/tests/reflection/acceleration-structure.slang
new file mode 100644
index 000000000..cfb7d232b
--- /dev/null
+++ b/tests/reflection/acceleration-structure.slang
@@ -0,0 +1,8 @@
+//TEST(64-bit):REFLECTION(filecheck=CHECK): -target cuda
+
+// CHECK: "binding": {"kind": "uniform", "offset": 0, "size": 8},
+uniform RaytracingAccelerationStructure accel;
+
+[numthreads(1,1,1)]
+void main()
+{}
\ No newline at end of file