Initial support for DXR payload access qualifiers (#1705)

This change adds initial support for a feature being proposed for inclusion in dxc: https://github.com/microsoft/DirectXShaderCompiler/pull/3171.

The main features are:

* A `[payload]` attribute that indicates which `struct` types are intended to be used as payloads. Consistent use of this attribute should mean that an application no longer needs to manually specify a maximum payload size when creating a ray-tracing pipeline.

* `read(...)` and `write(...)` qualifiers which can be attached to fields of `struct` types (usually `[payload]`-attributed types) to indicate which ray tracing pipeline stages are allowed read/write access to that part of the payload. Use of these qualifiers should allow an implementation to optimize storage of ray payload elements across RT pipeline stages.

The work in this change just adds basic parsing for these features, translation to matching IR decorations, and then emission of HLSL text based on those decorations.
Notable gaps in this first change include:

* No work is currently being done to validate access to ray payloads in RT entry points based on these qualifiers.

* The stage names in `read(...)` and `write(...)` are not being validated, and are being stored in the IR as text. These should probably use the `Stage` enumeration in some fashion, but we would need to have a way to encode the additional `caller` pseudo-stage that the feature uses.

* No work is currently being done to adjust or react to the chosen shader model when emitting HLSL code. We should *either* have these attributes force a switch to a higher shader model, *or* skip emission of these attributes if the chosen shader model / profile does not imply support for them.

* No tests are currently included for this work, because tests would rely on using a custom `dxcompiler.dll` build with the new feature supported.

10 files changed, 191 insertions, 1 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index 887852cbc..970ed1fa8 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -2098,3 +2098,6 @@ attribute_syntax [__requiresNVAPI] : RequiresNVAPIAttribute;
 
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [noinline] : NoInlineAttribute;
+
+__attributeTarget(StructDecl)
+attribute_syntax [payload] : PayloadAttribute;
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 3d0c3c6c1..10560dc84 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -292,6 +292,24 @@ class HLSLSimpleSemantic : public HLSLSemantic
     SLANG_AST_CLASS(HLSLSimpleSemantic)
 };
 
+// A semantic applied to a field of a ray-payload type, to control access
+class RayPayloadAccessSemantic : public HLSLSemantic
+{
+    SLANG_AST_CLASS(RayPayloadAccessSemantic)
+
+    List<Token> stageNameTokens;
+};
+
+class RayPayloadReadSemantic : public RayPayloadAccessSemantic
+{
+    SLANG_AST_CLASS(RayPayloadReadSemantic)
+};
+
+class RayPayloadWriteSemantic : public RayPayloadAccessSemantic
+{
+    SLANG_AST_CLASS(RayPayloadWriteSemantic)
+};
+
 
 // GLSL
 
@@ -948,5 +966,14 @@ class NoInlineAttribute : public Attribute
     SLANG_AST_CLASS(NoInlineAttribute)
 };
 
+    /// A `[payload]` attribute indicates that a `struct` type will be used as
+    /// a ray payload for `TraceRay()` calls, and thus also as input/output
+    /// for shaders in the ray tracing pipeline that might be invoked for
+    /// such a ray.
+    ///
+class PayloadAttribute : public Attribute
+{
+    SLANG_AST_CLASS(PayloadAttribute)
+};
 
 } // namespace Slang
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index d916019c7..92947214d 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2898,6 +2898,9 @@ void CLikeSourceEmitter::emitStruct(IRStructType* structType)
     }
 
     m_writer->emit("struct ");
+
+    emitPostKeywordTypeAttributes(structType);
+
     m_writer->emit(getName(structType));
     m_writer->emit("\n{\n");
     m_writer->indent();
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index 5846d99fd..a969fa334 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -253,6 +253,9 @@ public:
 
     void emitStruct(IRStructType* structType);
 
+        /// Emit type attributes that should appear after, e.g., a `struct` keyword
+    void emitPostKeywordTypeAttributes(IRInst* inst) { emitPostKeywordTypeAttributesImpl(inst); }
+
     void emitInterpolationModifiers(IRInst* varInst, IRType* valueType, IRVarLayout* layout);
 
     UInt getRayPayloadLocation(IRInst* inst);
@@ -368,6 +371,9 @@ public:
     void handleRequiredCapabilities(IRInst* inst);
     virtual void handleRequiredCapabilitiesImpl(IRInst* inst) { SLANG_UNUSED(inst); }
 
+    virtual void emitPostKeywordTypeAttributesImpl(IRInst* inst) { SLANG_UNUSED(inst); }
+
+
     void _emitArrayType(IRArrayType* arrayType, EDeclarator* declarator);
     void _emitUnsizedArrayType(IRUnsizedArrayType* arrayType, EDeclarator* declarator);
     void _emitType(IRType* type, EDeclarator* declarator);
diff --git a/source/slang/slang-emit-hlsl.cpp b/source/slang/slang-emit-hlsl.cpp
index a1da1bbdf..26a841190 100644
--- a/source/slang/slang-emit-hlsl.cpp
+++ b/source/slang/slang-emit-hlsl.cpp
@@ -915,6 +915,11 @@ void HLSLSourceEmitter::emitSemanticsImpl(IRInst* inst)
         return;
     }
 
