Add target option to force `scalar` layout for storage buffers. (#2135)

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

9 files changed, 69 insertions, 14 deletions

diff --git a/slang.h b/slang.h
index 1b8c0c8a4..99d7f0d62 100644
--- a/slang.h
+++ b/slang.h
@@ -1365,6 +1365,12 @@ extern "C"
         int targetIndex,
         SlangLineDirectiveMode  mode);
 
+    /*! @see slang::ICompileRequest::setTargetLineDirectiveMode */
+    SLANG_API void spSetTargetForceGLSLScalarBufferLayout(
+        SlangCompileRequest*    request,
+        int targetIndex,
+        bool forceScalarLayout);
+
     /*! @see slang::ICompileRequest::setCodeGenTarget */
     SLANG_API void spSetCodeGenTarget(
         SlangCompileRequest*    request,
@@ -3867,6 +3873,12 @@ namespace slang
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetLineDirectiveMode(
             SlangInt targetIndex,
             SlangLineDirectiveMode mode) = 0;
+
+            /** Set whether to use scalar buffer layouts for GLSL/Vulkan targets.
+                If true, the generated GLSL/Vulkan code will use `scalar` layout for storage buffers.
+                If false, the resulting code will std430 for storage buffers.
+            */
+        virtual SLANG_NO_THROW void SLANG_MCALL setTargetForceGLSLScalarBufferLayout(int targetIndex, bool forceScalarLayout) = 0;
     };
 
     #define SLANG_UUID_ICompileRequest ICompileRequest::getTypeGuid()
@@ -3901,6 +3913,10 @@ namespace slang
             /** The line directive mode for output source code.
             */
         SlangLineDirectiveMode lineDirectiveMode = SLANG_LINE_DIRECTIVE_MODE_DEFAULT;
+
+            /** Whether to force `scalar` layout for glsl shader storage buffers.
+            */
+        bool forceGLSLScalarBufferLayout = false;
     };
 
     typedef uint32_t SessionFlags;
diff --git a/source/slang/slang-api.cpp b/source/slang/slang-api.cpp
index e8e9b602d..ba0b4296a 100644
--- a/source/slang/slang-api.cpp
+++ b/source/slang/slang-api.cpp
@@ -269,6 +269,14 @@ SLANG_API void spSetLineDirectiveMode(
     request->setLineDirectiveMode(mode);
 }
 
+
+SLANG_API void spSetTargetForceGLSLScalarBufferLayout(
+    slang::ICompileRequest* request, int targetIndex, bool forceScalarLayout)
+{
+    SLANG_ASSERT(request);
+    request->setTargetForceGLSLScalarBufferLayout(targetIndex, forceScalarLayout);
+}
+
 SLANG_API void spSetTargetLineDirectiveMode(
     slang::ICompileRequest* request,
     int targetIndex,
diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index a97ec54d8..d73a03b42 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1461,6 +1461,11 @@ namespace Slang
         {
             dumpIntermediates = value;
         }
+        void setForceGLSLScalarBufferLayout(bool value)
+        {
+            forceGLSLScalarBufferLayout = value;
+        }
+
         void addCapability(CapabilityAtom capability);
 
         bool shouldEmitSPIRVDirectly()
@@ -1482,6 +1487,7 @@ namespace Slang
         LineDirectiveMode getLineDirectiveMode() { return lineDirectiveMode; }
         SlangTargetFlags getTargetFlags() { return targetFlags; }
         CapabilitySet getTargetCaps();
+        bool getForceGLSLScalarBufferLayout() { return forceGLSLScalarBufferLayout; }
 
         Session* getSession();
         MatrixLayoutMode getDefaultMatrixLayoutMode();
@@ -1503,6 +1509,7 @@ namespace Slang
         CapabilitySet           cookedCapabilities;
         LineDirectiveMode       lineDirectiveMode = LineDirectiveMode::Default;
         bool                    dumpIntermediates = false;
+        bool                    forceGLSLScalarBufferLayout = false;
     };
 
         /// Are we generating code for a D3D API?
