Make `-fvk-u-shift` work on AppendStructuredBuffer. (#3144)

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 16 insertions, 0 deletions

diff --git a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
index fa9f16223..eeafd1866 100644
--- a/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
+++ b/source/slang/slang-ir-lower-append-consume-structured-buffer.cpp
@@ -177,13 +177,29 @@ namespace Slang
                         IRBuilder subBuilder(typeUser);
                         IRVarLayout::Builder newVarLayoutBuilder(&subBuilder, typeLayout);
                         newVarLayoutBuilder.cloneEverythingButOffsetsFrom(varLayout);
+                        IRVarOffsetAttr* uavOffsetAttr = nullptr;
+                        IRVarOffsetAttr* descriptorTableSlotOffsetAttr = nullptr;
+
                         for (auto offsetAttr : varLayout->getOffsetAttrs())
                         {
+                            if (offsetAttr->getResourceKind() == LayoutResourceKind::UnorderedAccess)
+                                uavOffsetAttr = offsetAttr;
+                            else if (offsetAttr->getResourceKind() == LayoutResourceKind::DescriptorTableSlot)
+                                descriptorTableSlotOffsetAttr = offsetAttr;
                             auto info = newVarLayoutBuilder.findOrAddResourceInfo(offsetAttr->getResourceKind());
                             info->offset = offsetAttr->getOffset();
                             info->space = offsetAttr->getSpace();
                             info->kind = offsetAttr->getResourceKind();
                         }
+                        // If the user provided an layout offset for UAV but not for descriptor table slot, then
+                        // we use the UAV offset for the descriptor table slot offset.
+                        if (uavOffsetAttr && !descriptorTableSlotOffsetAttr)
+                        {
+                            auto info = newVarLayoutBuilder.findOrAddResourceInfo(LayoutResourceKind::DescriptorTableSlot);
+                            info->offset = uavOffsetAttr->getOffset();
+                            info->space = uavOffsetAttr->getSpace();
+                            info->kind = LayoutResourceKind::DescriptorTableSlot;
+                        }
                         auto newVarLayout = newVarLayoutBuilder.build();
                         subBuilder.addLayoutDecoration(typeUser, newVarLayout);
                         varLayout->removeAndDeallocate();