Support `-fvk-use-gl-layout` for ByteAddressBuffer load/store. (#3068)

* Support `-fvk-use-gl-layout` for ByteAddressBuffer load/store.

* Fix.

* Fix.

* Add test for unaligned load.

---------

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

10 files changed, 400 insertions, 4 deletions

diff --git a/source/slang/slang-hlsl-to-vulkan-layout-options.cpp b/source/slang/slang-hlsl-to-vulkan-layout-options.cpp
index 1a56be3d9..d15b60f3b 100644
--- a/source/slang/slang-hlsl-to-vulkan-layout-options.cpp
+++ b/source/slang/slang-hlsl-to-vulkan-layout-options.cpp
@@ -91,7 +91,8 @@ Index HLSLToVulkanLayoutOptions::getShift(Kind kind, Index set) const
 
 bool HLSLToVulkanLayoutOptions::hasState() const
 {
-    return canInferBindings() || hasGlobalsBinding() || shouldInvertY() || getUseOriginalEntryPointName();
+    return canInferBindings() || hasGlobalsBinding() || shouldInvertY() || getUseOriginalEntryPointName()
+        || shouldUseGLLayout();
 }
 
 HLSLToVulkanLayoutOptions::Binding HLSLToVulkanLayoutOptions::inferBinding(Kind kind, const Binding& inBinding) const
diff --git a/source/slang/slang-hlsl-to-vulkan-layout-options.h b/source/slang/slang-hlsl-to-vulkan-layout-options.h
index 2d24dfccb..c88805412 100644
--- a/source/slang/slang-hlsl-to-vulkan-layout-options.h
+++ b/source/slang/slang-hlsl-to-vulkan-layout-options.h
@@ -119,6 +119,8 @@ public:
         /// True if the compiler should invert the Y coordinate of any SV_Position output.
     bool shouldInvertY() const { return m_invertY; }
 
+    bool shouldUseGLLayout() const { return m_useGLLayout; }
+
     bool getUseOriginalEntryPointName() const { return m_useOriginalEntryPointName; }
 
         /// Given an kind and a binding infer the vulkan binding.
@@ -149,6 +151,8 @@ public:
 
     void setUseOriginalEntryPointName(bool value) { m_useOriginalEntryPointName = value; }
 
+    void setUseGLLayout(bool value) { m_useGLLayout = value; }
+
         /// Ctor
     HLSLToVulkanLayoutOptions();
     
@@ -178,6 +182,9 @@ protected:
 
         /// If set, will use the original entry point name in the generated SPIRV instead of "main".
     bool m_useOriginalEntryPointName = false;
+
+        /// If set, raw buffer load/stores will follow std430 layout.
+    bool m_useGLLayout = false;
 };
 
 } // namespace Slang
diff --git a/source/slang/slang-ir-byte-address-legalize.cpp b/source/slang/slang-ir-byte-address-legalize.cpp
index 2d53dcab7..14f985c3b 100644
--- a/source/slang/slang-ir-byte-address-legalize.cpp
+++ b/source/slang/slang-ir-byte-address-legalize.cpp
@@ -203,6 +203,15 @@ struct ByteAddressBufferLegalizationContext
         return false;
     }
 
+    SlangResult getOffset(TargetRequest* target, IRStructField* field, IRIntegerValue* outOffset)
+    {
+        if (target->getHLSLToVulkanLayoutOptions() && target->getHLSLToVulkanLayoutOptions()->shouldUseGLLayout())
+        {
+            return getStd430Offset(target, field, outOffset);
+        }
+        return getNaturalOffset(target, field, outOffset);
+    }
+
     // The core workhorse routine for the load case is `emitLegalLoad`,
     // which tries to emit load operations that read a value of the
     // given `type` from the given `buffer` at the required `baseOffset`
@@ -256,7 +265,7 @@ struct ByteAddressBufferLegalizationContext
                 // then we fail to legalize this load.
                 //
                 IRIntegerValue fieldOffset = 0;
-                SLANG_RETURN_NULL_ON_FAIL(getNaturalOffset(m_target, field, &fieldOffset));
+                SLANG_RETURN_NULL_ON_FAIL(getOffset(m_target, field, &fieldOffset));
 
                 // Otherwise, we load the field by recursively calling this function
                 // on the field type, with an adjusted immediate offset.
@@ -803,7 +812,7 @@ struct ByteAddressBufferLegalizationContext
                 auto fieldType = field->getFieldType();
 
                 IRIntegerValue fieldOffset;