@@ -2243,6 +2250,7 @@ namespace Slang
         virtual SLANG_NO_THROW void SLANG_MCALL setTargetFlags(int targetIndex, SlangTargetFlags flags) SLANG_OVERRIDE;
         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 setMatrixLayoutMode(SlangMatrixLayoutMode mode) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setDebugInfoLevel(SlangDebugInfoLevel level) SLANG_OVERRIDE;
         virtual SLANG_NO_THROW void SLANG_MCALL setOptimizationLevel(SlangOptimizationLevel level) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index c0106e9ad..03943abb3 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -98,6 +98,7 @@ CLikeSourceEmitter::CLikeSourceEmitter(const Desc& desc)
     m_targetCaps = desc.targetCaps;
 
     m_compileRequest = desc.compileRequest;
+    m_targetRequest = desc.targetRequest;
     m_entryPointStage = desc.entryPointStage;
     m_effectiveProfile = desc.effectiveProfile;
 }
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index ed0f69a2c..aa5a22697 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -41,6 +41,8 @@ public:
         ExtensionTracker* extensionTracker = nullptr;
 
         SourceWriter* sourceWriter = nullptr;
+
+        TargetRequest* targetRequest = nullptr;
     };
 
     enum
@@ -475,6 +477,7 @@ public:
     List<IRWitnessTableEntry*> getSortedWitnessTableEntries(IRWitnessTable* witnessTable);
     
     BackEndCompileRequest* m_compileRequest = nullptr;
+    TargetRequest* m_targetRequest = nullptr;
     IRModule* m_irModule = nullptr;
 
     // The stage for which we are emitting code.
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index d5175a423..441fc94d7 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -40,6 +40,11 @@ SlangResult GLSLSourceEmitter::init()
         default: break;
     }
 
+    if (m_targetRequest->getForceGLSLScalarBufferLayout())
+    {
+        m_glslExtensionTracker->requireExtension(
+            UnownedStringSlice::fromLiteral("GL_EXT_scalar_block_layout"));
+    }
     return SLANG_OK;
 }
 
