Add `vk::offset` to specify member offsets for push constants (#6797)

* Add struct member offset qualifier for SPIRV

* Implement for GLSL target and add tests

* clean up

* fix formatting

* fix typo

* renamed GLSLStructOffset to VkStructOffset and added emit-spirv-via-glsl test case

15 files changed, 159 insertions, 7 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index 481aba191..5911c997c 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -3710,6 +3710,13 @@ attribute_syntax [vk_location(location : int)] : GLSLLocationAttribute;
 __attributeTarget(VarDeclBase)
 attribute_syntax [vk_index(index : int)] : GLSLIndexAttribute;
 
+/// Declare offset for a struct field. Applies to all kinds of structs when targeting Vulkan.
+/// Applies only to structs that are directly used as interface blocks (such as push constants and uniforms)
+/// when targeting GLSL, as GLSL does not support the `offset` qualifier on regular structs.
+/// This attribute has no effect on other targets.
+__attributeTarget(VarDeclBase)
+attribute_syntax [vk_offset(index: int)] : VkStructOffsetAttribute;
+
 /// @deprecated
 /// Use `spirv_asm` instead for inline SPIR-V assembly.
 __attributeTarget(FuncDecl)
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 86c1b556c..563084361 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -955,6 +955,12 @@ class GLSLIndexAttribute : public GLSLSimpleIntegerLayoutAttribute
     SLANG_AST_CLASS(GLSLIndexAttribute)
 };
 
+// [[vk_offset]]
+class VkStructOffsetAttribute : public GLSLSimpleIntegerLayoutAttribute
+{
+    SLANG_AST_CLASS(VkStructOffsetAttribute)
+};
+
 // [[vk_spirv_instruction]]
 class SPIRVInstructionOpAttribute : public Attribute
 {
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index d94c77d6a..e0b203fb6 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -1266,7 +1266,6 @@ AttributeBase* SemanticsVisitor::checkAttribute(
     //
     // The attribute declaration will have one or more `AttributeTargetModifier`s
     // that each specify a syntax class that the attribute can be applied to.
-    // If any of these match `attrTarget`, then we are good.
     //
     bool validTarget = false;
     for (auto attrTargetMod : attrDecl->getModifiersOfType<AttributeTargetModifier>())
@@ -1277,6 +1276,18 @@ AttributeBase* SemanticsVisitor::checkAttribute(
             break;
         }
     }
+
+    // Some attributes impose constraints on where they can be placed that cannot be captured by the
+    // only checking the syntax class. Perform more checks here.
+    switch (attr->astNodeType)
+    {
+    // Allowed only on struct fields.
+    case ASTNodeType::VkStructOffsetAttribute:
+        auto targetDecl = as<Decl>(attrTarget);
+        validTarget = validTarget && targetDecl && as<StructDecl>(getParentDecl(targetDecl));
+        break;
+    };
+
     if (!validTarget)
     {
         getSink()->diagnose(attr, Diagnostics::attributeNotApplicable, attrName);
@@ -1327,6 +1338,7 @@ ASTNodeType getModifierConflictGroupKind(ASTNodeType modifierType)
     case ASTNodeType::GLSLLayoutModifierGroupBegin:
     case ASTNodeType::GLSLLayoutModifierGroupEnd:
     case ASTNodeType::GLSLBufferModifier:
+    case ASTNodeType::VkStructOffsetAttribute:
     case ASTNodeType::MemoryQualifierSetModifier:
     case ASTNodeType::GLSLWriteOnlyModifier:
     case ASTNodeType::GLSLReadOnlyModifier:
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index d337e09a0..18dbcf1cd 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -4413,7 +4413,7 @@ void CLikeSourceEmitter::emitStructDeclarationsBlock(
             }
         }
         emitSemanticsPrefix(fieldKey);
-        emitStructFieldAttributes(structType, ff);
+        emitStructFieldAttributes(structType, ff, allowOffsetLayout);
         emitMemoryQualifiers(fieldKey);
         emitType(fieldType, getName(fieldKey));
         emitSemantics(fieldKey, allowOffsetLayout);
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index ca915ab2d..28dcdcc4d 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -416,8 +416,8 @@ public:
     virtual void emitMemoryQualifiers(IRInst* /*varInst*/){};
     virtual void emitStructFieldAttributes(
         IRStructType* /* structType */,
-        IRStructField* /* field */
-    ){};
+        IRStructField* /* field */,
+        bool /* allowOffsetLayout */){};
     void emitInterpolationModifiers(IRInst* varInst, IRType* valueType, IRVarLayout* layout);
     void emitMeshShaderModifiers(IRInst* varInst);
     virtual void emitPackOffsetModifier(
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 51a22c4e2..848d9c08d 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -234,6 +234,26 @@ void GLSLSourceEmitter::emitMemoryQualifiers(IRInst* varDecl)
     _emitMemoryQualifierDecorations(varDecl);
 }
 