-                SLANG_RETURN_ON_FAIL(getNaturalOffset(m_target, field, &fieldOffset));
+                SLANG_RETURN_ON_FAIL(getOffset(m_target, field, &fieldOffset));
 
                 auto fieldVal = m_builder.emitFieldExtract(fieldType, value, field->getKey());
                 SLANG_RETURN_ON_FAIL(emitLegalStore(fieldType, buffer, baseOffset, immediateOffset + fieldOffset, fieldVal));
@@ -874,7 +883,6 @@ struct ByteAddressBufferLegalizationContext
     Result emitSimpleStore(IRType* type, IRInst* buffer, IRInst* baseOffset, IRIntegerValue immediateOfset, IRInst* value)
     {
         IRInst* offset = emitOffsetAddIfNeeded(baseOffset, immediateOfset);
-
         if( m_options.translateToStructuredBufferOps )
         {
             if( auto structuredBuffer = getEquivalentStructuredBuffer(type, buffer) )
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 69f3c4e0d..c26d5fe53 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -731,6 +731,12 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
         /// A `[naturalOffset(o)]` decoration is attached to a field to indicate that it has natural offset `o` in the parent type
     INST(NaturalOffsetDecoration, naturalOffset, 1, 0)
 
+        /// A `[std430SizeAndAlignment(s,a)]` decoration is attached to a type to indicate that is has std430 size `s` and alignment `a`
+    INST(Std430SizeAndAlignmentDecoration, naturalSizeAndAlignment, 2, 0)
+
+        /// A `[std430Offset(o)]` decoration is attached to a field to indicate that it has std430 offset `o` in the parent type
+    INST(Std430OffsetDecoration, naturalOffset, 1, 0)
+
     /* LinkageDecoration */
         INST(ImportDecoration, import, 1, 0)
         INST(ExportDecoration, export, 1, 0)
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 123cc33c6..2d8393698 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -595,6 +595,28 @@ struct IRNaturalOffsetDecoration : IRDecoration
     IRIntegerValue getOffset() { return getOffsetOperand()->getValue(); }
 };
 
+struct IRStd430SizeAndAlignmentDecoration : IRDecoration
+{
+    enum { kOp = kIROp_Std430SizeAndAlignmentDecoration };
+    IR_LEAF_ISA(Std430SizeAndAlignmentDecoration)
+
+    IRIntLit* getSizeOperand() { return cast<IRIntLit>(getOperand(0)); }
+    IRIntLit* getAlignmentOperand() { return cast<IRIntLit>(getOperand(1)); }
+
+    IRIntegerValue getSize() { return getSizeOperand()->getValue(); }
+    IRIntegerValue getAlignment() { return getAlignmentOperand()->getValue(); }
+};
+
+struct IRStd430OffsetDecoration : IRDecoration
+{
+    enum { kOp = kIROp_Std430OffsetDecoration };
+    IR_LEAF_ISA(Std430OffsetDecoration)
+
+    IRIntLit* getOffsetOperand() { return cast<IRIntLit>(getOperand(0)); }
+
+    IRIntegerValue getOffset() { return getOffsetOperand()->getValue(); }
+};
+
 struct IRBuiltinDecoration : IRDecoration
 {
     enum
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index ab24d3354..df2f50b30 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -352,4 +352,308 @@ Result getNaturalOffset(TargetRequest* target, IRStructField* field, IRIntegerVa
     return SLANG_FAIL;
 }
 
+
+//////////////////////////
+// Std430 Layout
+//////////////////////////
+
+static Result _calcStd430ArraySizeAndAlignment(
+    TargetRequest* target,
+    IRType* elementType,
+    IRInst* elementCountInst,
+    IRSizeAndAlignment* outSizeAndAlignment)
+{
+    auto elementCountLit = as<IRIntLit>(elementCountInst);
+    if (!elementCountLit)
+        return SLANG_FAIL;
+    auto elementCount = elementCountLit->getValue();
+
+    if (elementCount == 0)
+    {
+        *outSizeAndAlignment = IRSizeAndAlignment(0, 1);
+        return SLANG_OK;
+    }
+
+    IRSizeAndAlignment elementTypeLayout;
+    SLANG_RETURN_ON_FAIL(getStd430SizeAndAlignment(target, elementType, &elementTypeLayout));
+
+    auto elementStride = elementTypeLayout.getStride();
+
+    *outSizeAndAlignment = IRSizeAndAlignment(
+        elementStride * (elementCount - 1) + elementTypeLayout.size,
+        elementTypeLayout.alignment);
+    return SLANG_OK;
+}
+
+static Result _calcStd430SizeAndAlignment(
+    TargetRequest* target,
+    IRType* type,
+    IRSizeAndAlignment* outSizeAndAlignment)
+{
+    switch (type->getOp())
+    {
+
+#define CASE(TYPE, SIZE, ALIGNMENT)                                 \
+    case kIROp_##TYPE##Type:                                        \
+        *outSizeAndAlignment = IRSizeAndAlignment(SIZE, ALIGNMENT); \
+        return SLANG_OK                                             \
+        /* end */
+
+        // Most base types are "std430 aligned" (meaning alignment and size are the same)
+#define BASE(TYPE, SIZE) CASE(TYPE, SIZE, SIZE)
+
+        BASE(Int8, 1);
+        BASE(UInt8, 1);
+
+        BASE(Int16, 2);
+        BASE(UInt16, 2);
+        BASE(Half, 2);
+
+        BASE(Int, 4);
+        BASE(UInt, 4);
+        BASE(Float, 4);
+
+        BASE(Int64, 8);
+        BASE(UInt64, 8);
+        BASE(Double, 8);
+
+        // We are currently handling `bool` following the HLSL
+        // precednet of storing it in 4 bytes.
+        //
+        // TODO: It would be good to try to make this follow
+        // per-platform conventions, or at least to be able
+        // to use a 1-byte encoding where available.
+        //
+        BASE(Bool, 4);
+
+        // The Slang `void` type is treated as a zero-byte
+        // type, so that it does not influence layout at all.
+        //
+        CASE(Void, 0, 1);
+
+#undef CASE
+
+#undef CASE
+
+    case kIROp_StructType:
+    {
+        auto structType = cast<IRStructType>(type);
+        IRSizeAndAlignment structLayout;
+        for (auto field : structType->getFields())
+        {
+            IRSizeAndAlignment fieldTypeLayout;
+            SLANG_RETURN_ON_FAIL(getStd430SizeAndAlignment(target, field->getFieldType(), &fieldTypeLayout));
+
+            structLayout.size = align(structLayout.size, fieldTypeLayout.alignment);
+            structLayout.alignment = std::max(structLayout.alignment, fieldTypeLayout.alignment);
+
+            IRIntegerValue fieldOffset = structLayout.size;
+            if (auto module = type->getModule())
+            {
+                // If we are in a situation where attaching new
+                // decorations is possible, then we want to
+                // cache the field offset on the IR field
+                // instruction.
+                //
+                IRBuilder builder(module);
+
+                auto intType = builder.getIntType();
+                builder.addDecoration(
+                    field,
+                    kIROp_Std430OffsetDecoration,
+                    builder.getIntValue(intType, fieldOffset));
+            }
+
+            structLayout.size += fieldTypeLayout.size;
+        }
+        *outSizeAndAlignment = structLayout;
+        return SLANG_OK;
+    }
+    break;
+
+    case kIROp_ArrayType:
+    {
+        auto arrayType = cast<IRArrayType>(type);
+
+        return _calcStd430ArraySizeAndAlignment(
+            target,
+            arrayType->getElementType(),
+            arrayType->getElementCount(),
+            outSizeAndAlignment);
+    }
+    break;
+
+    case kIROp_VectorType:
+    {
+        auto vecType = cast<IRVectorType>(type);
+        auto elementCount = getIntegerValueFromInst(vecType->getElementCount());
+        auto alignmentMultiplier = elementCount;
+        if (elementCount == 3)
+            alignmentMultiplier = 4;
+        IRSizeAndAlignment sizeAndAlignment;
+        SLANG_RETURN_ON_FAIL(getStd430SizeAndAlignment(target, vecType->getElementType(), &sizeAndAlignment));
+        sizeAndAlignment.size *= (int)elementCount;
+        sizeAndAlignment.alignment *= (int)alignmentMultiplier;
+        *outSizeAndAlignment = sizeAndAlignment;
+        return SLANG_OK;
+    }
+    break;
+    case kIROp_AnyValueType:
+    {
+        auto anyValType = cast<IRAnyValueType>(type);
+        outSizeAndAlignment->size = getIntVal(anyValType->getSize());
+        outSizeAndAlignment->alignment = 4;
+        return SLANG_OK;
+    }
+    break;
+    case kIROp_TupleType:
+    {
+        auto tupleType = cast<IRTupleType>(type);
+        IRSizeAndAlignment resultLayout;
+        for (UInt i = 0; i < tupleType->getOperandCount(); i++)
+        {
+            auto elementType = tupleType->getOperand(i);
+            IRSizeAndAlignment fieldTypeLayout;
+            SLANG_RETURN_ON_FAIL(getStd430SizeAndAlignment(target, (IRType*)elementType, &fieldTypeLayout));
+            resultLayout.size = align(resultLayout.size, fieldTypeLayout.alignment);
+            resultLayout.alignment = std::max(resultLayout.alignment, fieldTypeLayout.alignment);
+        }
+        *outSizeAndAlignment = resultLayout;
+        return SLANG_OK;
+    }
+    break;
+    case kIROp_WitnessTableType:
+    case kIROp_WitnessTableIDType:
+    case kIROp_RTTIHandleType:
+    {
+        outSizeAndAlignment->size = kRTTIHandleSize;
+        outSizeAndAlignment->alignment = 4;
+        return SLANG_OK;
+    }
+    break;
+    case kIROp_InterfaceType:
+    {
+        auto interfaceType = cast<IRInterfaceType>(type);
+        auto size = SharedGenericsLoweringContext::getInterfaceAnyValueSize(interfaceType, interfaceType->sourceLoc);
+        size += kRTTIHeaderSize;
+        size = align(size, 4);
+        IRSizeAndAlignment resultLayout;
+        resultLayout.size = size;
+        resultLayout.alignment = 4;
+        *outSizeAndAlignment = resultLayout;
+        return SLANG_OK;
+    }
+    break;
+    case kIROp_MatrixType:
+    {
+        auto matType = cast<IRMatrixType>(type);
+        IRBuilder builder(type->getModule());
+        builder.setInsertBefore(matType);
+        auto rowType = builder.getVectorType(matType->getElementType(), matType->getColumnCount());
+        return _calcStd430ArraySizeAndAlignment(
+            target, rowType,
+            matType->getRowCount(),
+            outSizeAndAlignment);
+    }
+    break;
+    case kIROp_OutType:
+    case kIROp_InOutType:
+    case kIROp_RefType:
+    case kIROp_RawPointerType:
+    case kIROp_PtrType:
+    case kIROp_NativePtrType:
+    case kIROp_ComPtrType:
+    case kIROp_NativeStringType:
+    {
+        *outSizeAndAlignment = IRSizeAndAlignment(sizeof(void*), sizeof(void*));
+        return SLANG_OK;
+    }
+    break;
+    default:
+        break;
+    }
+
+    if (areResourceTypesBindlessOnTarget(target))
+    {
+        // TODO: need this to be based on target, instead of hard-coded
+        int pointerSize = sizeof(void*);
+
+        if (as<IRTextureType>(type))
+        {
+            *outSizeAndAlignment = IRSizeAndAlignment(pointerSize, pointerSize);
+            return SLANG_OK;
+        }
+        else if (as<IRSamplerStateTypeBase>(type))
+        {
+            *outSizeAndAlignment = IRSizeAndAlignment(pointerSize, pointerSize);
+            return SLANG_OK;
+        }
+        // TODO: the remaining cases for "bindless" resources on CPU/CUDA targets
+    }
+
+    return SLANG_FAIL;
+}
+
+Result getStd430SizeAndAlignment(TargetRequest* target, IRType* type, IRSizeAndAlignment* outSizeAndAlignment)
+{
+    if (auto decor = type->findDecoration<IRStd430SizeAndAlignmentDecoration>())
+    {
+        *outSizeAndAlignment = IRSizeAndAlignment(decor->getSize(), (int)decor->getAlignment());
+        return SLANG_OK;
+    }
+
+    IRSizeAndAlignment sizeAndAlignment;
+    SLANG_RETURN_ON_FAIL(_calcStd430SizeAndAlignment(target, type, &sizeAndAlignment));
+
+    if (auto module = type->getModule())
+    {
+        IRBuilder builder(module);
+
+        auto intType = builder.getIntType();
+        builder.addDecoration(
+            type,
+            kIROp_Std430SizeAndAlignmentDecoration,
+            builder.getIntValue(intType, sizeAndAlignment.size),
+            builder.getIntValue(intType, sizeAndAlignment.alignment));
+    }
+
+    *outSizeAndAlignment = sizeAndAlignment;
+    return SLANG_OK;
+}
+
+
+Result getStd430Offset(TargetRequest* target, IRStructField* field, IRIntegerValue* outOffset)
+{
+    if (auto decor = field->findDecoration<IRStd430OffsetDecoration>())
+    {
+        *outOffset = decor->getOffset();
+        return SLANG_OK;
+    }
+
+    // Offsets are computed as part of layout out types,
+    // so we expect that layout of the "parent" type
+    // of the field should add an offset to it if
+    // possible.
+
+    auto structType = as<IRStructType>(field->getParent());
+    if (!structType)
+        return SLANG_FAIL;
+
+    IRSizeAndAlignment structTypeLayout;
+    SLANG_RETURN_ON_FAIL(getStd430SizeAndAlignment(target, structType, &structTypeLayout));
+
+    if (auto decor = field->findDecoration<IRStd430OffsetDecoration>())
+    {
+        *outOffset = decor->getOffset();
+        return SLANG_OK;
+    }
+
+    // If attempting to lay out the parent type didn't
+    // cause the field to get an offset, then we are
+    // in an unexpected case with no easy answer.
+    //
+    return SLANG_FAIL;
+}
+
+
 }
diff --git a/source/slang/slang-ir-layout.h b/source/slang/slang-ir-layout.h
index 5c736f474..d4e4c570a 100644
--- a/source/slang/slang-ir-layout.h
+++ b/source/slang/slang-ir-layout.h
@@ -67,5 +67,21 @@ Result getNaturalSizeAndAlignment(TargetRequest* target, IRType* type, IRSizeAnd
     ///
 Result getNaturalOffset(TargetRequest* target, IRStructField* field, IRIntegerValue* outOffset);
 
+/// Compute (if necessary) and return the std430 size and alignment of `type`.
+///
+/// This operation may fail if `type` is not one that can be stored in
+/// general-purpose memory for the current target. In that case the
+/// type is considered to have no std430 layout.
+///
+Result getStd430SizeAndAlignment(TargetRequest* target, IRType* type, IRSizeAndAlignment* outSizeAndAlignment);
+
+/// Compute (if necessary) and return the std430 offset of `field`
+///
+/// This operation can fail if the parent type of `field` is not one
+/// that can be stored in general-purpose memory. In that case, the
+/// field is considered to have no std430 offset.
+///
+Result getStd430Offset(TargetRequest* target, IRStructField* field, IRIntegerValue* outOffset);
+
 }
 
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index b31a502e7..dbd6b97a4 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -94,6 +94,7 @@ enum class OptionKind
     VulkanBindGlobals,
     VulkanInvertY,
     VulkanUseEntryPointName,
+    VulkanUseGLLayout,
 
 
     GLSLForceScalarLayout,
@@ -501,6 +502,7 @@ void initCommandOptions(CommandOptions& options)
         "Places the $Globals cbuffer at descriptor set <descriptor-set> and binding <N>."},
         { OptionKind::VulkanInvertY, "-fvk-invert-y", nullptr, "Negates (additively inverts) SV_Position.y before writing to stage output."},
         { OptionKind::VulkanUseEntryPointName, "-fvk-use-entrypoint-name", nullptr, "Uses the entrypoint name from the source instead of 'main' in the spirv output."},
