[opt] Add struct-packing pass and unit test. (#5778)

This pass allows to re-assign offset layout decorations
to tightly pack a struct according to its packing rules.

Android.mk

@@ -183,6 +183,7 @@ SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/strip_debug_info_pass.cpp \
 		source/opt/strip_nonsemantic_info_pass.cpp \
 		source/opt/struct_cfg_analysis.cpp \
+		source/opt/struct_packing_pass.cpp \
 		source/opt/switch_descriptorset_pass.cpp \
 		source/opt/trim_capabilities_pass.cpp \
 		source/opt/type_manager.cpp \

include/spirv-tools/optimizer.hpp

@@ -956,6 +956,15 @@ Optimizer::PassToken CreateFixFuncCallArgumentsPass();
 // the unknown capability interacts with one of the trimmed capabilities.
 Optimizer::PassToken CreateTrimCapabilitiesPass();
 
+// Creates a struct-packing pass.
+// This pass re-assigns all offset layout decorators to tightly pack
+// the struct with OpName matching `structToPack` according to the given packing
+// rule. Accepted packing rules are: std140, std140EnhancedLayout, std430,
+// std430EnhancedLayout, hlslCbuffer, hlslCbufferPackOffset, scalar,
+// scalarEnhancedLayout.
+Optimizer::PassToken CreateStructPackingPass(const char* structToPack,
+                                             const char* packingRule);
+
 // Creates a switch-descriptorset pass.
 // This pass changes any DescriptorSet decorations with the value |ds_from| to
 // use the new value |ds_to|.

source/opt/CMakeLists.txt

@@ -240,6 +240,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   strip_debug_info_pass.cpp
   strip_nonsemantic_info_pass.cpp
   struct_cfg_analysis.cpp
+  struct_packing_pass.cpp
   switch_descriptorset_pass.cpp
   trim_capabilities_pass.cpp
   type_manager.cpp

source/opt/optimizer.cpp

@@ -571,6 +571,26 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag,
              pass_args.c_str());
       return false;
     }
