Implement `-fvk-use-dx-layout` (#4912)

* Implement `-fvk-use-dx-layout`

Fixes: #4126

Changes:
* Added fvk-use-dx-layout
* Modified `HLSLConstantBufferLayoutRulesImpl` for correctness (ex: Array is always 16 byte aligned)
* Added kFXCShaderResourceLayoutRulesFamilyImpl and kFXCConstantBufferLayoutRulesFamilyImpl to handle fvk-use-dx-layout
* Added `ConstantBufferLayoutRules` to manage constant buffer rules
* Added `alignCompositeElementOfNonAggregate`/`alignCompositeElementOfAggregate` to handle forced alignment of composites for ConstantBuffers
* `StructuredBuffer` rules are mostly equal to `scalar` layout, not much was needed to be changed to support this behavior.

* seperate legacy constant buffer and how Slang does constant-buffer normally

* undo an addition

* remove accidental test

* Address review and fix

Address review and remove GLSL support since GLSL requires a seperate legalization (need to linearlize structs like with `legalizeMetalIR` to assign explicit offsets)

* comments

* remove aggregate and non-aggregate logic

We don't need this distinction for the logic

---------

Co-authored-by: Yong He <yonghe@outlook.com>

12 files changed, 122 insertions, 11 deletions

diff --git a/source/slang-record-replay/util/emum-to-string.h b/source/slang-record-replay/util/emum-to-string.h
index f06631719..3c7568818 100644
--- a/source/slang-record-replay/util/emum-to-string.h
+++ b/source/slang-record-replay/util/emum-to-string.h
@@ -175,6 +175,7 @@ namespace SlangRecord
         CASE(VulkanUseGLLayout);
         CASE(VulkanEmitReflection);
         CASE(GLSLForceScalarLayout);
+        CASE(ForceDXLayout);
         CASE(EnableEffectAnnotations);
         CASE(EmitSpirvViaGLSL);
         CASE(EmitSpirvDirectly);
diff --git a/source/slang/slang-compiler-options.cpp b/source/slang/slang-compiler-options.cpp
index ea2593713..3325a313a 100644
--- a/source/slang/slang-compiler-options.cpp
+++ b/source/slang/slang-compiler-options.cpp
@@ -110,6 +110,7 @@ namespace Slang
                     break;
                 case CompilerOptionName::EmitSpirvDirectly:
                 case CompilerOptionName::GLSLForceScalarLayout:
+                case CompilerOptionName::ForceDXLayout:
                 case CompilerOptionName::MatrixLayoutRow:
                 case CompilerOptionName::MatrixLayoutColumn:
                 case CompilerOptionName::VulkanInvertY:
diff --git a/source/slang/slang-compiler-options.h b/source/slang/slang-compiler-options.h
index dae6d49e5..f2bf467e2 100644
--- a/source/slang/slang-compiler-options.h
+++ b/source/slang/slang-compiler-options.h
@@ -346,6 +346,11 @@ namespace Slang
             return getBoolOption(CompilerOptionName::GLSLForceScalarLayout);
         }
 
+        bool shouldUseDXLayout()
+        {
+            return getBoolOption(CompilerOptionName::ForceDXLayout);
+        }
+
         bool shouldDumpIntermediates()
         {
             return getBoolOption(CompilerOptionName::DumpIntermediates);
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index ac85f175e..5ebefb888 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -2807,6 +2807,7 @@ namespace Slang
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetFloatingPointMode(int targetIndex, SlangFloatingPointMode mode) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetMatrixLayoutMode(int targetIndex, SlangMatrixLayoutMode mode) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceGLSLScalarBufferLayout(int targetIndex, bool value) SLANG_OVERRIDE;
+        virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceDXLayout(int targetIndex, bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetGenerateWholeProgram(int targetIndex, bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetEmbedDXIL(int targetIndex, bool value) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setMatrixLayoutMode(SlangMatrixLayoutMode mode) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 20cbbac02..8f4501c2b 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/source/slang/slang-emit-spirv.cpp
@@ -3993,7 +3993,7 @@ struct SPIRVEmitContext
                 auto rule = IRTypeLayoutRules::get(layoutRuleName);
                 IRSizeAndAlignment elementSizeAlignment;
                 getSizeAndAlignment(m_targetProgram->getOptionSet(), rule, matrixType->getElementType(), &elementSizeAlignment);
-
+                IRIntegerValue matrixMinorVectorCount = 0;
                 // Reminder: the meaning of row/column major layout
                 // in our semantics is the *opposite* of what GLSL/SPIRV
                 // calls them, because what they call "columns"
@@ -4007,10 +4007,7 @@ struct SPIRVEmitContext
                         spvStructID,
                         SpvLiteralInteger::from32(id),
                         SpvDecorationRowMajor);
-
-                    auto vectorSize = rule->getVectorSizeAndAlignment(elementSizeAlignment, getIntVal(matrixType->getRowCount()));
-                    vectorSize = rule->alignCompositeElement(vectorSize);
-                    matrixStride = vectorSize.getStride();
+                    matrixMinorVectorCount = getIntVal(matrixType->getRowCount());
                 }
                 else
                 {
@@ -4020,12 +4017,15 @@ struct SPIRVEmitContext
                         spvStructID,
                         SpvLiteralInteger::from32(id),
                         SpvDecorationColMajor);
-
-                    auto vectorSize = rule->getVectorSizeAndAlignment(elementSizeAlignment, getIntVal(matrixType->getColumnCount()));
-                    vectorSize = rule->alignCompositeElement(vectorSize);
-                    matrixStride = vectorSize.getStride();
+                    matrixMinorVectorCount = getIntVal(matrixType->getColumnCount());
                 }
 