+    if( auto readAccessSemantic = inst->findDecoration<IRStageReadAccessDecoration>())
+        _emitStageAccessSemantic(readAccessSemantic, "read");
+    if( auto writeAccessSemantic = inst->findDecoration<IRStageWriteAccessDecoration>())
+        _emitStageAccessSemantic(writeAccessSemantic, "write");
+
     if (auto layoutDecoration = inst->findDecoration<IRLayoutDecoration>())
     {
         auto layout = layoutDecoration->getLayout();
@@ -932,6 +937,32 @@ void HLSLSourceEmitter::emitSemanticsImpl(IRInst* inst)
     }
 }
 
+void HLSLSourceEmitter::_emitStageAccessSemantic(IRStageAccessDecoration* decoration, const char* name)
+{
+    Int stageCount = decoration->getStageCount();
+    if(stageCount == 0)
+        return;
+
+    m_writer->emit(" : ");
+    m_writer->emit(name);
+    m_writer->emit("(");
+    for( Int i = 0; i < stageCount; ++i )
+    {
+        if(i != 0) m_writer->emit(", ");
+        m_writer->emit(decoration->getStageName(i));
+    }
+    m_writer->emit(")");
+}
+
+void HLSLSourceEmitter::emitPostKeywordTypeAttributesImpl(IRInst* inst)
+{
+    if( auto payloadDecoration = inst->findDecoration<IRPayloadDecoration>() )
+    {
+        m_writer->emit("[payload] ");
+    }
+}
+
+
 void HLSLSourceEmitter::emitSimpleFuncParamImpl(IRParam* param)
 {
     if (auto decor = param->findDecoration<IRGeometryInputPrimitiveTypeDecoration>())
diff --git a/source/slang/slang-emit-hlsl.h b/source/slang/slang-emit-hlsl.h
index 614ee58bc..59449e519 100644
--- a/source/slang/slang-emit-hlsl.h
+++ b/source/slang/slang-emit-hlsl.h
@@ -53,6 +53,9 @@ protected:
 
     virtual void emitGlobalInstImpl(IRInst* inst) SLANG_OVERRIDE;
 
+    virtual void emitPostKeywordTypeAttributesImpl(IRInst* inst) SLANG_OVERRIDE;
+
+
         // Emit a single `register` semantic, as appropriate for a given resource-type-specific layout info
         // Keyword to use in the uniform case (`register` for globals, `packoffset` inside a `cbuffer`)
     void _emitHLSLRegisterSemantic(LayoutResourceKind kind, EmitVarChain* chain, char const* uniformSemanticSpelling = "register");
@@ -70,7 +73,8 @@ protected:
 
     void _emitHLSLDecorationSingleString(const char* name, IRFunc* entryPoint, IRStringLit* val);
     void _emitHLSLDecorationSingleInt(const char* name, IRFunc* entryPoint, IRIntLit* val);
-    
+
+    void _emitStageAccessSemantic(IRStageAccessDecoration* decoration, const char* name);
 };
 
 }
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 082de1c42..f73dbd006 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -605,6 +605,13 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
         /// Applie to an IR function and signals that inlining should not be performed unless unavoidable.
     INST(NoInlineDecoration, noInline, 0, 0)
 
+    INST(PayloadDecoration, payload, 0, 0)
+
+    /* StageAccessDecoration */
+        INST(StageReadAccessDecoration, stageReadAccess, 0, 0)
+        INST(StageWriteAccessDecoration, stageWriteAccess, 0, 0)
+    INST_RANGE(StageAccessDecoration, StageReadAccessDecoration, StageWriteAccessDecoration)
+
     INST(SemanticDecoration, semantic, 2, 0)
 
     INST_RANGE(Decoration, HighLevelDeclDecoration, SemanticDecoration)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index bde7be24c..334815d54 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -575,6 +575,30 @@ struct IRSemanticDecoration : public IRDecoration
     int getSemanticIndex() { return int(getIntVal(getSemanticIndexOperand())); }
 };
 
+struct IRStageAccessDecoration : public IRDecoration
+{
+    IR_PARENT_ISA(StageAccessDecoration)
+
+    Int getStageCount() { return (Int) getOperandCount(); }
+    IRStringLit* getStageOperand(Int index) { return cast<IRStringLit>(getOperand(index)); }
+    UnownedStringSlice getStageName(Int index) { return getStageOperand(index)->getStringSlice(); }
+};
+
+struct IRStageReadAccessDecoration : public IRStageAccessDecoration
+{
+    IR_LEAF_ISA(StageReadAccessDecoration)
+};
+
+struct IRStageWriteAccessDecoration : public IRStageAccessDecoration
+{
+    IR_LEAF_ISA(StageWriteAccessDecoration)
+};
+
+struct IRPayloadDecoration : public IRDecoration
+{
+    IR_LEAF_ISA(PayloadDecoration)
+};
+
     /// An attribute that can be attached to another instruction as an operand.
     ///
     /// Attributes serve a similar role to decorations, in that both are ways
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 58a5b0ed1..f0066b543 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -6191,6 +6191,11 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         addNameHint(context, irStruct, decl);
         addLinkageDecoration(context, irStruct, decl);
 