+        { OptionKind::VulkanUseGLLayout, "-fvk-use-gl-layout", nullptr, "Use std430 layout instead of D3D buffer layout for raw buffer load/stores."},
         { OptionKind::EnableEffectAnnotations,
          "-enable-effect-annotations", nullptr, 
          "Enables support for legacy effect annotation syntax."},
@@ -2025,6 +2027,12 @@ SlangResult OptionsParser::_parse(
                 m_hlslToVulkanLayoutOptions->setUseOriginalEntryPointName(true);
                 break;
             }
+            case OptionKind::VulkanUseGLLayout:
+            {
+                // -fvk-use-gl-layout
+                m_hlslToVulkanLayoutOptions->setUseGLLayout(true);
+                break;
+            }
             case OptionKind::Profile: SLANG_RETURN_ON_FAIL(_parseProfile(arg)); break;
             case OptionKind::Capability:
             {
diff --git a/tests/hlsl/byte-buffer-load-std430.slang b/tests/hlsl/byte-buffer-load-std430.slang
new file mode 100644
index 000000000..b526250eb
--- /dev/null
+++ b/tests/hlsl/byte-buffer-load-std430.slang
@@ -0,0 +1,22 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type -xslang -fvk-use-gl-layout
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+//TEST_INPUT:set inBuffer = ubuffer(data=[1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0], stride=4)
+RWByteAddressBuffer inBuffer;
+
+struct MyStruct
+{
+    float v0;
+    float2 v1;
+    float2 v2;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    var v = inBuffer.Load<MyStruct>(0);
+    var g = inBuffer.Load<float4>(4);
+    outputBuffer[dispatchThreadID.x] = v.v2.x + v.v1.x + g.x; // expect 10.0
+}
diff --git a/tests/hlsl/byte-buffer-load-std430.slang.expected.txt b/tests/hlsl/byte-buffer-load-std430.slang.expected.txt
new file mode 100644
index 000000000..0c174890f
--- /dev/null
+++ b/tests/hlsl/byte-buffer-load-std430.slang.expected.txt
@@ -0,0 +1,2 @@
+type: float
+10.000000