+                // We need the size of our vector. To get the stride we need to know how 'big'
+                // each vector element is inside an array, due to this we align our vector
+                // as if a composite.
+                auto vectorSize = rule->getVectorSizeAndAlignment(elementSizeAlignment, matrixMinorVectorCount);
+                vectorSize = rule->alignCompositeElement(vectorSize);
+                matrixStride = vectorSize.getStride();
                 emitOpMemberDecorateMatrixStride(
                     getSection(SpvLogicalSectionID::Annotations),
                     nullptr,
diff --git a/source/slang/slang-ir-layout.cpp b/source/slang/slang-ir-layout.cpp
index 6a4e9360a..3a2471930 100644
--- a/source/slang/slang-ir-layout.cpp
+++ b/source/slang/slang-ir-layout.cpp
@@ -443,6 +443,33 @@ struct NaturalLayoutRules : IRTypeLayoutRules
     }
 };
 
+struct ConstantBufferLayoutRules : IRTypeLayoutRules
+{
+    ConstantBufferLayoutRules()
+    {
+        ruleName = IRTypeLayoutRuleName::D3DConstantBuffer;
+    }
+
+    /// Next member only aligns to 16 if the next member is an array/matrix/struct
+    virtual IRSizeAndAlignment alignCompositeElement(IRSizeAndAlignment currentSize)
+    {
+        // Matrix/Array/Struct should be aligned on a new register
+        return IRSizeAndAlignment(currentSize.size, 16);
+    }
+
+    virtual IRIntegerValue adjustOffsetForNextAggregateMember(IRIntegerValue currentSize, IRIntegerValue lastElementAlignment)
+    {
+        SLANG_UNUSED(lastElementAlignment);
+        return currentSize;
+    }
+
+    virtual IRSizeAndAlignment getVectorSizeAndAlignment(IRSizeAndAlignment element, IRIntegerValue count)
+    {
+        IRIntegerValue countForAlignment = count;
+        return IRSizeAndAlignment((int)(element.size * count), (int)(element.size * countForAlignment));
+    }
+};
+
 struct Std430LayoutRules : IRTypeLayoutRules
 {
     Std430LayoutRules()
@@ -534,6 +561,13 @@ IRTypeLayoutRules* IRTypeLayoutRules::getNatural()
     static NaturalLayoutRules rules;
     return &rules;
 }
+
+IRTypeLayoutRules* IRTypeLayoutRules::getConstantBuffer()
+{
+    static ConstantBufferLayoutRules rules;
+    return &rules;
+}
+
 IRTypeLayoutRules* IRTypeLayoutRules::get(IRTypeLayoutRuleName name)
 {
     switch (name)
@@ -541,6 +575,7 @@ IRTypeLayoutRules* IRTypeLayoutRules::get(IRTypeLayoutRuleName name)
         case IRTypeLayoutRuleName::Std430: return getStd430();
         case IRTypeLayoutRuleName::Std140: return getStd140();
         case IRTypeLayoutRuleName::Natural: return getNatural();
+        case IRTypeLayoutRuleName::D3DConstantBuffer: return getConstantBuffer();
         default: return nullptr;
     }
 }
diff --git a/source/slang/slang-ir-layout.h b/source/slang/slang-ir-layout.h
index 09da2a8f9..c7e5d2b31 100644
--- a/source/slang/slang-ir-layout.h
+++ b/source/slang/slang-ir-layout.h
@@ -58,12 +58,19 @@ struct IRTypeLayoutRules
 {
 public:
     IRTypeLayoutRuleName ruleName;
+
+    /// Align composite based on rule. Type is aligned assuming
+    /// it is apart of a composite (array, struct, matrix, etc...)
     virtual IRSizeAndAlignment alignCompositeElement(IRSizeAndAlignment elementSize) = 0;
+
+    /// Get alignment and size of a vector given components of vector.
+    /// This alignment is not assuming this vector is a member of a struct.
     virtual IRSizeAndAlignment getVectorSizeAndAlignment(IRSizeAndAlignment element, IRIntegerValue count) = 0;
     virtual IRIntegerValue adjustOffsetForNextAggregateMember(IRIntegerValue currentSize, IRIntegerValue lastElementAlignment) = 0;
     static IRTypeLayoutRules* getStd430();
     static IRTypeLayoutRules* getStd140();
     static IRTypeLayoutRules* getNatural();
+    static IRTypeLayoutRules* getConstantBuffer();
     static IRTypeLayoutRules* get(IRTypeLayoutRuleName name);
 };
 
