Callable shader fix and explicit payload locations for GLSL (#2185)

* Fixed the callable shader payload type for GLSL.

* Added location parameters to the __vulkanRayPayload and __vulkanCallablePayload attributes.
The default value is -1 which means use the old auto-assignment logic.

* Fixed the vkray/callable-caller test.

8 files changed, 58 insertions, 10 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index c3e41b8fc..de83eac73 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -2129,10 +2129,10 @@ attribute_syntax [numthreads(x: int, y: int = 1, z: int = 1)]   : NumThreadsAttr
 
 //
 __attributeTarget(VarDeclBase)
-attribute_syntax [__vulkanRayPayload] : VulkanRayPayloadAttribute;
+attribute_syntax [__vulkanRayPayload(location : int = -1)] : VulkanRayPayloadAttribute;
 
 __attributeTarget(VarDeclBase)
-attribute_syntax [__vulkanCallablePayload] : VulkanCallablePayloadAttribute;
+attribute_syntax [__vulkanCallablePayload(location : int = -1)] : VulkanCallablePayloadAttribute;
 
 __attributeTarget(VarDeclBase)
 attribute_syntax [__vulkanHitAttributes] : VulkanHitAttributesAttribute;
diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 971efcfcc..e404aa626 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -4500,7 +4500,7 @@ __generic<Payload>
 __specialized_for_target(glsl)
 void CallShader(uint shaderIndex, inout Payload payload)
 {
-    [__vulkanRayPayload]
+    [__vulkanCallablePayload]
     static Payload p;
 
     p = payload;
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 9a450cc31..6a142454d 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -761,23 +761,27 @@ class EntryPointAttribute : public Attribute
     Stage stage;
 };
 
-// A `[__vulkanRayPayload]` attribute, which is used in the
+// A `[__vulkanRayPayload(location)]` attribute, which is used in the
 // standard library implementation to indicate that a variable
 // actually represents the input/output interface for a Vulkan
 // ray tracing shader to pass per-ray payload information.
 class VulkanRayPayloadAttribute : public Attribute 
 {
     SLANG_AST_CLASS(VulkanRayPayloadAttribute)
+
+    int location;
 };
 
 
-// A `[__vulkanCallablePayload]` attribute, which is used in the
+// A `[__vulkanCallablePayload(location)]` attribute, which is used in the
 // standard library implementation to indicate that a variable
 // actually represents the input/output interface for a Vulkan
 // ray tracing shader to pass payload information to/from a callee.
 class VulkanCallablePayloadAttribute : public Attribute 
 {
     SLANG_AST_CLASS(VulkanCallablePayloadAttribute)
+
+    int location;
 };
 
 
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index e597407e2..320b980e9 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -539,6 +539,24 @@ namespace Slang
             }
             dllImportAttr->modulePath = libraryName;
         }
+        else if (auto rayPayloadAttr = as<VulkanRayPayloadAttribute>(attr))
+        {
+            SLANG_ASSERT(attr->args.getCount() == 1);
+            auto val = checkConstantIntVal(attr->args[0]);
+
+            if (!val) return false;
+
+            rayPayloadAttr->location = (int32_t)val->value;
+        }
+        else if (auto callablePayloadAttr = as<VulkanCallablePayloadAttribute>(attr))
+        {
+            SLANG_ASSERT(attr->args.getCount() == 1);
+            auto val = checkConstantIntVal(attr->args[0]);
+
+            if (!val) return false;
+
+            callablePayloadAttr->location = (int32_t)val->value;
+        }
         else
         {
             if(attr->args.getCount() == 0)
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 562a532cd..2d30c2109 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2945,6 +2945,14 @@ void CLikeSourceEmitter::emitInterpolationModifiers(IRInst* varInst, IRType* val
 
 UInt CLikeSourceEmitter::getRayPayloadLocation(IRInst* inst)
 {
+    if (auto rayPayloadDecoration = inst->findDecoration<IRVulkanRayPayloadDecoration>())
+    {
+        int explicitLocation = int(getIntVal(rayPayloadDecoration->getOperand(0)));
+
+        if (explicitLocation >= 0)
+            return UInt(explicitLocation);
+    }
+
     auto& map = m_mapIRValueToRayPayloadLocation;
     UInt value = 0;
     if(map.TryGetValue(inst, value))
@@ -2957,6 +2965,14 @@ UInt CLikeSourceEmitter::getRayPayloadLocation(IRInst* inst)
 
 UInt CLikeSourceEmitter::getCallablePayloadLocation(IRInst* inst)
 {
+    if (auto callablePayloadDecoration = inst->findDecoration<IRVulkanCallablePayloadDecoration>())
+    {
+        int explicitLocation = int(getIntVal(callablePayloadDecoration->getOperand(0)));
+
+        if (explicitLocation >= 0)
+            return UInt(explicitLocation);
+    }
+
     auto& map = m_mapIRValueToCallablePayloadLocation;
     UInt value = 0;
     if(map.TryGetValue(inst, value))
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 8b4eaf5cb..e4318177f 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -2891,6 +2891,16 @@ public:
     {
         addDecoration(inst, kIROp_SequentialIDDecoration, getIntValue(getUIntType(), id));
     }
+
+    void addVulkanRayPayloadDecoration(IRInst* inst, int location)
+    {
+        addDecoration(inst, kIROp_VulkanRayPayloadDecoration, getIntValue(getIntType(), location));
+    }
+
+    void addVulkanCallablePayloadDecoration(IRInst* inst, int location)
+    {
+        addDecoration(inst, kIROp_VulkanCallablePayloadDecoration, getIntValue(getIntType(), location));
+    }
 };
 
 void addHoistableInst(
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 6cdf20af4..842a1ce5f 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -1876,13 +1876,13 @@ void addVarDecorations(
         {
             builder->addInterpolationModeDecoration(inst, IRInterpolationMode::Centroid);
         }
-        else if(as<VulkanRayPayloadAttribute>(mod))
+        else if(auto rayPayloadAttr = as<VulkanRayPayloadAttribute>(mod))
         {
-            builder->addSimpleDecoration<IRVulkanRayPayloadDecoration>(inst);
+            builder->addVulkanRayPayloadDecoration(inst, rayPayloadAttr->location);
         }
-        else if(as<VulkanCallablePayloadAttribute>(mod))
+        else if(auto callablePayloadAttr = as<VulkanCallablePayloadAttribute>(mod))
         {
-            builder->addSimpleDecoration<IRVulkanCallablePayloadDecoration>(inst);
+            builder->addVulkanCallablePayloadDecoration(inst, callablePayloadAttr->location);
         }
         else if(as<VulkanHitAttributesAttribute>(mod))
         {
diff --git a/tests/vkray/callable-caller.slang.glsl b/tests/vkray/callable-caller.slang.glsl
index 91b788655..0b7a9677b 100644
--- a/tests/vkray/callable-caller.slang.glsl
+++ b/tests/vkray/callable-caller.slang.glsl
@@ -20,7 +20,7 @@ struct MaterialPayload_0
 };
 
 layout(location = 0)
-rayPayloadNV
+callableDataNV
 MaterialPayload_0 p_0;
 
 void CallShader_0(uint shaderIndex_1, inout MaterialPayload_0 payload_0)