+        if( auto payloadAttribute = decl->findModifier<PayloadAttribute>() )
+        {
+            subBuilder->addDecoration(irStruct, kIROp_PayloadDecoration);
+        }
+
         subBuilder->setInsertInto(irStruct);
 
         // A `struct` that inherits from another `struct` must start
@@ -6261,6 +6266,20 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         return LoweredValInfo::simple(finishOuterGenerics(subBuilder, irStruct, outerGeneric));
     }
 
+    void lowerRayPayloadAccessModifier(IRInst* inst, RayPayloadAccessSemantic* semantic, IROp op)
+    {
+        auto builder = getBuilder();
+
+        List<IRInst*> operands;
+        for(auto stageNameToken : semantic->stageNameTokens)
+        {
+            IRInst* stageName = builder->getStringValue(stageNameToken.getContent());
+            operands.add(stageName);
+        }
+
+        builder->addDecoration(inst, op, operands.getBuffer(), operands.getCount());
+    }
+
     LoweredValInfo lowerMemberVarDecl(VarDecl* fieldDecl)
     {
         // Each field declaration in the AST translates into
@@ -6285,6 +6304,15 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             builder->addSemanticDecoration(irFieldKey, semanticModifier->name.getName()->text.getUnownedSlice());
         }
 
+        if( auto readModifier = fieldDecl->findModifier<RayPayloadReadSemantic>() )
+        {
+            lowerRayPayloadAccessModifier(irFieldKey, readModifier, kIROp_StageReadAccessDecoration);
+        }
+        if( auto writeModifier = fieldDecl->findModifier<RayPayloadWriteSemantic>())
+        {
+            lowerRayPayloadAccessModifier(irFieldKey, writeModifier, kIROp_StageWriteAccessDecoration);
+        }
+
         // We allow a field to be marked as a target intrinsic,
         // so that we can override its mangled name in the
         // output for the chosen target.
diff --git a/source/slang/slang-parser.cpp b/source/slang/slang-parser.cpp
index 8ea51645a..68d06714d 100644
--- a/source/slang/slang-parser.cpp
+++ b/source/slang/slang-parser.cpp
@@ -2202,6 +2202,40 @@ namespace Slang
         parser->sink->diagnose(semantic, Diagnostics::packOffsetNotSupported);
     }
 
+    static RayPayloadAccessSemantic* _parseRayPayloadAccessSemantic(Parser* parser, RayPayloadAccessSemantic* semantic)
+    {
+        parser->FillPosition(semantic);
+
+        // Read the keyword that introduced the semantic
+        semantic->name = parser->ReadToken(TokenType::Identifier);
+
+        parser->ReadToken(TokenType::LParent);
+
+        for(;;)
+        {
+            if(AdvanceIfMatch(parser, TokenType::RParent))
+                break;
+
+            auto stageName = parser->ReadToken(TokenType::Identifier);
+            semantic->stageNameTokens.add(stageName);
+
+            if(AdvanceIfMatch(parser, TokenType::RParent))
+                break;
+
+            expect(parser, TokenType::Comma);
+        }
+
+        return semantic;
+    }
+
+    template<typename T>
+    static T* _parseRayPayloadAccessSemantic(Parser* parser)
+    {
+        T* semantic = parser->astBuilder->create<T>();
+        _parseRayPayloadAccessSemantic(parser, semantic);
+        return semantic;
+    }
+
     //
     // semantic ::= identifier ( '(' args ')' )?
     //
@@ -2222,6 +2256,14 @@ namespace Slang
             parseHLSLPackOffsetSemantic(parser, semantic);
             return semantic;
         }
+        else if( parser->LookAheadToken("read") && parser->LookAheadToken(TokenType::LParent, 1) )
+        {
+            return _parseRayPayloadAccessSemantic<RayPayloadReadSemantic>(parser);
+        }
+        else if( parser->LookAheadToken("write") && parser->LookAheadToken(TokenType::LParent, 1) )
+        {
+            return _parseRayPayloadAccessSemantic<RayPayloadWriteSemantic>(parser);
+        }
         else if (parser->LookAheadToken(TokenType::Identifier))
         {
             HLSLSimpleSemantic* semantic = parser->astBuilder->create<HLSLSimpleSemantic>();
@@ -3414,6 +3456,21 @@ namespace Slang
         FillPosition(rs);
         ReadToken("struct");
 
+        // The `struct` keyword may optionally be followed by
+        // attributes that appertain to the struct declaration
+        // itself, and not to any variables declared using this
+        // type specifier.
+        //
+        // TODO: We don't yet correctly associate attributes with
+        // a variable decarlation vs. a struct type when a variable
+        // is declared with a struct type specified.
+        //
+        if(LookAheadToken(TokenType::LBracket))
+        {
+            Modifier** modifierLink = &rs->modifiers.first;
+            ParseSquareBracketAttributes(this, &modifierLink);
+        }
+
         // TODO: support `struct` declaration without tag
         rs->nameAndLoc = expectIdentifier(this);