diff --git a/source/slang/slang-ir-lower-buffer-element-type.cpp b/source/slang/slang-ir-lower-buffer-element-type.cpp
index cc24a0e81..a480ae673 100644
--- a/source/slang/slang-ir-lower-buffer-element-type.cpp
+++ b/source/slang/slang-ir-lower-buffer-element-type.cpp
@@ -899,6 +899,16 @@ namespace Slang
         if (target->getOptionSet().shouldUseScalarLayout())
             return IRTypeLayoutRules::getNatural();
 
+        if (target->getOptionSet().shouldUseDXLayout())
+        {
+            if (as<IRUniformParameterGroupType>(bufferType))
+            {
+                return IRTypeLayoutRules::getConstantBuffer();
+            }
+            else
+                return IRTypeLayoutRules::getNatural();
+        }
+
         // The default behavior is to use std140 for constant buffers and std430 for other buffers.
         switch (bufferType->getOp())
         {
diff --git a/source/slang/slang-ir.h b/source/slang/slang-ir.h
index 4a3a04404..036c7f3a7 100644
--- a/source/slang/slang-ir.h
+++ b/source/slang/slang-ir.h
@@ -564,6 +564,7 @@ enum class IRTypeLayoutRuleName
     Scalar = Natural,
     Std430,
     Std140,
+    D3DConstantBuffer,
     _Count,
 };
 
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index 33103442d..12b32998f 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -400,6 +400,7 @@ void initCommandOptions(CommandOptions& options)
         { OptionKind::GLSLForceScalarLayout,
          "-force-glsl-scalar-layout,-fvk-use-scalar-layout", nullptr,
          "Make data accessed through ConstantBuffer, ParameterBlock, StructuredBuffer, ByteAddressBuffer and general pointers follow the 'scalar' layout when targeting GLSL or SPIRV."},