+  } else if (pass_name == "struct-packing") {
+    if (pass_args.size() == 0) {
+      Error(consumer(), nullptr, {},
+            "--struct-packing requires a name:rule argument.");
+      return false;
+    }
+
+    auto separator_pos = pass_args.find(':');
+    if (separator_pos == std::string::npos || separator_pos == 0 ||
+        separator_pos + 1 == pass_args.size()) {
+      Errorf(consumer(), nullptr, {},
+             "Invalid argument for --struct-packing: %s", pass_args.c_str());
+      return false;
+    }
+
+    const std::string struct_name = pass_args.substr(0, separator_pos);
+    const std::string rule_name = pass_args.substr(separator_pos + 1);
+
+    RegisterPass(
+        CreateStructPackingPass(struct_name.c_str(), rule_name.c_str()));
   } else if (pass_name == "switch-descriptorset") {
     if (pass_args.size() == 0) {
       Error(consumer(), nullptr, {},
@@ -1152,6 +1172,14 @@ Optimizer::PassToken CreateTrimCapabilitiesPass() {
       MakeUnique<opt::TrimCapabilitiesPass>());
 }
 
+Optimizer::PassToken CreateStructPackingPass(const char* structToPack,
+                                             const char* packingRule) {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::StructPackingPass>(
+          structToPack,
+          opt::StructPackingPass::ParsePackingRuleFromString(packingRule)));
+}
+
 Optimizer::PassToken CreateSwitchDescriptorSetPass(uint32_t from, uint32_t to) {
   return MakeUnique<Optimizer::PassToken::Impl>(
       MakeUnique<opt::SwitchDescriptorSetPass>(from, to));

source/opt/passes.h

@@ -83,6 +83,7 @@
 #include "source/opt/strength_reduction_pass.h"
 #include "source/opt/strip_debug_info_pass.h"
 #include "source/opt/strip_nonsemantic_info_pass.h"
+#include "source/opt/struct_packing_pass.h"
 #include "source/opt/switch_descriptorset_pass.h"
 #include "source/opt/trim_capabilities_pass.h"
 #include "source/opt/unify_const_pass.h"

source/opt/struct_packing_pass.cpp

@@ -0,0 +1,482 @@
+// Copyright (c) 2024 Epic Games, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "struct_packing_pass.h"
+
+#include <algorithm>
+
+#include "source/opt/instruction.h"
+#include "source/opt/ir_context.h"
+
+namespace spvtools {
+namespace opt {
+
+/*
+Std140 packing rules from the original GLSL 140 specification (see
+https://registry.khronos.org/OpenGL/extensions/ARB/ARB_uniform_buffer_object.txt)
+
+When using the "std140" storage layout, structures will be laid out in
+buffer storage with its members stored in monotonically increasing order
+based on their location in the declaration. A structure and each
+structure member have a base offset and a base alignment, from which an
+aligned offset is computed by rounding the base offset up to a multiple of
+the base alignment. The base offset of the first member of a structure is
+taken from the aligned offset of the structure itself. The base offset of
+all other structure members is derived by taking the offset of the last
+basic machine unit consumed by the previous member and adding one. Each
+structure member is stored in memory at its aligned offset. The members
+of a top-level uniform block are laid out in buffer storage by treating
+the uniform block as a structure with a base offset of zero.
+
+(1) If the member is a scalar consuming <N> basic machine units, the
+    base alignment is <N>.
+
+(2) If the member is a two- or four-component vector with components
+    consuming <N> basic machine units, the base alignment is 2<N> or
+    4<N>, respectively.
+
+(3) If the member is a three-component vector with components consuming
+    <N> basic machine units, the base alignment is 4<N>.
+
+(4) If the member is an array of scalars or vectors, the base alignment
+    and array stride are set to match the base alignment of a single
+    array element, according to rules (1), (2), and (3), and rounded up
+    to the base alignment of a vec4. The array may have padding at the
+    end; the base offset of the member following the array is rounded up
+    to the next multiple of the base alignment.
+
+(5) If the member is a column-major matrix with <C> columns and <R>
+    rows, the matrix is stored identically to an array of <C> column
+    vectors with <R> components each, according to rule (4).
+
+(6) If the member is an array of <S> column-major matrices with <C>
+    columns and <R> rows, the matrix is stored identically to a row of
+    <S>*<C> column vectors with <R> components each, according to rule
+    (4).
+
+(7) If the member is a row-major matrix with <C> columns and <R> rows,
+    the matrix is stored identically to an array of <R> row vectors
+    with <C> components each, according to rule (4).
+
+(8) If the member is an array of <S> row-major matrices with <C> columns
+    and <R> rows, the matrix is stored identically to a row of <S>*<R>
+    row vectors with <C> components each, according to rule (4).
+
+(9) If the member is a structure, the base alignment of the structure is
+    <N>, where <N> is the largest base alignment value of any of its
+    members, and rounded up to the base alignment of a vec4. The
+    individual members of this sub-structure are then assigned offsets
+    by applying this set of rules recursively, where the base offset of
+    the first member of the sub-structure is equal to the aligned offset
+    of the structure. The structure may have padding at the end; the
+    base offset of the member following the sub-structure is rounded up
+    to the next multiple of the base alignment of the structure.
+
+(10) If the member is an array of <S> structures, the <S> elements of
+    the array are laid out in order, according to rule (9).
+*/
+
+static bool isPackingVec4Padded(StructPackingPass::PackingRules rules) {
+  switch (rules) {
+    case StructPackingPass::PackingRules::Std140:
+    case StructPackingPass::PackingRules::Std140EnhancedLayout:
+    case StructPackingPass::PackingRules::HlslCbuffer:
+    case StructPackingPass::PackingRules::HlslCbufferPackOffset:
+      return true;
+    default:
+      return false;
+  }
+}
+
+static bool isPackingScalar(StructPackingPass::PackingRules rules) {
+  switch (rules) {
+    case StructPackingPass::PackingRules::Scalar:
+    case StructPackingPass::PackingRules::ScalarEnhancedLayout:
+      return true;
+    default:
+      return false;
+  }
+}
+
+static bool isPackingHlsl(StructPackingPass::PackingRules rules) {
+  switch (rules) {
+    case StructPackingPass::PackingRules::HlslCbuffer:
+    case StructPackingPass::PackingRules::HlslCbufferPackOffset:
+      return true;
+    default:
+      return false;
+  }
+}
+
+static uint32_t getPackedBaseSize(const analysis::Type& type) {
+  switch (type.kind()) {
+    case analysis::Type::kBool:
+      return 1;
+    case analysis::Type::kInteger:
+      return type.AsInteger()->width() / 8;
+    case analysis::Type::kFloat:
+      return type.AsFloat()->width() / 8;
+    case analysis::Type::kVector:
+      return getPackedBaseSize(*type.AsVector()->element_type());
+    case analysis::Type::kMatrix:
+      return getPackedBaseSize(*type.AsMatrix()->element_type());
+    default:
+      break;  // we only expect bool, int, float, vec, and mat here
+  }
+  assert(0 && "Unrecognized type to get base size");
+  return 0;
+}
+
+static uint32_t getScalarElementCount(const analysis::Type& type) {
+  switch (type.kind()) {
+    case analysis::Type::kVector:
+      return type.AsVector()->element_count();
+    case analysis::Type::kMatrix:
+      return getScalarElementCount(*type.AsMatrix()->element_type());
+    case analysis::Type::kStruct:
+      assert(0 && "getScalarElementCount() does not recognized struct types");
+      return 0;
+    default:
+      return 1;
+  }
+}
+
+// Aligns the specified value to a multiple of alignment, whereas the
+// alignment must be a power-of-two.
+static uint32_t alignPow2(uint32_t value, uint32_t alignment) {
+  return (value + alignment - 1) & ~(alignment - 1);
+}
+
+void StructPackingPass::buildConstantsMap() {
+  constantsMap_.clear();
+  for (Instruction* instr : context()->module()->GetConstants()) {
+    constantsMap_[instr->result_id()] = instr;
+  }
+}
+
+uint32_t StructPackingPass::getPackedAlignment(
+    const analysis::Type& type) const {
+  switch (type.kind()) {
+    case analysis::Type::kArray: {
+      // Get alignment of base type and round up to minimum alignment
+      const uint32_t minAlignment = isPackingVec4Padded(packingRules_) ? 16 : 1;
+      return std::max<uint32_t>(
+          minAlignment, getPackedAlignment(*type.AsArray()->element_type()));
+    }
+    case analysis::Type::kStruct: {
+      // Rule 9. Struct alignment is maximum alignmnet of its members
+      uint32_t alignment = 1;
+
+      for (const analysis::Type* elementType :
+           type.AsStruct()->element_types()) {
+        alignment =
+            std::max<uint32_t>(alignment, getPackedAlignment(*elementType));
+      }
+
+      if (isPackingVec4Padded(packingRules_))
+        alignment = std::max<uint32_t>(alignment, 16u);
+
+      return alignment;
+    }
+    default: {
+      const uint32_t baseAlignment = getPackedBaseSize(type);
+
+      // Scalar block layout always uses alignment for the most basic component
+      if (isPackingScalar(packingRules_)) return baseAlignment;
+
+      if (const analysis::Matrix* matrixType = type.AsMatrix()) {
+        // Rule 5/7
+        if (isPackingVec4Padded(packingRules_) ||
+            matrixType->element_count() == 3)
+          return baseAlignment * 4;
+        else
+          return baseAlignment * matrixType->element_count();
+      } else if (const analysis::Vector* vectorType = type.AsVector()) {
+        // Rule 1
+        if (vectorType->element_count() == 1) return baseAlignment;
+
+        // Rule 2
+        if (vectorType->element_count() == 2 ||
+            vectorType->element_count() == 4)
+          return baseAlignment * vectorType->element_count();
+
+        // Rule 3
+        if (vectorType->element_count() == 3) return baseAlignment * 4;
+      } else {
+        // Rule 1
+        return baseAlignment;
+      }
+    }
+  }
+  assert(0 && "Unrecognized type to get packed alignment");
+  return 0;
+}
+
+static uint32_t getPadAlignment(const analysis::Type& type,
+                                uint32_t packedAlignment) {
+  // The next member following a struct member is aligned to the base alignment
+  // of a previous struct member.
+  return type.kind() == analysis::Type::kStruct ? packedAlignment : 1;
+}
+
+uint32_t StructPackingPass::getPackedSize(const analysis::Type& type) const {
+  switch (type.kind()) {
+    case analysis::Type::kArray: {
+      if (const analysis::Array* arrayType = type.AsArray()) {
+        uint32_t size =
+            getPackedArrayStride(*arrayType) * getArrayLength(*arrayType);
+
+        // For arrays of vector and matrices in HLSL, the last element has a
+        // size depending on its vector/matrix size to allow packing other
+        // vectors in the last element.
+        const analysis::Type* arraySubType = arrayType->element_type();
+        if (isPackingHlsl(packingRules_) &&
+            arraySubType->kind() != analysis::Type::kStruct) {
+          size -= (4 - getScalarElementCount(*arraySubType)) *
+                  getPackedBaseSize(*arraySubType);
+        }
+        return size;
+      }
+      break;
+    }
+    case analysis::Type::kStruct: {
+      uint32_t size = 0;
+      uint32_t padAlignment = 1;
+      for (const analysis::Type* memberType :
+           type.AsStruct()->element_types()) {
+        const uint32_t packedAlignment = getPackedAlignment(*memberType);
+        const uint32_t alignment =
+            std::max<uint32_t>(packedAlignment, padAlignment);
+        padAlignment = getPadAlignment(*memberType, packedAlignment);
+        size = alignPow2(size, alignment);
+        size += getPackedSize(*memberType);
+      }
+      return size;
+    }
+    default: {
+      const uint32_t baseAlignment = getPackedBaseSize(type);
+      if (isPackingScalar(packingRules_)) {
+        return getScalarElementCount(type) * baseAlignment;
+      } else {
+        uint32_t size = 0;
+        if (const analysis::Matrix* matrixType = type.AsMatrix()) {
+          const analysis::Vector* matrixSubType =
+              matrixType->element_type()->AsVector();
+          assert(matrixSubType != nullptr &&
+                 "Matrix sub-type is expected to be a vector type");
+          if (isPackingVec4Padded(packingRules_) ||
+              matrixType->element_count() == 3)
+            size = matrixSubType->element_count() * baseAlignment * 4;
+          else
+            size = matrixSubType->element_count() * baseAlignment *
+                   matrixType->element_count();
+
+          // For matrices in HLSL, the last element has a size depending on its
+          // vector size to allow packing other vectors in the last element.
+          if (isPackingHlsl(packingRules_)) {
+            size -= (4 - matrixSubType->element_count()) *
+                    getPackedBaseSize(*matrixSubType);
+          }
+        } else if (const analysis::Vector* vectorType = type.AsVector()) {
+          size = vectorType->element_count() * baseAlignment;
+        } else {
+          size = baseAlignment;
+        }
+        return size;
+      }
+    }
+  }
+  assert(0 && "Unrecognized type to get packed size");
+  return 0;
+}
+
+uint32_t StructPackingPass::getPackedArrayStride(
+    const analysis::Array& arrayType) const {
+  // Array stride is equal to aligned size of element type
+  const uint32_t elementSize = getPackedSize(*arrayType.element_type());
+  const uint32_t alignment = getPackedAlignment(arrayType);
+  return alignPow2(elementSize, alignment);
+}
+
+uint32_t StructPackingPass::getArrayLength(
+    const analysis::Array& arrayType) const {
+  return getConstantInt(arrayType.LengthId());
+}
+
+uint32_t StructPackingPass::getConstantInt(spv::Id id) const {
+  auto it = constantsMap_.find(id);
+  assert(it != constantsMap_.end() &&
+         "Failed to map SPIR-V instruction ID to constant value");
+  [[maybe_unused]] const analysis::Type* constType =
+      context()->get_type_mgr()->GetType(it->second->type_id());
+  assert(constType != nullptr &&
+         "Failed to map SPIR-V instruction result type to definition");
+  assert(constType->kind() == analysis::Type::kInteger &&
+         "Failed to map SPIR-V instruction result type to integer type");
+  return it->second->GetOperand(2).words[0];
+}
+
+StructPackingPass::PackingRules StructPackingPass::ParsePackingRuleFromString(
+    const std::string& s) {
+  if (s == "std140") return PackingRules::Std140;
+  if (s == "std140EnhancedLayout") return PackingRules::Std140EnhancedLayout;
+  if (s == "std430") return PackingRules::Std430;
+  if (s == "std430EnhancedLayout") return PackingRules::Std430EnhancedLayout;
+  if (s == "hlslCbuffer") return PackingRules::HlslCbuffer;
+  if (s == "hlslCbufferPackOffset") return PackingRules::HlslCbufferPackOffset;
+  if (s == "scalar") return PackingRules::Scalar;
+  if (s == "scalarEnhancedLayout") return PackingRules::ScalarEnhancedLayout;
+  return PackingRules::Undefined;
+}
+
+StructPackingPass::StructPackingPass(const char* structToPack,
+                                     PackingRules rules)
+    : structToPack_{structToPack != nullptr ? structToPack : ""},
+      packingRules_{rules} {}
+
+Pass::Status StructPackingPass::Process() {
+  if (packingRules_ == PackingRules::Undefined) {
+    if (consumer()) {
+      consumer()(SPV_MSG_ERROR, "", {0, 0, 0},
+                 "Cannot pack struct with undefined rule");
+    }
+    return Status::Failure;
+  }
+
+  // Build Id-to-instruction map for easier access
+  buildConstantsMap();
+
+  // Find structure of interest
+  const uint32_t structIdToPack = findStructIdByName(structToPack_.c_str());
+
+  const Instruction* structDef =
+      context()->get_def_use_mgr()->GetDef(structIdToPack);
+  if (structDef == nullptr || structDef->opcode() != spv::Op::OpTypeStruct) {
+    if (consumer()) {
+      const std::string message =
+          "Failed to find struct with name " + structToPack_;
+      consumer()(SPV_MSG_ERROR, "", {0, 0, 0}, message.c_str());
+    }
+    return Status::Failure;
+  }
+
+  // Find all struct member types
+  std::vector<const analysis::Type*> structMemberTypes =
+      findStructMemberTypes(*structDef);
+
+  return assignStructMemberOffsets(structIdToPack, structMemberTypes);
+}
+
+uint32_t StructPackingPass::findStructIdByName(const char* structName) const {
+  for (Instruction& instr : context()->module()->debugs2()) {
+    if (instr.opcode() == spv::Op::OpName &&
+        instr.GetOperand(1).AsString() == structName) {
+      return instr.GetOperand(0).AsId();
+    }
+  }
+  return 0;
+}
+
+std::vector<const analysis::Type*> StructPackingPass::findStructMemberTypes(
+    const Instruction& structDef) const {
+  // Found struct type to pack, now collect all types of its members
+  assert(structDef.NumOperands() > 0 &&
+         "Number of operands in OpTypeStruct instruction must not be zero");
+  const uint32_t numMembers = structDef.NumOperands() - 1;
+  std::vector<const analysis::Type*> structMemberTypes;
+  structMemberTypes.resize(numMembers);
+  for (uint32_t i = 0; i < numMembers; ++i) {
+    const spv::Id memberTypeId = structDef.GetOperand(1 + i).AsId();
+    if (const analysis::Type* memberType =
+            context()->get_type_mgr()->GetType(memberTypeId)) {
+      structMemberTypes[i] = memberType;
+    }
+  }
+  return structMemberTypes;
+}
+
+Pass::Status StructPackingPass::assignStructMemberOffsets(
+    uint32_t structIdToPack,
+    const std::vector<const analysis::Type*>& structMemberTypes) {
+  // Returns true if the specified instruction is a OpMemberDecorate for the
+  // struct we're looking for with an offset decoration
+  auto isMemberOffsetDecoration =
+      [structIdToPack](const Instruction& instr) -> bool {
+    return instr.opcode() == spv::Op::OpMemberDecorate &&
+           instr.GetOperand(0).AsId() == structIdToPack &&
+           static_cast<spv::Decoration>(instr.GetOperand(2).words[0]) ==
+               spv::Decoration::Offset;
+  };
+
+  bool modified = false;
+
+  // Find and re-assign all member offset decorations
+  for (auto it = context()->module()->annotation_begin(),
+            itEnd = context()->module()->annotation_end();
+       it != itEnd; ++it) {
+    if (isMemberOffsetDecoration(*it)) {
+      // Found first member decoration with offset, we expect all other
+      // offsets right after the first one
+      uint32_t prevMemberIndex = 0;
+      uint32_t currentOffset = 0;
+      uint32_t padAlignment = 1;
+      do {
+        const uint32_t memberIndex = it->GetOperand(1).words[0];
+        if (memberIndex < prevMemberIndex) {
+          // Failure: we expect all members to appear in consecutive order
+          return Status::Failure;
+        }
+
+        // Apply alignment rules to current offset
+        const analysis::Type& memberType = *structMemberTypes[memberIndex];
+        uint32_t packedAlignment = getPackedAlignment(memberType);
+        uint32_t packedSize = getPackedSize(memberType);
+
+        if (isPackingHlsl(packingRules_)) {
+          // If a member crosses vec4 boundaries, alignment is size of vec4
+          if (currentOffset / 16 != (currentOffset + packedSize - 1) / 16)
+            packedAlignment = std::max<uint32_t>(packedAlignment, 16u);
+        }
+
+        const uint32_t alignment =
+            std::max<uint32_t>(packedAlignment, padAlignment);
+        currentOffset = alignPow2(currentOffset, alignment);
+        padAlignment = getPadAlignment(memberType, packedAlignment);
+
+        // Override packed offset in instruction
+        if (it->GetOperand(3).words[0] < currentOffset) {
+          // Failure: packing resulted in higher offset for member than
+          // previously generated
+          return Status::Failure;
+        }
+
+        it->GetOperand(3).words[0] = currentOffset;
+        modified = true;
+
+        // Move to next member
+        ++it;
+        prevMemberIndex = memberIndex;
+        currentOffset += packedSize;
+      } while (it != itEnd && isMemberOffsetDecoration(*it));
+
+      // We're done with all decorations for the struct of interest
+      break;
+    }
+  }
+
+  return modified ? Status::SuccessWithChange : Status::SuccessWithoutChange;
+}
+
+}  // namespace opt
+}  // namespace spvtools

source/opt/struct_packing_pass.h

@@ -0,0 +1,81 @@
+// Copyright (c) 2024 Epic Games, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SOURCE_OPT_STRUCT_PACKING_PASS_
+#define SOURCE_OPT_STRUCT_PACKING_PASS_
+
+#include <unordered_map>
+
+#include "source/opt/ir_context.h"
+#include "source/opt/module.h"
+#include "source/opt/pass.h"
+
+namespace spvtools {
+namespace opt {
+
+// This pass re-assigns all field offsets under the specified packing rules.
+class StructPackingPass final : public Pass {
+ public:
+  enum class PackingRules {
+    Undefined,
+    Std140,
+    Std140EnhancedLayout,
+    Std430,
+    Std430EnhancedLayout,
+    HlslCbuffer,
+    HlslCbufferPackOffset,
+    Scalar,
+    ScalarEnhancedLayout,
+  };
+
+  static PackingRules ParsePackingRuleFromString(const std::string& s);
+
+  StructPackingPass(const char* structToPack, PackingRules rules);
+  const char* name() const override { return "struct-packing"; }
+  Status Process() override;
+
+  IRContext::Analysis GetPreservedAnalyses() override {
+    return IRContext::kAnalysisCombinators | IRContext::kAnalysisCFG |
+           IRContext::kAnalysisDominatorAnalysis |
+           IRContext::kAnalysisLoopAnalysis | IRContext::kAnalysisNameMap |
+           IRContext::kAnalysisScalarEvolution |
+           IRContext::kAnalysisStructuredCFG | IRContext::kAnalysisConstants |
+           IRContext::kAnalysisDebugInfo | IRContext::kAnalysisLiveness;
+  }
+
+ private:
+  void buildConstantsMap();
+  uint32_t findStructIdByName(const char* structName) const;
+  std::vector<const analysis::Type*> findStructMemberTypes(
+      const Instruction& structDef) const;
+  Status assignStructMemberOffsets(
+      uint32_t structIdToPack,
+      const std::vector<const analysis::Type*>& structMemberTypes);
+
+  uint32_t getPackedAlignment(const analysis::Type& type) const;
+  uint32_t getPackedSize(const analysis::Type& type) const;
+  uint32_t getPackedArrayStride(const analysis::Array& arrayType) const;
+  uint32_t getArrayLength(const analysis::Array& arrayType) const;
+  uint32_t getConstantInt(spv::Id id) const;
+
+ private:
+  std::string structToPack_;
+  PackingRules packingRules_ = PackingRules::Undefined;
+  std::unordered_map<spv::Id, Instruction*> constantsMap_;
+};
+
+}  // namespace opt
+}  // namespace spvtools
+
+#endif  // SOURCE_OPT_STRUCT_PACKING_PASS_

test/opt/CMakeLists.txt

@@ -104,6 +104,7 @@ add_spvtools_unittest(TARGET opt
        strip_debug_info_test.cpp
        strip_nonsemantic_info_test.cpp
        struct_cfg_analysis_test.cpp
+       struct_packing_test.cpp
        switch_descriptorset_test.cpp
        trim_capabilities_pass_test.cpp
        type_manager_test.cpp

test/opt/struct_packing_test.cpp

@@ -0,0 +1,242 @@
+// Copyright (c) 2024 Epic Games, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <string>
+
+#include "gmock/gmock.h"
+#include "source/opt/struct_packing_pass.h"
+#include "test/opt/pass_fixture.h"
+#include "test/opt/pass_utils.h"
+
+namespace spvtools {
+namespace opt {
+namespace {
+
+using StructPackingTest = PassTest<::testing::Test>;
+
+TEST_F(StructPackingTest, PackSimpleStructStd140) {
+  // #version 420
+  //
+  // layout(std140, binding = 0) uniform Globals {
+  //   layout(offset = 16) vec3 a_xyz;
+  //   float a_w;
+  //   layout(offset = 128) vec3 b_xyz;
+  //   int b_w;
+  // };
+  //
+  // void main() {}
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main"
+OpExecutionMode %main OriginLowerLeft
+OpSource GLSL 420
+OpName %main "main"
+OpName %Globals "Globals"
+OpMemberName %Globals 0 "a_xyz"
+OpMemberName %Globals 1 "a_w"
+OpMemberName %Globals 2 "b_xyz"
+OpMemberName %Globals 3 "b_w"
+OpName %_ ""
+; CHECK: OpMemberDecorate %Globals 0 Offset 0
+OpMemberDecorate %Globals 0 Offset 16
+; CHECK: OpMemberDecorate %Globals 1 Offset 12
+OpMemberDecorate %Globals 1 Offset 28
+; CHECK: OpMemberDecorate %Globals 2 Offset 16
+OpMemberDecorate %Globals 2 Offset 128
+; CHECK: OpMemberDecorate %Globals 3 Offset 28
+OpMemberDecorate %Globals 3 Offset 140
+OpDecorate %Globals Block
+OpDecorate %_ DescriptorSet 0
+OpDecorate %_ Binding 0
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v3float = OpTypeVector %float 3
+%int = OpTypeInt 32 1
+%Globals = OpTypeStruct %v3float %float %v3float %int
+%_ptr_Uniform_Globals = OpTypePointer Uniform %Globals
+%_ = OpVariable %_ptr_Uniform_Globals Uniform
+%main = OpFunction %void None %3
+%5 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<StructPackingPass>(
+      spirv, true, "Globals", StructPackingPass::PackingRules::Std140);
+}
+
+TEST_F(StructPackingTest, PackSimpleStructWithPaddingStd140) {
+  // #version 420
+  //
+  // layout(std140, binding = 0) uniform Globals {
+  //   layout(offset = 16) vec3 a_xyz;
+  //   float a_w;
+  //   float b_x_padding_yzw;
+  //   layout(offset = 128) vec3 c_xyz;
+  //   int c_w;
+  // };
+  //
+  // void main() {}
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main"
+OpExecutionMode %main OriginLowerLeft
+OpSource GLSL 420
+OpName %main "main"
+OpName %Globals "Globals"
+OpMemberName %Globals 0 "a_xyz"
+OpMemberName %Globals 1 "a_w"
+OpMemberName %Globals 2 "b_x_padding_yzw"
+OpMemberName %Globals 3 "c_xyz"
+OpMemberName %Globals 4 "c_w"
+OpName %_ ""
+; CHECK: OpMemberDecorate %Globals 0 Offset 0
+OpMemberDecorate %Globals 0 Offset 16
+; CHECK: OpMemberDecorate %Globals 1 Offset 12
+OpMemberDecorate %Globals 1 Offset 28
+; CHECK: OpMemberDecorate %Globals 2 Offset 16
+OpMemberDecorate %Globals 2 Offset 32
+; CHECK: OpMemberDecorate %Globals 3 Offset 32
+OpMemberDecorate %Globals 3 Offset 128
+; CHECK: OpMemberDecorate %Globals 4 Offset 44
+OpMemberDecorate %Globals 4 Offset 140
+OpDecorate %Globals Block
+OpDecorate %_ DescriptorSet 0
+OpDecorate %_ Binding 0
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v3float = OpTypeVector %float 3
+%int = OpTypeInt 32 1
+%Globals = OpTypeStruct %v3float %float %float %v3float %int
+%_ptr_Uniform_Globals = OpTypePointer Uniform %Globals
+%_ = OpVariable %_ptr_Uniform_Globals Uniform
+%main = OpFunction %void None %3
+%5 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<StructPackingPass>(
+      spirv, true, "Globals", StructPackingPass::PackingRules::Std140);
+}
+
+TEST_F(StructPackingTest, PackSimpleScalarArrayStd140) {
+  // #version 420
+  //
+  // layout(std140, binding = 0) uniform Globals {
+  //   layout(offset = 16) float a[2];
+  //   layout(offset = 128) float b[2]; // Must become offset 32 with std140
+  // };
+  //
+  // void main() {}
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main"
+OpExecutionMode %main OriginLowerLeft
+OpSource GLSL 420
+OpName %main "main"
+OpName %Globals "Globals"
+OpMemberName %Globals 0 "a"
+OpMemberName %Globals 1 "b"
+OpName %_ ""
+OpDecorate %_arr_float_uint_2 ArrayStride 16
+OpDecorate %_arr_float_uint_2_0 ArrayStride 16
+; CHECK: OpMemberDecorate %Globals 0 Offset 0
+OpMemberDecorate %Globals 0 Offset 16
+; CHECK: OpMemberDecorate %Globals 1 Offset 32
+OpMemberDecorate %Globals 1 Offset 128
+OpDecorate %Globals Block
+OpDecorate %_ DescriptorSet 0
+OpDecorate %_ Binding 0
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%uint = OpTypeInt 32 0
+%uint_2 = OpConstant %uint 2
+%_arr_float_uint_2 = OpTypeArray %float %uint_2
+%_arr_float_uint_2_0 = OpTypeArray %float %uint_2
+%Globals = OpTypeStruct %_arr_float_uint_2 %_arr_float_uint_2_0
+%_ptr_Uniform_Globals = OpTypePointer Uniform %Globals
+%_ = OpVariable %_ptr_Uniform_Globals Uniform
+%main = OpFunction %void None %3
+%5 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<StructPackingPass>(
+      spirv, true, "Globals", StructPackingPass::PackingRules::Std140);
+}
+
+TEST_F(StructPackingTest, PackSimpleScalarArrayStd430) {
+  // #version 430
+  //
+  // layout(std430, binding = 0) buffer Globals {
+  //   layout(offset = 16) float a[2];
+  //   layout(offset = 128) float b[2]; // Must become offset 8 with std430
+  // };
+  //
+  // void main() {}
+  const std::string spirv = R"(
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main"
+OpExecutionMode %main OriginLowerLeft
+OpSource GLSL 430
+OpName %main "main"
+OpName %Globals "Globals"
+OpMemberName %Globals 0 "a"
+OpMemberName %Globals 1 "b"
+OpName %_ ""
+OpDecorate %_arr_float_uint_2 ArrayStride 4
+OpDecorate %_arr_float_uint_2_0 ArrayStride 4
+; CHECK: OpMemberDecorate %Globals 0 Offset 0
+OpMemberDecorate %Globals 0 Offset 16
+; CHECK: OpMemberDecorate %Globals 1 Offset 8
+OpMemberDecorate %Globals 1 Offset 128
+OpDecorate %Globals BufferBlock
+OpDecorate %_ DescriptorSet 0
+OpDecorate %_ Binding 0
+%void = OpTypeVoid
+%3 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%uint = OpTypeInt 32 0
+%uint_2 = OpConstant %uint 2
+%_arr_float_uint_2 = OpTypeArray %float %uint_2
+%_arr_float_uint_2_0 = OpTypeArray %float %uint_2
+%Globals = OpTypeStruct %_arr_float_uint_2 %_arr_float_uint_2_0
+%_ptr_Uniform_Globals = OpTypePointer Uniform %Globals
+%_ = OpVariable %_ptr_Uniform_Globals Uniform
+%main = OpFunction %void None %3
+%5 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<StructPackingPass>(
+      spirv, true, "Globals", StructPackingPass::PackingRules::Std430);
+}
+
+}  // namespace
+}  // namespace opt
+}  // namespace spvtools

tools/opt/opt.cpp

@@ -515,6 +515,10 @@ Options (in lexicographical order):)",
                covers reflection information defined by
                SPV_GOOGLE_hlsl_functionality1 and SPV_KHR_non_semantic_info)");
   printf(R"(
+  --struct-packing=name:rule
+               Re-assign layout offsets to a given struct according to
+               its packing rules.)");
+  printf(R"(
   --switch-descriptorset=<from>:<to>
                Switch any DescriptoSet decorations using the value <from> to
                the new value <to>.)");

utils/check_copyright.py

@@ -42,7 +42,8 @@ AUTHORS = ['The Khronos Group Inc.',
            'Mostafa Ashraf',
            'Shiyu Liu',
            'ZHOU He',
-           'Nintendo']
+           'Nintendo',
+           'Epic Games, Inc.']
 CURRENT_YEAR = 2023
 
 FIRST_YEAR = 2014