+
+void GLSLSourceEmitter::emitStructFieldAttributes(
+    IRStructType* structType,
+    IRStructField* field,
+    bool allowOffsetLayout)
+{
+    SLANG_UNUSED(structType);
+    auto structKey = field->getKey();
+
+    if (allowOffsetLayout)
+    {
+        if (auto offsetDecoration = structKey->findDecoration<IRVkStructOffsetDecoration>())
+        {
+            m_writer->emit("layout(offset = ");
+            m_writer->emit(offsetDecoration->getOffset()->getValue());
+            m_writer->emit(") ");
+        }
+    }
+}
+
 void GLSLSourceEmitter::_emitGLSLStructuredBuffer(
     IRGlobalParam* varDecl,
     IRHLSLStructuredBufferTypeBase* structuredBufferType)
diff --git a/source/slang/slang-emit-glsl.h b/source/slang/slang-emit-glsl.h
index 8308a9954..c68c08b54 100644
--- a/source/slang/slang-emit-glsl.h
+++ b/source/slang/slang-emit-glsl.h
@@ -53,6 +53,10 @@ protected:
         IRPackOffsetDecoration* decoration) SLANG_OVERRIDE;
 
     virtual void emitMemoryQualifiers(IRInst* varInst) SLANG_OVERRIDE;
+    virtual void emitStructFieldAttributes(
+        IRStructType* structType,
+        IRStructField* field,
+        bool allowOffsetLayout) SLANG_OVERRIDE;
     virtual void emitMeshShaderModifiersImpl(IRInst* varInst) SLANG_OVERRIDE;
     virtual void emitSimpleTypeImpl(IRType* type) SLANG_OVERRIDE;
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount)
diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index 8a3a3cc01..d8a243b32 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -280,8 +280,13 @@ void WGSLSourceEmitter::emitSemanticsPrefixImpl(IRInst* inst)
     }
 }
 