+        { OptionKind::ForceDXLayout, "-fvk-use-dx-layout", nullptr, "Pack members using FXCs member packing rules when targeting GLSL or SPIRV." },
         { OptionKind::VulkanBindShift, vkShiftNames.getBuffer(), "-fvk-<vulkan-shift>-shift <N> <space>", 
         "For example '-fvk-b-shift <N> <space>' shifts by N the inferred binding numbers for all resources in 'b' registers of space <space>. "
         "For a resource attached with :register(bX, <space>) but not [vk::binding(...)], "
@@ -2054,6 +2055,11 @@ SlangResult OptionsParser::_parse(
                 getCurrentTarget()->optionSet.add(CompilerOptionName::GLSLForceScalarLayout, true);
                 break;
             }
+            case OptionKind::ForceDXLayout:
+            {
+                getCurrentTarget()->optionSet.add(CompilerOptionName::ForceDXLayout, true);
+                break;
+            }
             case OptionKind::EnableEffectAnnotations:
             {
                 m_compileRequest->setEnableEffectAnnotations(true);
@@ -2780,6 +2786,11 @@ SlangResult OptionsParser::_parse(
                 m_compileRequest->setTargetForceGLSLScalarBufferLayout(targetID, true);
             }
 
+            if (rawTarget.optionSet.shouldUseDXLayout())
+            {
+                m_compileRequest->setTargetForceDXLayout(targetID, true);
+            }
+
             if (rawTarget.optionSet.getBoolOption(CompilerOptionName::GenerateWholeProgram))
             {
                 m_compileRequest->setTargetGenerateWholeProgram(targetID, true);
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index f5fbfafdf..7d6d047d3 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -400,6 +400,12 @@ struct HLSLConstantBufferLayoutRulesImpl : DefaultLayoutRulesImpl
     }
 };
 
+/// GLSL fvk-use-dx-layout for `ShaderResource`
+struct FXCShaderResourceLayoutRulesImpl : DefaultLayoutRulesImpl
+{
+    // Currently this FXC layout is equal to how we compute 'DefaultLayoutRulesImpl'
+};
+
 /* CPU layout requires that all sizes are a multiple of alignment.
 */
 struct CPULayoutRulesImpl : DefaultLayoutRulesImpl
@@ -894,6 +900,7 @@ struct CUDARayTracingLayoutRulesImpl : DefaultVaryingLayoutRulesImpl
 DefaultLayoutRulesImpl kDefaultLayoutRulesImpl;
 Std140LayoutRulesImpl kStd140LayoutRulesImpl;
 Std430LayoutRulesImpl kStd430LayoutRulesImpl;
+FXCShaderResourceLayoutRulesImpl kFXCShaderResourceLayoutRulesImpl;
 HLSLConstantBufferLayoutRulesImpl kHLSLConstantBufferLayoutRulesImpl;
 HLSLStructuredBufferLayoutRulesImpl kHLSLStructuredBufferLayoutRulesImpl;
 
@@ -1206,6 +1213,18 @@ LayoutRulesImpl kScalarLayoutRulesImpl_ = {
     &kGLSLObjectLayoutRulesImpl,
 };
 
+LayoutRulesImpl kFXCShaderResourceLayoutRulesFamilyImpl = {
+    &kGLSLLayoutRulesFamilyImpl,
+    &kFXCShaderResourceLayoutRulesImpl,
+    &kGLSLObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kFXCConstantBufferLayoutRulesFamilyImpl = {
+    &kGLSLLayoutRulesFamilyImpl,
+    &kHLSLConstantBufferLayoutRulesImpl,
+    &kGLSLObjectLayoutRulesImpl,
+};
+
 LayoutRulesImpl kGLSLAnyValueLayoutRulesImpl_ = {
     &kGLSLLayoutRulesFamilyImpl,
     &kDefaultLayoutRulesImpl,
@@ -1300,6 +1319,9 @@ LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getConstantBufferRules(CompilerOptio
 {
     if (compilerOptions.shouldUseScalarLayout())
         return &kScalarLayoutRulesImpl_;
+    else if (compilerOptions.shouldUseDXLayout())
+        return &kFXCConstantBufferLayoutRulesFamilyImpl;
+
     return &kStd140LayoutRulesImpl_;
 }
 
@@ -1307,6 +1329,9 @@ LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getParameterBlockRules(CompilerOptio
 {
     if (compilerOptions.shouldUseScalarLayout())
         return &kScalarLayoutRulesImpl_;
+    else if (compilerOptions.shouldUseDXLayout())
+        return &kFXCConstantBufferLayoutRulesFamilyImpl;
+
     return &kStd140LayoutRulesImpl_;
 }
 
@@ -1324,6 +1349,9 @@ LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getTextureBufferRules(CompilerOption
 {
     if (compilerOptions.shouldUseScalarLayout())
         return &kScalarLayoutRulesImpl_;
+    else if (compilerOptions.shouldUseDXLayout())
+        return &kFXCConstantBufferLayoutRulesFamilyImpl;
+
     return &kStd430LayoutRulesImpl_;
 
 }
@@ -1347,6 +1375,9 @@ LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getShaderStorageBufferRules(Compiler
 {
     if (compilerOptions.shouldUseScalarLayout())
         return &kScalarLayoutRulesImpl_;
+    else if (compilerOptions.shouldUseDXLayout())
+        return &kFXCShaderResourceLayoutRulesFamilyImpl;
+
     return &kStd430LayoutRulesImpl_;
 }
 
@@ -1365,10 +1396,13 @@ LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getHitAttributesParameterRules()
     return &kGLSLHitAttributesParameterLayoutRulesImpl_;
 }
 
-LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getStructuredBufferRules(CompilerOptionSet& options)
+LayoutRulesImpl* GLSLLayoutRulesFamilyImpl::getStructuredBufferRules(CompilerOptionSet& compilerOptions)
 {
-    if (options.shouldUseScalarLayout())
+    if (compilerOptions.shouldUseScalarLayout())
         return &kScalarLayoutRulesImpl_;
+    else if (compilerOptions.shouldUseDXLayout())
+        return &kFXCShaderResourceLayoutRulesFamilyImpl;
+
     return &kGLSLStructuredBufferLayoutRulesImpl_;
 }
 
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index f08a6fc44..3072ef0a7 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -5854,6 +5854,11 @@ void EndToEndCompileRequest::setTargetForceGLSLScalarBufferLayout(int targetInde
     getTargetOptionSet(targetIndex).set(CompilerOptionName::GLSLForceScalarLayout, value);
 }
 
+void EndToEndCompileRequest::setTargetForceDXLayout(int targetIndex, bool value)
+{
+    getTargetOptionSet(targetIndex).set(CompilerOptionName::ForceDXLayout, value);
+}
+
 void EndToEndCompileRequest::setTargetFloatingPointMode(int targetIndex, SlangFloatingPointMode  mode)
 {
     getTargetOptionSet(targetIndex).set(CompilerOptionName::FloatingPointMode, FloatingPointMode(mode));