Add callable shader support for Vulkan ray tracing (shader-slang#718)

* Add callable shader support for Vulkan ray tracing

This change extends the previous work to update Vulkan ray tracing support for the finished `GL_NV_ray_tracing` spec.
One of the features missing in the experimental extension that was added to the final spec is "callable shaders," which allow ray tracing shaders to call other shaders as general-purpose subroutines.

Most of the implementation work here mirrors what was done for the `TraceRay()` function to map it to `traceNV()`.
We map the generic `CallShader<P>` function to the non-generic `executeCallableNV`, with a payload identifier that indicates a specific global variable of type `P` (the global variable being generated from a `static` local in `CallShader`). A new modifier is added to identify the payload structure, and the parameter binding/layout logic introduces a new resource kind for callable-shader payload data (where previously the logic had assumed ray and callable payloads should use the same resource kind).

Two test shaders are included: one for the callable shader (`callable.slang`) and one for a ray generation shader that calls it (`callable-caller.slang`). Just for kicks, the payload data type is defined in a shared file so that we can be sure the two agree (trying to emulate what might be good practice, and ensure that ray tracing support works together with other Slang mechanisms).

* Typo fix: assocaited->associated

One instance was found in review, but I went ahead and fixed a bunch since I seem to make this typo a lot.

* Typo fix: defintiion->definition

Author: Tim Foley

examples/model-viewer/shaders.slang

@@ -155,7 +155,7 @@ struct SimpleMaterial : IMaterial
     // To satisfy the requirements of the `IMaterial` interface, our
     // material type needs to provide a suitable `BRDF` type. We
     // do this by using a simple `typedef`, although a nested
-    // `struct` type can also satisfy an assocaited type requirement.
+    // `struct` type can also satisfy an associated type requirement.
     //
     // A future version of the Slang compiler may allow the "right"
     // associated type definition to be inferred from the signature
@@ -459,7 +459,7 @@ float4 fragmentMain(
     // from different light sources.
     //
     // Note that the return type here is `TMaterial.BRDF`,
-    // which is the `BRDF` type *assocaited* with the (unknown)
+    // which is the `BRDF` type *associated* with the (unknown)
     // `TMaterial` type. When `TMaterial` gets substituted for
     // a concrete type later (e.g., `SimpleMaterial`) this
     // will resolve to a concrete type too (e.g., `SimpleMaterial.BRDF`

slang.h

@@ -1378,6 +1378,7 @@ extern "C"
 
         SLANG_PARAMETER_CATEGORY_RAY_PAYLOAD,
         SLANG_PARAMETER_CATEGORY_HIT_ATTRIBUTES,
+        SLANG_PARAMETER_CATEGORY_CALLABLE_PAYLOAD,
 
         //
         SLANG_PARAMETER_CATEGORY_COUNT,
@@ -1681,6 +1682,7 @@ namespace slang
 
         RayPayload = SLANG_PARAMETER_CATEGORY_RAY_PAYLOAD,
         HitAttributes = SLANG_PARAMETER_CATEGORY_HIT_ATTRIBUTES,
+        CallablePayload = SLANG_PARAMETER_CATEGORY_CALLABLE_PAYLOAD,
 
         // DEPRECATED:
         VertexInput = SLANG_PARAMETER_CATEGORY_VERTEX_INPUT,

source/slang-glslang/slang-glslang.cpp

@@ -102,6 +102,7 @@ static int glslang_compileGLSLToSPIRV(glslang_CompileRequest* request)
     CASE(ANY_HIT,           AnyHitNV);
     CASE(CLOSEST_HIT,       ClosestHitNV);
     CASE(MISS,              MissNV);
+    CASE(CALLABLE,          CallableNV);
 
 #undef CASE
 

source/slang/core.meta.slang

@@ -1233,6 +1233,9 @@ attribute_syntax [numthreads(x: int, y: int = 1, z: int = 1)]   : NumThreadsAttr
 __attributeTarget(VarDeclBase)
 attribute_syntax [__vulkanRayPayload] : VulkanRayPayloadAttribute;
 
+__attributeTarget(VarDeclBase)
+attribute_syntax [__vulkanCallablePayload] : VulkanCallablePayloadAttribute;
+
 __attributeTarget(VarDeclBase)
 attribute_syntax [__vulkanHitAttributes] : VulkanHitAttributesAttribute;
 

source/slang/core.meta.slang.h

@@ -1252,6 +1252,9 @@ SLANG_RAW("__attributeTarget(VarDeclBase)\n")
 SLANG_RAW("attribute_syntax [__vulkanRayPayload] : VulkanRayPayloadAttribute;\n")
 SLANG_RAW("\n")
 SLANG_RAW("__attributeTarget(VarDeclBase)\n")
+SLANG_RAW("attribute_syntax [__vulkanCallablePayload] : VulkanCallablePayloadAttribute;\n")
+SLANG_RAW("\n")
+SLANG_RAW("__attributeTarget(VarDeclBase)\n")
 SLANG_RAW("attribute_syntax [__vulkanHitAttributes] : VulkanHitAttributesAttribute;\n")
 SLANG_RAW("\n")
 SLANG_RAW("__attributeTarget(FunctionDeclBase)\n")

source/slang/emit.cpp

@@ -148,6 +148,7 @@ struct SharedEmitContext
     DiagnosticSink* getSink() { return &entryPoint->compileRequest->mSink; }
 
     Dictionary<IRInst*, UInt> mapIRValueToRayPayloadLocation;
+    Dictionary<IRInst*, UInt> mapIRValueToCallablePayloadLocation;
 };
 
 struct EmitContext
@@ -3315,7 +3316,7 @@ struct EmitVisitor
                         case 'P':
                             {
                                 // The `$XP` case handles looking up
-                                // the assocaited `location` for a variable
+                                // the associated `location` for a variable
                                 // used as the argument ray payload at a
                                 // trace call site.
 
@@ -3329,6 +3330,23 @@ struct EmitVisitor
                             }
                             break;
 
+                        case 'C':
+                            {
+                                // The `$XC` case handles looking up
+                                // the associated `location` for a variable
+                                // used as the argument callable payload at a
+                                // call site.
+
+                                UInt argIndex = 0;
+                                SLANG_RELEASE_ASSERT(argCount > argIndex);
+                                auto arg = args[argIndex].get();
+                                auto argLoad = as<IRLoad>(arg);
+                                SLANG_RELEASE_ASSERT(argLoad);
+                                auto argVar = argLoad->getOperand(0);
+                                Emit(getCallablePayloadLocation(ctx, argVar));
+                            }
+                            break;
+
                         case 'T':
                             {
                                 // The `$XT` case handles selecting between
@@ -5347,6 +5365,20 @@ struct EmitVisitor
         return value;
     }
 
+    UInt getCallablePayloadLocation(
+        EmitContext*    ctx,
+        IRInst*         inst)
+    {
+        auto& map = ctx->shared->mapIRValueToCallablePayloadLocation;
+        UInt value = 0;
+        if(map.TryGetValue(inst, value))
+            return value;
+
+        value = map.Count();
+        map.Add(inst, value);
+        return value;
+    }
+
     void emitIRVarModifiers(
         EmitContext*    ctx,
         VarLayout*      layout,
@@ -5365,6 +5397,13 @@ struct EmitVisitor
             emit(")\n");
             emit("rayPayloadNV\n");
         }
+        if(varDecl->findDecoration<IRVulkanCallablePayloadDecoration>())
+        {
+            emit("layout(location = ");
+            Emit(getCallablePayloadLocation(ctx, varDecl));
+            emit(")\n");
+            emit("callableDataNV\n");
+        }
         if(varDecl->findDecoration<IRVulkanHitAttributesDecoration>())
         {
             emit("hitAttributeNV\n");
@@ -5520,6 +5559,12 @@ struct EmitVisitor
                     }
                     break;
 
+                case LayoutResourceKind::CallablePayload:
+                    {
+                        emit("callableDataInNV ");
+                    }
+                    break;
+
                 case LayoutResourceKind::HitAttributes:
                     {
                         emit("hitAttributeNV ");

source/slang/hlsl.meta.slang

@@ -1355,8 +1355,38 @@ struct BuiltInTriangleIntersectionAttributes
 // 10.3 - Intrinsics
 
 // 10.3.1
-__target_intrinsic(glsl, "callableShadersAreNotYetAvailableInVulkan")
-void CallShader<param_t>(uint ShaderIndex, inout param_t Parameter);
+
+void CallShader<Payload>(uint shaderIndex, inout Payload payload);
+
+// `executeCallableNV` is the GLSL intrinsic that will be used to implement
+// `CallShader()` for GLSL-based targets.
+//
+__target_intrinsic(glsl, "executeCallableNV")
+void __executeCallableNV(uint shaderIndex, int payloadLocation);
+
+// Next is the custom intrinsic that will compute the payload location
+// for a type being used in a `CallShader()` call for GLSL-based targets.
+//
+__generic<Payload>
+__target_intrinsic(glsl, "$XC")
+[__readNone]
+int __callablePayloadLocation(Payload payload);
+
+// Now we provide a hard-coded definition of `CallShader()` for GLSL-based
+// targets, which maps the generic HLSL operation into the non-generic
+// GLSL equivalent.
+//
+__generic<Payload>
+__specialized_for_target(glsl)
+void CallShader(uint shaderIndex, inout Payload payload)
+{
+    [__vulkanRayPayload]
+    static Payload p;
+
+    p = payload;
+    __executeCallableNV(shaderIndex, __callablePayloadLocation(p));
+    payload = p;
+}
 
 // 10.3.2
 void TraceRay<payload_t>(

source/slang/hlsl.meta.slang.h

@@ -1403,8 +1403,38 @@ SLANG_RAW("\n")
 SLANG_RAW("// 10.3 - Intrinsics\n")
 SLANG_RAW("\n")
 SLANG_RAW("// 10.3.1\n")
-SLANG_RAW("__target_intrinsic(glsl, \"callableShadersAreNotYetAvailableInVulkan\")\n")
-SLANG_RAW("void CallShader<param_t>(uint ShaderIndex, inout param_t Parameter);\n")
+SLANG_RAW("\n")
+SLANG_RAW("void CallShader<Payload>(uint shaderIndex, inout Payload payload);\n")
+SLANG_RAW("\n")
+SLANG_RAW("// `executeCallableNV` is the GLSL intrinsic that will be used to implement\n")
+SLANG_RAW("// `CallShader()` for GLSL-based targets.\n")
+SLANG_RAW("//\n")
+SLANG_RAW("__target_intrinsic(glsl, \"executeCallableNV\")\n")
+SLANG_RAW("void __executeCallableNV(uint shaderIndex, int payloadLocation);\n")
+SLANG_RAW("\n")
+SLANG_RAW("// Next is the custom intrinsic that will compute the payload location\n")
+SLANG_RAW("// for a type being used in a `CallShader()` call for GLSL-based targets.\n")
+SLANG_RAW("//\n")
+SLANG_RAW("__generic<Payload>\n")
+SLANG_RAW("__target_intrinsic(glsl, \"$XC\")\n")
+SLANG_RAW("[__readNone]\n")
+SLANG_RAW("int __callablePayloadLocation(Payload payload);\n")
+SLANG_RAW("\n")
+SLANG_RAW("// Now we provide a hard-coded definition of `CallShader()` for GLSL-based\n")
+SLANG_RAW("// targets, which maps the generic HLSL operation into the non-generic\n")
+SLANG_RAW("// GLSL equivalent.\n")
+SLANG_RAW("//\n")
+SLANG_RAW("__generic<Payload>\n")
+SLANG_RAW("__specialized_for_target(glsl)\n")
+SLANG_RAW("void CallShader(uint shaderIndex, inout Payload payload)\n")
+SLANG_RAW("{\n")
+SLANG_RAW("    [__vulkanRayPayload]\n")
+SLANG_RAW("    static Payload p;\n")
+SLANG_RAW("\n")
+SLANG_RAW("    p = payload;\n")
+SLANG_RAW("    __executeCallableNV(shaderIndex, __callablePayloadLocation(p));\n")
+SLANG_RAW("    payload = p;\n")
+SLANG_RAW("}\n")
 SLANG_RAW("\n")
 SLANG_RAW("// 10.3.2\n")
 SLANG_RAW("void TraceRay<payload_t>(\n")

source/slang/ir-insts.h

@@ -125,6 +125,14 @@ struct IRVulkanRayPayloadDecoration : IRDecoration
     enum { kDecorationOp = kIRDecorationOp_VulkanRayPayload };
 };
 
+/// A decoration that indicates that a variable represents
+/// a vulkan callable shader payload, and should have a location assigned
+/// to it.
+struct IRVulkanCallablePayloadDecoration : IRDecoration
+{
+    enum { kDecorationOp = kIRDecorationOp_VulkanCallablePayload };
+};
+
 /// A decoration that indicates that a variable represents
 /// vulkan hit attributes, and should have a location assigned
 /// to it.
@@ -978,7 +986,7 @@ IRGlobalValue* getSpecializedGlobalValueForDeclRef(
 struct ExtensionUsageTracker;
 
 // Clone the IR values reachable from the given entry point
-// into the IR module assocaited with the specialization state.
+// into the IR module associated with the specialization state.
 // When multiple definitions of a symbol are found, the one
 // that is best specialized for the given `targetReq` will be
 // used.

source/slang/ir-serialize.cpp

@@ -704,6 +704,7 @@ Result IRSerialWriter::write(IRModule* module, SourceManager* sourceManager, Opt
                     break;
                 }
                 case kIRDecorationOp_VulkanRayPayload:
+                case kIRDecorationOp_VulkanCallablePayload:
                 case kIRDecorationOp_VulkanHitAttributes:
                 case kIRDecorationOp_ReadNone:
                 {
@@ -1560,6 +1561,12 @@ IRDecoration* IRSerialReader::_createDecoration(const Ser::Inst& srcInst)
             SLANG_ASSERT(srcInst.m_payloadType == PayloadType::Empty);
             return decor;
         }
+        case kIRDecorationOp_VulkanCallablePayload:
+        {
+            auto decor = createEmptyDecoration<IRVulkanCallablePayloadDecoration>(m_module);
+            SLANG_ASSERT(srcInst.m_payloadType == PayloadType::Empty);
+            return decor;
+        }
         case kIRDecorationOp_VulkanHitAttributes:
         {
             auto decor = createEmptyDecoration<IRVulkanHitAttributesDecoration>(m_module);

source/slang/ir.cpp

@@ -2934,6 +2934,11 @@ namespace Slang
                     dump(context, "\n[__vulkanRayPayload]");
                 }
                 break;
+            case kIRDecorationOp_VulkanCallablePayload:
+                {
+                    dump(context, "\n[__vulkanCallPayload]");
+                }
+                break;
             case kIRDecorationOp_VulkanHitAttributes:
                 {
                     dump(context, "\n[__vulkanHitAttributes]");
@@ -5377,6 +5382,12 @@ namespace Slang
                 }
                 break;
 
+            case kIRDecorationOp_VulkanCallablePayload:
+                {
+                    context->builder->addDecoration<IRVulkanCallablePayloadDecoration>(clonedValue);
+                }
+                break;
+
             case kIRDecorationOp_VulkanHitAttributes:
                 {
                     context->builder->addDecoration<IRVulkanHitAttributesDecoration>(clonedValue);

source/slang/ir.h

@@ -158,6 +158,7 @@ enum IRDecorationOp : uint16_t
     kIRDecorationOp_RequireGLSLVersion,
     kIRDecorationOp_RequireGLSLExtension,
     kIRDecorationOp_ReadNone,
+    kIRDecorationOp_VulkanCallablePayload,
 
     kIRDecorationOp_CountOf
 };

source/slang/lower-to-ir.cpp

@@ -1260,6 +1260,10 @@ void addVarDecorations(
         {
             builder->addDecoration<IRVulkanRayPayloadDecoration>(inst);
         }
+        else if(mod.As<VulkanCallablePayloadAttribute>())
+        {
+            builder->addDecoration<IRVulkanCallablePayloadDecoration>(inst);
+        }
         else if(mod.As<VulkanHitAttributesAttribute>())
         {
             builder->addDecoration<IRVulkanHitAttributesDecoration>(inst);

source/slang/modifier-defs.h

@@ -393,6 +393,12 @@ END_SYNTAX_CLASS()
 // ray tracing shader to pass per-ray payload information.
 SIMPLE_SYNTAX_CLASS(VulkanRayPayloadAttribute, Attribute)
 
+// A `[__vulkanCallablePayload]` 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.
+SIMPLE_SYNTAX_CLASS(VulkanCallablePayloadAttribute, Attribute)
+
 // A `[__vulkanHitAttributes]` attribute, which is used in the
 // standard library implementation to indicate that a variable
 // actually represents the output interface for a Vulkan

source/slang/options.cpp

@@ -103,7 +103,7 @@ struct OptionsParser
 
     // An entry point represents a function to be checked and possibly have
     // code generated in one of our translation units. An entry point
-    // needs to have an assocaited stage, which might come via the
+    // needs to have an associated stage, which might come via the
     // `-stage` command line option, or a `[shader("...")]` attribute
     // in the source code.
     //
@@ -1176,7 +1176,7 @@ struct OptionsParser
             // need to support output formats that can store multiple
             // entry points in one file).
 
-            // If an output doesn't have a target assocaited with
+            // If an output doesn't have a target associated with
             // it, then search for the target with the matching format.
             if( rawOutput.targetIndex == -1 )
             {

source/slang/parameter-binding.cpp

@@ -1910,13 +1910,19 @@ static RefPtr<TypeLayout> processEntryPointParameter(
             break;
 
         case Stage::AnyHit:
-        case Stage::Callable:
         case Stage::ClosestHit:
         case Stage::Miss:
             // `in out` or `out` parameter is payload
             return CreateTypeLayout(context->layoutContext.with(
                 context->getRulesFamily()->getRayPayloadParameterRules()),
                 type);
+
+        case Stage::Callable:
+            // `in out` or `out` parameter is payload
+            return CreateTypeLayout(context->layoutContext.with(
+                context->getRulesFamily()->getCallablePayloadParameterRules()),
+                type);
+
         }
     }
     else

source/slang/preprocessor.cpp

@@ -478,7 +478,7 @@ static PreprocessorMacro* LookupMacro(PreprocessorEnvironment* environment, Name
 
 static PreprocessorEnvironment* GetCurrentEnvironment(Preprocessor* preprocessor)
 {
-    // The environment we will use for looking up a macro is assocaited
+    // The environment we will use for looking up a macro is associated
     // with the current input stream (because it may include entries
     // for macro arguments).
     //