-void WGSLSourceEmitter::emitStructFieldAttributes(IRStructType* structType, IRStructField* field)
+void WGSLSourceEmitter::emitStructFieldAttributes(
+    IRStructType* structType,
+    IRStructField* field,
+    bool allowOffsetLayout)
 {
+    SLANG_UNUSED(allowOffsetLayout);
+
     // Tint emits errors unless we explicitly spell out the layout in some cases, so emit
     // offset and align attribtues for all fields.
     IRSizeAndAlignmentDecoration* const sizeAndAlignmentDecoration =
diff --git a/source/slang/slang-emit-wgsl.h b/source/slang/slang-emit-wgsl.h
index a29f39a1d..551e3b0ba 100644
--- a/source/slang/slang-emit-wgsl.h
+++ b/source/slang/slang-emit-wgsl.h
@@ -42,8 +42,10 @@ public:
     virtual void _emitType(IRType* type, DeclaratorInfo* declarator) SLANG_OVERRIDE;
     virtual void emitFrontMatterImpl(TargetRequest* targetReq) SLANG_OVERRIDE;
     virtual void emitSemanticsPrefixImpl(IRInst* inst) SLANG_OVERRIDE;
-    virtual void emitStructFieldAttributes(IRStructType* structType, IRStructField* field)
-        SLANG_OVERRIDE;
+    virtual void emitStructFieldAttributes(
+        IRStructType* structType,
+        IRStructField* field,
+        bool allowOffsetLayout) SLANG_OVERRIDE;
     virtual void emitCallArg(IRInst* inst) SLANG_OVERRIDE;
     virtual void emitInterpolationModifiersImpl(
         IRInst* varInst,
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 3de40d2c0..8fb1bc9ac 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -1008,6 +1008,7 @@ INST_RANGE(BindingQuery, GetRegisterIndex, GetRegisterSpace)
     INST(GlobalInputDecoration, input, 0, 0)
     INST(GLSLLocationDecoration, glslLocation, 1, 0)
     INST(GLSLOffsetDecoration, glslOffset, 1, 0)
+    INST(VkStructOffsetDecoration, vkStructOffset, 1, 0)
     INST(PayloadDecoration, payload, 0, 0)
     INST(RayPayloadDecoration, raypayload, 0, 0)
 
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 645b13e00..21d5e1c23 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -502,6 +502,12 @@ struct IRGLSLOffsetDecoration : IRDecoration
     IRIntLit* getOffset() { return cast<IRIntLit>(getOperand(0)); }
 };
 
+struct IRVkStructOffsetDecoration : IRDecoration
+{
+    IR_LEAF_ISA(VkStructOffsetDecoration)
+    IRIntLit* getOffset() { return cast<IRIntLit>(getOperand(0)); }
+};
+
 struct IRNVAPIMagicDecoration : IRDecoration
 {
     enum
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index ec5389f56..df643a3c1 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -159,6 +159,12 @@ static Result _calcSizeAndAlignment(
                 // subsequent offsets
                 SLANG_ASSERT(!seenFinalUnsizedArrayField);
 
+                if (auto offsetDecor =
+                        field->getKey()->findDecoration<IRVkStructOffsetDecoration>())
+                {
+                    offset = offsetDecor->getOffset()->getValue();
+                }
+
                 IRSizeAndAlignment fieldTypeLayout;
                 SLANG_RETURN_ON_FAIL(
                     getSizeAndAlignment(optionSet, rules, field->getFieldType(), &fieldTypeLayout));
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 7bc7d4fb4..d764cade5 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -2360,6 +2360,13 @@ void addVarDecorations(IRGenContext* context, IRInst* inst, Decl* decl)
                 kIROp_GLSLOffsetDecoration,
                 builder->getIntValue(builder->getIntType(), glslOffsetMod->offset));
         }
+        else if (auto glslStructOffsetMod = as<VkStructOffsetAttribute>(mod))
+        {
+            builder->addDecoration(
+                inst,
+                kIROp_VkStructOffsetDecoration,
+                builder->getIntValue(builder->getIntType(), glslStructOffsetMod->value));
+        }
         else if (auto hlslSemantic = as<HLSLSimpleSemantic>(mod))
         {
             builder->addSemanticDecoration(inst, hlslSemantic->name.getContent());
diff --git a/tests/diagnostics/vk-offset.slang b/tests/diagnostics/vk-offset.slang
new file mode 100644
index 000000000..5dbd821c6
--- /dev/null
+++ b/tests/diagnostics/vk-offset.slang
@@ -0,0 +1,34 @@
+//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK): -target spirv -emit-spirv-directly -entry vertexMain -stage vertex
+
+// CHECK: error 31002: attribute 'vk_offset' is not valid here
+// CHECK-NEXT: [vk::offset(16)]] struct S1
+// CHECK: error 31002: attribute 'vk_offset' is not valid here
+// CHECK-NEXT: [vk::offset(8)]] VertexOutput output;
+
+[[vk::offset(16)]] struct S1
+{
+    [[vk::offset(32)]]
+    float2 a;
+
+    float3 b;
+
+    [[vk::offset(16)]]
+    float4 c;
+  }
+
+[[vk::push_constant]]
+S1 pc;
+
+struct VertexOutput
+{
+    float3 position : SV_Position;
+}
+
+[shader("vertex")]
+VertexOutput vertexMain()
+{
+    [[vk::offset(8)]] VertexOutput output;
+    output.position = float3(pc.a.x, pc.b.y, pc.c.z);
+    return output;
+}
+
diff --git a/tests/spirv/vk-offset.slang b/tests/spirv/vk-offset.slang
new file mode 100644
index 000000000..64cec1b55
--- /dev/null
+++ b/tests/spirv/vk-offset.slang
@@ -0,0 +1,42 @@
+//TEST:SIMPLE(filecheck=CHECK_SPV): -target spirv -emit-spirv-directly -entry vertexMain -stage vertex
+//TEST:SIMPLE(filecheck=CHECK_SPV): -target spirv -emit-spirv-via-glsl -entry vertexMain -stage vertex
+//TEST:SIMPLE(filecheck=CHECK_GLSL): -target glsl -entry vertexMain -stage vertex
+
+struct S1
+{
+    [[vk::offset(32)]]
+    float2 a;
+
+    float3 b;
+
+    [[vk::offset(16)]]
+    float4 c;
+  }
+
+[[vk::push_constant]]
+S1 pc;
+
+struct VertexOutput
+{
+    float3 position : SV_Position;
+}
+
+// CHECK_SPV: OpDecorate [[STRUCT:%[a-zA-Z0-9_]+]] Block
+// CHECK_SPV-NEXT: OpMemberDecorate [[STRUCT]] 0 Offset 32
+// CHECK_SPV-NEXT: OpMemberDecorate [[STRUCT]] 1 Offset 48
+// CHECK_SPV-NEXT: OpMemberDecorate [[STRUCT]] 2 Offset 16
+
+// CHECK_GLSL: layout(std430) uniform
+// CHECK_GLSL-NEXT: {
+// CHECK_GLSL-NEXT: layout(offset = 32) vec2 a
+// CHECK_GLSL-NEXT: vec3 b
+// CHECK_GLSL-NEXT: layout(offset = 16) vec4 c
+
+[shader("vertex")]
+VertexOutput vertexMain()
+{
+    VertexOutput output;
+    output.position = float3(pc.a.x, pc.b.y, pc.c.z);
+    return output;
+}
+