@@ -115,7 +120,8 @@ void GLSLSourceEmitter::_emitGLSLStructuredBuffer(IRGlobalParam* varDecl, IRHLSL
     // TODO: we should require either the extension or the version...
     _requireGLSLVersion(430);
 
-    m_writer->emit("layout(std430");
+    m_writer->emit("layout(");
+    m_writer->emit(m_targetRequest->getForceGLSLScalarBufferLayout() ? "scalar" : "std430");
 
     auto layout = getVarLayout(varDecl);
     if (layout)
@@ -189,7 +195,8 @@ void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByt
     // TODO: we should require either the extension or the version...
     _requireGLSLVersion(430);
 
-    m_writer->emit("layout(std430");
+    m_writer->emit("layout(");
+    m_writer->emit(m_targetRequest->getForceGLSLScalarBufferLayout() ? "scalar" : "std430");
 
     auto layout = getVarLayout(varDecl);
     if (layout)
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 780fa1155..378732fb3 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -774,6 +774,7 @@ SlangResult emitEntryPointsSourceFromIR(
     CLikeSourceEmitter::Desc desc;
 
     desc.compileRequest = compileRequest;
+    desc.targetRequest = targetRequest;
     desc.target = target;
     // TODO(DG): Can't assume a single entry point stage for multiple entry points
     if (entryPointIndices.getCount() == 1)
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 3880d1114..c1834eacb 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -899,6 +899,7 @@ void Linkage::addTarget(
     target->addTargetFlags(desc.flags);
     target->setTargetProfile(Profile(desc.profile));
     target->setLineDirectiveMode(LineDirectiveMode(desc.lineDirectiveMode));
+    target->setForceGLSLScalarBufferLayout(desc.forceGLSLScalarBufferLayout);
 }
 
 #if 0
@@ -4187,6 +4188,11 @@ void EndToEndCompileRequest::setTargetFlags(int targetIndex, SlangTargetFlags fl
     getLinkage()->targets[targetIndex]->addTargetFlags(flags);
 }
 
+void EndToEndCompileRequest::setTargetForceGLSLScalarBufferLayout(int targetIndex, bool value)
+{
+    getLinkage()->targets[targetIndex]->setForceGLSLScalarBufferLayout(value);
+}
+
 void EndToEndCompileRequest::setTargetFloatingPointMode(int targetIndex, SlangFloatingPointMode  mode)
 {
     getLinkage()->targets[targetIndex]->setFloatingPointMode(FloatingPointMode(mode));
@@ -4848,5 +4854,3 @@ SlangResult EndToEndCompileRequest::getEntryPoint(SlangInt entryPointIndex, slan
 }
 
 } // namespace Slang
-
-
diff --git a/tools/gfx/d3d12/render-d3d12.cpp b/tools/gfx/d3d12/render-d3d12.cpp
index 8457889df..12f8bafa6 100644
--- a/tools/gfx/d3d12/render-d3d12.cpp
+++ b/tools/gfx/d3d12/render-d3d12.cpp
@@ -2570,8 +2570,8 @@ public:
 
             m_layout = layout;
 
-            m_constantBufferTransientHeap = nullptr;
-            m_constantBufferTransientHeapVersion = 0;
+            m_cachedTransientHeap = nullptr;
+            m_cachedTransientHeapVersion = 0;
             m_isConstantBufferDirty = true;
 
             // If the layout tells us that there is any uniform data,
@@ -2763,8 +2763,8 @@ public:
 
         bool shouldAllocateConstantBuffer(TransientResourceHeapImpl* transientHeap)
         {
-            return m_isConstantBufferDirty || m_constantBufferTransientHeap != transientHeap ||
-                m_constantBufferTransientHeapVersion != transientHeap->getVersion();
+            return m_isConstantBufferDirty || m_cachedTransientHeap != transientHeap ||
+                m_cachedTransientHeapVersion != transientHeap->getVersion();
         }
 
         /// Ensure that the `m_ordinaryDataBuffer` has been created, if it is needed
@@ -2780,8 +2780,8 @@ public:
                 return SLANG_OK;
             }
             m_isConstantBufferDirty = false;
-            m_constantBufferTransientHeap = encoder->m_transientHeap;
-            m_constantBufferTransientHeapVersion = encoder->m_transientHeap->getVersion();
+            m_cachedTransientHeap = encoder->m_transientHeap;
+            m_cachedTransientHeapVersion = encoder->m_transientHeap->getVersion();
 
             // Computing the size of the ordinary data buffer is *not* just as simple
             // as using the size of the `m_ordinayData` array that we store. The reason
@@ -2930,6 +2930,11 @@ public:
             BindingOffset const&    offset,
             ShaderObjectLayoutImpl* specializedLayout)
         {
+            if (m_cachedTransientHeap == context->transientHeap &&
+                m_cachedTransientHeapVersion == m_cachedTransientHeap->getVersion())
+            {
+
+            }
             // The first step to binding an object as a parameter block is to allocate a descriptor
             // set (consisting of zero or one resource descriptor table and zero or one sampler
             // descriptor table) to represent its values.
@@ -3142,6 +3147,8 @@ public:
         }
             /// A CPU-memory descriptor set holding any descriptors used to represent the resources/samplers in this object's state
         DescriptorSet m_descriptorSet;
+            /// A cached descriptor set on GPU heap.
+        DescriptorSet m_cachedGPUDescriptorSet;
 
         ShortList<RefPtr<Resource>, 8> m_boundResources;
         List<D3D12_GPU_VIRTUAL_ADDRESS> m_rootArguments;
@@ -3155,10 +3162,10 @@ public:
 
         /// Dirty bit tracking whether the constant buffer needs to be updated.
         bool m_isConstantBufferDirty = true;
-        /// The transient heap from which the constant buffer is allocated.
-        TransientResourceHeapImpl* m_constantBufferTransientHeap;
-        /// The version of the transient heap when the constant buffer is allocated.
-        uint64_t m_constantBufferTransientHeapVersion;
+        /// The transient heap from which the constant buffer and descriptor set is allocated.
+        TransientResourceHeapImpl* m_cachedTransientHeap;
+        /// The version of the transient heap when the constant buffer and descriptor set is allocated.
+        uint64_t m_cachedTransientHeapVersion;
 
         /// Get the layout of this shader object with specialization arguments considered
         ///