[spirv] Handle nested structs for [[vk::offset]] (#1399)

Also update the SPIR-V doc about [[vk::offset]

docs/SPIR-V.rst

@@ -248,6 +248,9 @@ The namespace ``vk`` will be used for all Vulkan attributes:
 - ``push_constant``: For marking a variable as the push constant block. Allowed
   on global variables of struct type. At most one variable can be marked as
   ``push_constant`` in a shader.
+- ``offset(X)``: For manually layout struct members. Annotating a struct member
+  with this attribute will force the compiler to put the member at offset ``X``
+  w.r.t. the beginning of the struct. Only allowed on struct members.
 - ``constant_id(X)``: For marking a global constant as a specialization constant.
   Allowed on global variables of boolean/integer/float types.
 - ``input_attachment_index(X)``: To associate the Xth entry in the input pass
@@ -713,6 +716,21 @@ We will have the following offsets for each member:
 ``g_float2_2``   160    160    176    112    88    128
 ============== ====== ====== ====== ====== ====== ======
 
+If the above layout rules do not satisfy your needs and you want to manually
+control the layout of struct members, you can use either
+
+* The native HLSL ``:packoffset()`` attribute: only available for cbuffers; or
+* The Vulkan-specific ``[[vk::offset()]]`` attribute: applies to all resources.
+
+``[[vk::offset]]`` overrules ``:packoffset``. Attaching ``[[vk::offset]]``
+to a struct memeber affects all variables of the struct type in question. So
+sharing the same struct definition having ``[[vk::offset]]`` annotations means
+also sharing the layout.
+
+These attributes give great flexibility but also responsibility to the
+developer; the compiler will just take in what is specified in the source code
+and emit it to SPIR-V with no error checking.
+
 ``cbuffer`` and ``ConstantBuffer``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

tools/clang/lib/SPIRV/TypeTranslator.cpp

@@ -1305,11 +1305,12 @@ llvm::SmallVector<const Decoration *, 4> TypeTranslator::getLayoutDecorations(
 
     if (rule == LayoutRule::RelaxedGLSLStd140 ||
         rule == LayoutRule::RelaxedGLSLStd430 ||
-        rule == LayoutRule::FxcCTBuffer)
+        rule == LayoutRule::FxcCTBuffer) {
       alignUsingHLSLRelaxedLayout(fieldType, memberSize, &memberAlignment,
                                   &offset);
-    else
+    } else {
       offset = roundToPow2(offset, memberAlignment);
+    }
 
     // The vk::offset attribute takes precedence over all.
     if (const auto *offsetAttr = decl->getAttr<VKOffsetAttr>()) {
@@ -1832,11 +1833,20 @@ TypeTranslator::getAlignmentAndSize(QualType type, LayoutRule rule,
 
       if (rule == LayoutRule::RelaxedGLSLStd140 ||
           rule == LayoutRule::RelaxedGLSLStd430 ||
-          rule == LayoutRule::FxcCTBuffer)
+          rule == LayoutRule::FxcCTBuffer) {
         alignUsingHLSLRelaxedLayout(field->getType(), memberSize,
                                     &memberAlignment, &structSize);
-      else
+      } else {
         structSize = roundToPow2(structSize, memberAlignment);
+      }
+
+      // Reset the current offset to the one specified in the source code
+      // if exists. It's debatable whether we should do sanity check here.
+      // If the developers want manually control the layout, we leave
+      // everything to them.
+      if (const auto *offsetAttr = field->getAttr<VKOffsetAttr>()) {
+        structSize = offsetAttr->getOffset();
+      }
 
       // The base alignment of the structure is N, where N is the largest
       // base alignment value of any of its members...

tools/clang/test/CodeGenSPIRV/vk.layout.attr.offset.hlsl

@@ -0,0 +1,50 @@
+// Run: %dxc -T vs_6_0 -E main -fvk-use-dx-layout
+
+// Make sure that:
+// * [[vk::offset]] on an internal struct affects layout of an external struct.
+// * [[vk::offset]] affects all variables of the struct type containing it.
+// * We follow the normal rules for the fields without [[vk::offset]] specified.
+
+// CHECK: OpMemberDecorate %S 0 Offset 0
+// CHECK: OpMemberDecorate %S 1 Offset 32
+// CHECK: OpMemberDecorate %S 2 Offset 64
+// CHECK: OpMemberDecorate %S 3 Offset 80
+
+// CHECK: OpMemberDecorate %type_ConstantBuffer_T 0 Offset 4
+// CHECK: OpMemberDecorate %type_ConstantBuffer_T 1 Offset 16
+// CHECK: OpMemberDecorate %type_ConstantBuffer_T 2 Offset 116
+
+// CHECK: OpMemberDecorate %S_0 0 Offset 0
+// CHECK: OpMemberDecorate %S_0 1 Offset 32
+// CHECK: OpMemberDecorate %S_0 2 Offset 64
+// CHECK: OpMemberDecorate %S_0 3 Offset 76
+
+// CHECK: OpMemberDecorate %T 0 Offset 4
+// CHECK: OpMemberDecorate %T 1 Offset 8
+// CHECK: OpMemberDecorate %T 2 Offset 92
+
+// CHECK: %type_ConstantBuffer_T = OpTypeStruct %float %S %float
+// CHECK: %T = OpTypeStruct %float %S_0 %float
+
+struct S {
+    float  a;
+    [[vk::offset(32)]]
+    float2 b;
+    [[vk::offset(64)]]
+    float3 c;
+    float  d[2];
+};
+
+struct T {
+    [[vk::offset(4)]]
+    float a;
+    S     s;
+    float b;
+};
+
+ConstantBuffer<T>   MyCBuffer;
+StructuredBuffer<T> MySBuffer;
+
+float4 main() : A {
+    return MyCBuffer.s.a + MySBuffer[0].s.a;
+}

tools/clang/unittests/SPIRV/CodeGenSPIRVTest.cpp

@@ -1378,6 +1378,7 @@ TEST_F(FileTest, VulkanStructuredBufferCounter) {
 
 TEST_F(FileTest, VulkanPushConstant) { runFileTest("vk.push-constant.hlsl"); }
 TEST_F(FileTest, VulkanPushConstantOffset) {
+  // Checks the behavior of [[vk::offset]] with [[vk::push_constant]]
   runFileTest("vk.push-constant.offset.hlsl");
 }
 TEST_F(FileTest, VulkanPushConstantAnonymousStruct) {
@@ -1460,6 +1461,11 @@ TEST_F(FileTest, VulkanLayoutVectorRelaxedLayout) {
   runFileTest("vk.layout.vector.relaxed.hlsl");
 }
 
+TEST_F(FileTest, VulkanLayoutVkOffsetAttr) {
+  // Checks the behavior of [[vk::offset]]
+  runFileTest("vk.layout.attr.offset.hlsl");
+}
+
 TEST_F(FileTest, VulkanLayoutPushConstantStd430) {
   runFileTest("vk.layout.push-constant.std430.hlsl");
 }