From 3f04389a1806c74355236ecd0a2493d701b19d87 Mon Sep 17 00:00:00 2001
From: Malcolm Bechard <mbechard@users.noreply.github.com>
Date: Wed, 22 Sep 2021 13:41:10 -0400
Subject: [PATCH] Auto push constant blocks (#2764)

* add ability to upgrade uniform block to push constants assuming it fits within the size limit imposed by the caller
* allow selecting the packing for the auto push constants
* check the size using the potential layout packing of the push constants
---
 glslang/Include/Types.h                     | 10 ++++++++++
 glslang/MachineIndependent/iomapper.cpp     | 31 +++++++++++++++++++++++++++++
 glslang/MachineIndependent/iomapper.h       | 16 ++++++++++++++-
 glslang/MachineIndependent/linkValidate.cpp |  2 +-
 4 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/glslang/Include/Types.h b/glslang/Include/Types.h
index a6bf191..f4f3f34 100644
--- a/glslang/Include/Types.h
+++ b/glslang/Include/Types.h
@@ -741,6 +741,16 @@ public:
         }
     }
 
+    bool isUniform() const
+    {
+        switch (storage) {
+        case EvqUniform:
+            return true;
+        default:
+            return false;
+        }
+    }
+
     bool isIo() const
     {
         switch (storage) {
diff --git a/glslang/MachineIndependent/iomapper.cpp b/glslang/MachineIndependent/iomapper.cpp
index 7e12864..3486ea6 100644
--- a/glslang/MachineIndependent/iomapper.cpp
+++ b/glslang/MachineIndependent/iomapper.cpp
@@ -1633,6 +1633,37 @@ bool TGlslIoMapper::doMap(TIoMapResolver* resolver, TInfoSink& infoSink) {
             return TVarEntryInfo::TOrderByPriority()(p1.second, p2.second);
         });
         resolver->endResolve(EShLangCount);
+        if (autoPushConstantBlockName.length()) {
+            bool upgraded = false;
+            for (size_t stage = 0; stage < EShLangCount; stage++) {
+                if (intermediates[stage] != nullptr) {
+                    TVarLiveMap** pUniformVarMap = uniformResolve.uniformVarMap;
+                    auto at = pUniformVarMap[stage]->find(autoPushConstantBlockName);
+                    if (at == pUniformVarMap[stage]->end())
+                        continue;
+                    TQualifier& qualifier = at->second.symbol->getQualifier();
+                    if (!qualifier.isUniform())
+                        continue;
+                    TType& t = at->second.symbol->getWritableType();
+                    int size, stride;
+                    TIntermediate::getBaseAlignment(t, size, stride, autoPushConstantBlockPacking,
+                                                    qualifier.layoutMatrix == ElmRowMajor);
+                    if (size <= int(autoPushConstantMaxSize)) {
+                        qualifier.setBlockStorage(EbsPushConstant);
+                        qualifier.layoutPacking = autoPushConstantBlockPacking;
+                        upgraded = true;
+                    }
+                }
+            }
+            // If it's been upgraded to push_constant, then remove it from the uniformVector
+            // so it doesn't get a set/binding assigned to it.
+            if (upgraded) {
+                auto at = std::find_if(uniformVector.begin(), uniformVector.end(),
+                                       [this](const TVarLivePair& p) { return p.first == autoPushConstantBlockName; });
+                if (at != uniformVector.end())
+                    uniformVector.erase(at);
+            }
+        }
         for (size_t stage = 0; stage < EShLangCount; stage++) {
             if (intermediates[stage] != nullptr) {
                 // traverse each stage, set new location to each input/output and unifom symbol, set new binding to
diff --git a/glslang/MachineIndependent/iomapper.h b/glslang/MachineIndependent/iomapper.h
index 07357c2..843ea73 100644
--- a/glslang/MachineIndependent/iomapper.h
+++ b/glslang/MachineIndependent/iomapper.h
@@ -291,7 +291,7 @@ public:
     bool virtual doMap(TIoMapResolver*, TInfoSink&) { return true; }
 };
 
-// I/O mapper for OpenGL
+// I/O mapper for GLSL
 class TGlslIoMapper : public TIoMapper {
 public:
     TGlslIoMapper() {
@@ -301,6 +301,8 @@ public:
         memset(intermediates, 0, sizeof(TIntermediate*) * (EShLangCount + 1));
         profile = ENoProfile;
         version = 0;
+        autoPushConstantMaxSize = 128;
+        autoPushConstantBlockPacking = ElpStd430;
     }
     virtual ~TGlslIoMapper() {
         for (size_t stage = 0; stage < EShLangCount; stage++) {
@@ -320,6 +322,13 @@ public:
                 intermediates[stage] = nullptr;
         }
     }
+	// If set, the uniform block with the given name will be changed to be backed by
+	// push_constant if it's size is <= maxSize
+    void setAutoPushConstantBlock(const char* name, unsigned int maxSize, TLayoutPacking packing) {
+        autoPushConstantBlockName = name;
+        autoPushConstantMaxSize = maxSize;
+        autoPushConstantBlockPacking = packing;
+    }
     // grow the reflection stage by stage
     bool addStage(EShLanguage, TIntermediate&, TInfoSink&, TIoMapResolver*) override;
     bool doMap(TIoMapResolver*, TInfoSink&) override;
@@ -329,6 +338,11 @@ public:
     bool hadError = false;
     EProfile profile;
     int version;
+
+private:
+    TString autoPushConstantBlockName;
+    unsigned int autoPushConstantMaxSize;
+    TLayoutPacking autoPushConstantBlockPacking;
 };
 
 } // end namespace glslang
diff --git a/glslang/MachineIndependent/linkValidate.cpp b/glslang/MachineIndependent/linkValidate.cpp
index 4a68130..4edd2a9 100644
--- a/glslang/MachineIndependent/linkValidate.cpp
+++ b/glslang/MachineIndependent/linkValidate.cpp
@@ -1934,7 +1934,7 @@ int TIntermediate::getBaseAlignment(const TType& type, int& size, int& stride, T
     }
 
     // rule 9
-    if (type.getBasicType() == EbtStruct) {
+    if (type.getBasicType() == EbtStruct || type.getBasicType() == EbtBlock) {
         const TTypeList& memberList = *type.getStruct();
 
         size = 0;
-- 
2.7.4