Fix: Vulkan buffer alignment

brenocq · Dec 28, 2023 · 7e0cd07 · 7e0cd07
1 parent d9b3ce4
commit 7e0cd07
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 16 deletions.
diff --git a/src/atta/graphics/bufferLayout.cpp b/src/atta/graphics/bufferLayout.cpp
@@ -171,7 +171,7 @@ uint32_t BufferLayout::Element::componentCountFromType(Type type) {
 //----------------------------------//
 //---------- BufferLayout ----------//
 //----------------------------------//
-BufferLayout::BufferLayout(AlignmentType alignmentType) : _alignmentType(alignmentType) {}
+BufferLayout::BufferLayout(AlignmentType alignmentType) : _alignmentType(alignmentType), _forceNextAlign(0) {}
 
 void BufferLayout::setAlignmentType(AlignmentType alignmentType) { _alignmentType = alignmentType; }
 
@@ -187,14 +187,21 @@ void BufferLayout::push(Element::Type type, std::string name, uint32_t customAli
     if (_elements.empty())
         e.offset = 0;
     else {
-        uint32_t align = customAlign ? customAlign : Element::alignmentFromType(_alignmentType, type);
+        e.align = std::max(customAlign, _forceNextAlign);
+        e.align = e.align ? e.align : Element::alignmentFromType(_alignmentType, type);
+        _forceNextAlign = 0;
         uint32_t offset = _elements.back().offset + _elements.back().size;
-        e.offset = (offset + align - 1) & ~(align - 1);
+        e.offset = (offset + e.align - 1) & ~(e.align - 1);
     }
 
     _elements.push_back(e);
 }
 
+void BufferLayout::pushStructArrayAlign() {
+    if (_alignmentType == AlignmentType::STD140 || _alignmentType == AlignmentType::STD430)
+        _forceNextAlign = 16;
+}
+
 const std::vector<BufferLayout::Element>& BufferLayout::getElements() const { return _elements; }
 
 bool BufferLayout::exists(std::string name) const {

diff --git a/src/atta/graphics/bufferLayout.h b/src/atta/graphics/bufferLayout.h
@@ -27,6 +27,7 @@ class BufferLayout final {
         Type type;
         std::string name;
         uint32_t offset;
+        uint32_t align;
         uint32_t size;
 
         /// Convert string to type ("float" -> FLOAT, "sampler2D" -> SAMPLER_2D, ...)
@@ -60,6 +61,13 @@ class BufferLayout final {
      */
     void push(Element::Type type, std::string name, uint32_t customAlign = 0);
 
+    /**
+     * @brief Push struct/array alignment
+     *
+     * Ensure buffer alignment for the start/end of the struct/array
+     */
+    void pushStructArrayAlign();
+
     /// Get elements
     const std::vector<Element>& getElements() const;
 
@@ -78,6 +86,7 @@ class BufferLayout final {
   private:
     AlignmentType _alignmentType;   ///< How the elements should be aligned in the buffer
     std::vector<Element> _elements; ///< Elements in the buffer
+    uint32_t _forceNextAlign;       ///< Used to force next element alignment when pushing
 };
 
 } // namespace atta::graphics

diff --git a/src/atta/graphics/shader.cpp b/src/atta/graphics/shader.cpp
@@ -325,9 +325,12 @@ void Shader::parseCustomTypes() {
             BufferLayout::Element::Type type = BufferLayout::Element::typeFromString(memberType);
             for (size_t a = 0; a < std::max(size_t(1), memberArraySize); a++) {
                 std::string fullMemberName = memberName;
-                // Handle array types;
-                if (memberArraySize)
+                // Handle array types
+                if (memberArraySize) {
                     fullMemberName += "[" + std::to_string(a) + "]";
+                    // Array start alignment
+                    structLayout.pushStructArrayAlign();
+                }
 
                 if (type != BufferLayout::Element::Type::NONE) {
                     // Handle basic type
@@ -346,16 +349,23 @@ void Shader::parseCustomTypes() {
                         continue;
                     }
 
+                    // Handle struct start alignment
+                    structLayout.pushStructArrayAlign();
+
                     const BufferLayout& memberStructLayout = _customTypes[memberType];
                     for (size_t i = 0; i < memberStructLayout.getElements().size(); i++) {
                         const BufferLayout::Element& e = memberStructLayout.getElements().at(i);
                         std::string name = fullMemberName + "." + e.name;
-                        if (i == 0)
-                            structLayout.push(e.type, name, memberStructLayout.getAlignment());
-                        else
-                            structLayout.push(e.type, name);
+                        structLayout.push(e.type, name, e.align);
                     }
+
+                    // Handle struct end alignment
+                    structLayout.pushStructArrayAlign();
                 }
+
+                // Handle array end alignment
+                if (memberArraySize)
+                    structLayout.pushStructArrayAlign();
             }
 
             memberStart = memberMatch.suffix().first;
@@ -418,6 +428,14 @@ void Shader::populateDescriptorLayouts() {
             if (arraySize)
                 fullName += "[" + std::to_string(i) + "]";
 
+            // Handle array start alignment
+            if (arraySize && type != "sampler2D" && type != "samplerCube") {
+                if (freq == "perFrame")
+                    _perFrameLayout.pushStructArrayAlign();
+                else if (freq == "perDraw")
+                    _perDrawLayout.pushStructArrayAlign();
+            }
+
             if (t != BufferLayout::Element::Type::NONE) {
                 // Handle basic types
                 if (type == "sampler2D" || type == "samplerCube") {
@@ -435,15 +453,35 @@ void Shader::populateDescriptorLayouts() {
                 // Handle custom types
                 // NOTE: There should be no sampler2D or samplerCube in the struct
                 const BufferLayout& structLayout = _customTypes[type];
+
+                // Handle struct start alignment
+                if (freq == "perFrame")
+                    _perFrameLayout.pushStructArrayAlign();
+                else if (freq == "perDraw")
+                    _perDrawLayout.pushStructArrayAlign();
+
                 for (size_t i = 0; i < structLayout.getElements().size(); i++) {
                     const BufferLayout::Element& e = structLayout.getElements().at(i);
                     std::string memberName = fullName + "." + e.name;
-                    uint32_t alignment = i == 0 ? structLayout.getAlignment() : 0;
                     if (freq == "perFrame")
-                        _perFrameLayout.push(e.type, memberName, alignment);
+                        _perFrameLayout.push(e.type, memberName, e.align);
                     else if (freq == "perDraw")
-                        _perDrawLayout.push(e.type, memberName, alignment);
+                        _perDrawLayout.push(e.type, memberName, e.align);
                 }
+
+                // Handle struct end alignment
+                if (freq == "perFrame")
+                    _perFrameLayout.pushStructArrayAlign();
+                else if (freq == "perDraw")
+                    _perDrawLayout.pushStructArrayAlign();
+            }
+
+            // Handle array end alignment
+            if (arraySize && type != "sampler2D" && type != "samplerCube") {
+                if (freq == "perFrame")
+                    _perFrameLayout.pushStructArrayAlign();
+                else if (freq == "perDraw")
+                    _perDrawLayout.pushStructArrayAlign();
             }
         }
 
@@ -465,10 +503,6 @@ void Shader::populateDescriptorLayouts() {
         // Move to the next match
         start = match.suffix().first;
     }
-
-    // TODO debug offsets
-    // for (auto e : _perFrameLayout.getElements())
-    //    LOG_DEBUG("Shader", "[r]$0 [y]offset $1", e.name, e.offset);
 }
 
 Shader::LayoutMember::LayoutMember(std::string type_, std::string name_) : type(type_), name(name_), isArray(false), arraySize(0) {}