bitfields: refactorize struct offset computation (#4781)

This is some preliminary work to bitfield addition into the SPIR-V
backend. This commit should not change the behavior of our backend.

Signed-off-by: Nathan Gauër <[email protected]>

Signed-off-by: Nathan Gauër <[email protected]>

tools/clang/lib/SPIRV/AlignmentSizeCalculator.cpp

@@ -36,7 +36,7 @@ namespace spirv {
 
 void AlignmentSizeCalculator::alignUsingHLSLRelaxedLayout(
     QualType fieldType, uint32_t fieldSize, uint32_t fieldAlignment,
-    uint32_t *currentOffset) {
+    uint32_t *currentOffset) const {
   QualType vecElemType = {};
   const bool fieldIsVecType = isVectorType(fieldType, &vecElemType);
 
@@ -64,7 +64,7 @@ void AlignmentSizeCalculator::alignUsingHLSLRelaxedLayout(
 
 std::pair<uint32_t, uint32_t> AlignmentSizeCalculator::getAlignmentAndSize(
     QualType type, SpirvLayoutRule rule, llvm::Optional<bool> isRowMajor,
-    uint32_t *stride) {
+    uint32_t *stride) const {
   // std140 layout rules:
 
   // 1. If the member is a scalar consuming N basic machine units, the base

tools/clang/lib/SPIRV/AlignmentSizeCalculator.h

@@ -38,7 +38,7 @@ public:
   /// size contains post-paddings required by the given type.
   std::pair<uint32_t, uint32_t>
   getAlignmentAndSize(QualType type, SpirvLayoutRule rule,
-                      llvm::Optional<bool> isRowMajor, uint32_t *stride);
+                      llvm::Optional<bool> isRowMajor, uint32_t *stride) const;
 
   /// \brief Aligns currentOffset properly to allow packing vectors in the HLSL
   /// way: using the element type's alignment as the vector alignment, as long
@@ -47,11 +47,12 @@ public:
   /// calculated without considering the HLSL vector relaxed rule.
   void alignUsingHLSLRelaxedLayout(QualType fieldType, uint32_t fieldSize,
                                    uint32_t fieldAlignment,
-                                   uint32_t *currentOffset);
+                                   uint32_t *currentOffset) const;
 
   /// \brief Returns true if we use row-major matrix for type. Otherwise,
   /// returns false.
-  bool useRowMajor(llvm::Optional<bool> isRowMajor, clang::QualType type) {
+  bool useRowMajor(llvm::Optional<bool> isRowMajor,
+                   clang::QualType type) const {
     return isRowMajor.hasValue() ? isRowMajor.getValue()
                                  : isRowMajorMatrix(spvOptions, type);
   }
@@ -60,7 +61,7 @@ private:
   /// Emits error to the diagnostic engine associated with this visitor.
   template <unsigned N>
   DiagnosticBuilder emitError(const char (&message)[N],
-                              SourceLocation srcLoc = {}) {
+                              SourceLocation srcLoc = {}) const {
     const auto diagId = astContext.getDiagnostics().getCustomDiagID(
         clang::DiagnosticsEngine::Error, message);
     return astContext.getDiagnostics().Report(srcLoc, diagId);

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -14,6 +14,9 @@
 #include "clang/SPIRV/AstTypeProbe.h"
 #include "clang/SPIRV/SpirvFunction.h"
 
+namespace clang {
+namespace spirv {
+
 namespace {
 /// Returns the :packoffset() annotation on the given decl. Returns nullptr if
 /// the decl does not have one.
@@ -30,10 +33,64 @@ inline uint32_t roundToPow2(uint32_t val, uint32_t pow2) {
   return (val + pow2 - 1) & ~(pow2 - 1);
 }
 
-} // end anonymous namespace
+// This method sorts a field list in the following order:
+//  - fields with register annotation first, sorted by register index.
+//  - then fields without annotation, in order of declaration.
+std::vector<const HybridStructType::FieldInfo *>
+sortFields(llvm::ArrayRef<HybridStructType::FieldInfo> fields) {
+  std::vector<const HybridStructType::FieldInfo *> output;
+  output.resize(fields.size());
+
+  auto back_inserter = output.rbegin();
+  std::map<uint32_t, const HybridStructType::FieldInfo *> fixed_fields;
+  for (auto it = fields.rbegin(); it < fields.rend(); it++) {
+    if (it->registerC) {
+      fixed_fields.insert({it->registerC->RegisterNumber, &*it});
+    } else {
+      *back_inserter = &*it;
+      back_inserter++;
+    }
+  }
 
-namespace clang {
-namespace spirv {
+  auto front_inserter = output.begin();
+  for (const auto &item : fixed_fields) {
+    *front_inserter = item.second;
+    front_inserter++;
+  }
+  return output;
+}
+
+// Correctly determine a field offset/size/padding depending on its neighbors
+// and other rules.
+void setDefaultFieldOffsetAndSize(
+    const AlignmentSizeCalculator &alignmentCalc, const SpirvLayoutRule rule,
+    const uint32_t previousFieldEnd,
+    const HybridStructType::FieldInfo *currentField,
+    StructType::FieldInfo *field) {
+
+  const auto &fieldType = currentField->astType;
+  uint32_t memberAlignment = 0, memberSize = 0, stride = 0;
+  std::tie(memberAlignment, memberSize) = alignmentCalc.getAlignmentAndSize(
+      fieldType, rule, /*isRowMajor*/ llvm::None, &stride);
+  field->sizeInBytes = memberSize;
+
+  const uint32_t baseOffset = previousFieldEnd;
+
+  // The next avaiable location after laying out the previous members
+  if (rule != SpirvLayoutRule::RelaxedGLSLStd140 &&
+      rule != SpirvLayoutRule::RelaxedGLSLStd430 &&
+      rule != SpirvLayoutRule::FxcCTBuffer) {
+    field->offset = roundToPow2(baseOffset, memberAlignment);
+    return;
+  }
+
+  uint32_t newOffset = previousFieldEnd;
+  alignmentCalc.alignUsingHLSLRelaxedLayout(fieldType, memberSize,
+                                            memberAlignment, &newOffset);
+  field->offset = newOffset;
+}
+
+} // end anonymous namespace
 
 bool LowerTypeVisitor::visit(SpirvFunction *fn, Phase phase) {
   if (phase == Visitor::Phase::Done) {
@@ -803,155 +860,146 @@ LowerTypeVisitor::translateSampledTypeToImageFormat(QualType sampledType,
   return spv::ImageFormat::Unknown;
 }
 
+StructType::FieldInfo
+LowerTypeVisitor::lowerField(const HybridStructType::FieldInfo *field,
+                             SpirvLayoutRule rule) {
+  auto fieldType = field->astType;
+  // Lower the field type fist. This call will populate proper matrix
+  // majorness information.
+  StructType::FieldInfo loweredField(
+      lowerType(fieldType, rule, /*isRowMajor*/ llvm::None, {}), field->name);
+
+  // Set RelaxedPrecision information for the lowered field.
+  if (isRelaxedPrecisionType(fieldType, spvOptions)) {
+    loweredField.isRelaxedPrecision = true;
+  }
+  if (field->isPrecise) {
+    loweredField.isPrecise = true;
+  }
+
+  // We only need layout information for structures with non-void layout rule.
+  if (rule == SpirvLayoutRule::Void) {
+    return loweredField;
+  }
+
+  // Each structure-type member that is a matrix or array-of-matrices must be
+  // decorated with
+  // * A MatrixStride decoration, and
+  // * one of the RowMajor or ColMajor Decorations.
+  if (const auto *arrayType = astContext.getAsConstantArrayType(fieldType)) {
+    // We have an array of matrices as a field, we need to decorate
+    // MatrixStride on the field. So skip possible arrays here.
+    fieldType = arrayType->getElementType();
+  }
+
+  // Non-floating point matrices are represented as arrays of vectors, and
+  // therefore ColMajor and RowMajor decorations should not be applied to
+  // them.
+  QualType elemType = {};
+  if (isMxNMatrix(fieldType, &elemType) && elemType->isFloatingType()) {
+    uint32_t stride = 0;
+    alignmentCalc.getAlignmentAndSize(fieldType, rule,
+                                      /*isRowMajor*/ llvm::None, &stride);
+    loweredField.matrixStride = stride;
+    loweredField.isRowMajor = isRowMajorMatrix(spvOptions, fieldType);
+  }
+  return loweredField;
+}
+
 llvm::SmallVector<StructType::FieldInfo, 4>
 LowerTypeVisitor::populateLayoutInformation(
     llvm::ArrayRef<HybridStructType::FieldInfo> fields, SpirvLayoutRule rule) {
 
-  // The resulting vector of fields with proper layout information.
-  llvm::SmallVector<StructType::FieldInfo, 4> loweredFields;
-  llvm::SmallVector<StructType::FieldInfo, 4> result;
-
-  using RegisterFieldPair =
-      std::pair<uint32_t, const HybridStructType::FieldInfo *>;
-  struct RegisterFieldPairLess {
-    bool operator()(const RegisterFieldPair &obj1,
-                    const RegisterFieldPair &obj2) const {
-      return obj1.first < obj2.first;
-    }
-  };
-  std::set<RegisterFieldPair, RegisterFieldPairLess> registerCSet;
-  std::vector<const HybridStructType::FieldInfo *> sortedFields;
-  llvm::DenseMap<const HybridStructType::FieldInfo *, uint32_t> fieldToIndexMap;
-
-  // First, check to see if any of the structure members had 'register(c#)'
-  // location semantics. If so, members that do not have the 'register(c#)'
-  // assignment should be allocated after the *highest explicit address*.
-  // Example:
-  // float x : register(c10);   // Offset = 160 (10 * 16)
-  // float y;                   // Offset = 164 (160 + 4)
-  // float z: register(c1);     // Offset = 16  (1  * 16)
-  for (const auto &field : fields)
-    if (field.registerC)
-      registerCSet.insert(
-          RegisterFieldPair(field.registerC->RegisterNumber, &field));
-  for (const auto &pair : registerCSet)
-    sortedFields.push_back(pair.second);
-  for (const auto &field : fields)
-    if (!field.registerC)
-      sortedFields.push_back(&field);
-
-  uint32_t offset = 0;
-  for (const auto *fieldPtr : sortedFields) {
-    // The field can only be FieldDecl (for normal structs) or VarDecl (for
-    // HLSLBufferDecls).
-    const auto field = *fieldPtr;
-    auto fieldType = field.astType;
-    fieldToIndexMap[fieldPtr] = loweredFields.size();
-
-    // Lower the field type fist. This call will populate proper matrix
-    // majorness information.
-    StructType::FieldInfo loweredField(
-        lowerType(fieldType, rule, /*isRowMajor*/ llvm::None, {}), field.name);
-
-    // Set RelaxedPrecision information for the lowered field.
-    if (isRelaxedPrecisionType(fieldType, spvOptions)) {
-      loweredField.isRelaxedPrecision = true;
-    }
-
-    // Set 'precise' information for the lowered field.
-    if (field.isPrecise) {
-      loweredField.isPrecise = true;
-    }
-
-    // We only need layout information for strcutres with non-void layout rule.
+  auto fieldVisitor = [this,
+                       &rule](const StructType::FieldInfo *previousField,
+                              const HybridStructType::FieldInfo *currentField) {
+    StructType::FieldInfo loweredField = lowerField(currentField, rule);
+    // We only need layout information for structures with non-void layout rule.
     if (rule == SpirvLayoutRule::Void) {
-      loweredFields.push_back(loweredField);
-      continue;
+      return loweredField;
     }
 
-    uint32_t memberAlignment = 0, memberSize = 0, stride = 0;
-    std::tie(memberAlignment, memberSize) = alignmentCalc.getAlignmentAndSize(
-        fieldType, rule, /*isRowMajor*/ llvm::None, &stride);
-
-    // The next avaiable location after laying out the previous members
-    const uint32_t nextLoc = offset;
-
-    if (rule == SpirvLayoutRule::RelaxedGLSLStd140 ||
-        rule == SpirvLayoutRule::RelaxedGLSLStd430 ||
-        rule == SpirvLayoutRule::FxcCTBuffer) {
-      alignmentCalc.alignUsingHLSLRelaxedLayout(fieldType, memberSize,
-                                                memberAlignment, &offset);
-    } else {
-      offset = roundToPow2(offset, memberAlignment);
-    }
+    const uint32_t previousFieldEnd =
+        previousField ? previousField->offset.getValue() +
+                            previousField->sizeInBytes.getValue()
+                      : 0;
+    setDefaultFieldOffsetAndSize(alignmentCalc, rule, previousFieldEnd,
+                                 currentField, &loweredField);
 
     // The vk::offset attribute takes precedence over all.
-    if (field.vkOffsetAttr) {
-      offset = field.vkOffsetAttr->getOffset();
+    if (currentField->vkOffsetAttr) {
+      loweredField.offset = currentField->vkOffsetAttr->getOffset();
+      return loweredField;
     }
+
     // The :packoffset() annotation takes precedence over normal layout
     // calculation.
-    else if (field.packOffsetAttr) {
-      const uint32_t packOffset = field.packOffsetAttr->Subcomponent * 16 +
-                                  field.packOffsetAttr->ComponentOffset * 4;
+    if (currentField->packOffsetAttr) {
+      const uint32_t offset = currentField->packOffsetAttr->Subcomponent * 16 +
+                              currentField->packOffsetAttr->ComponentOffset * 4;
       // Do minimal check to make sure the offset specified by packoffset does
       // not cause overlap.
-      if (packOffset < nextLoc) {
+      if (offset < previousFieldEnd) {
         emitError("packoffset caused overlap with previous members",
-                  field.packOffsetAttr->Loc);
-      } else {
-        offset = packOffset;
+                  currentField->packOffsetAttr->Loc);
       }
+
+      loweredField.offset = offset;
+      return loweredField;
     }
+
     // The :register(c#) annotation takes precedence over normal layout
     // calculation.
-    else if (field.registerC) {
-      offset = 16 * field.registerC->RegisterNumber;
+    if (currentField->registerC) {
+      const uint32_t offset = 16 * currentField->registerC->RegisterNumber;
       // Do minimal check to make sure the offset specified by :register(c#)
       // does not cause overlap.
-      if (offset < nextLoc) {
+      if (offset < previousFieldEnd) {
         emitError(
             "found offset overlap when processing register(c%0) assignment",
-            field.registerC->Loc)
-            << field.registerC->RegisterNumber;
+            currentField->registerC->Loc)
+            << currentField->registerC->RegisterNumber;
       }
-    }
 
-    // Each structure-type member must have an Offset Decoration.
-    loweredField.offset = offset;
-    loweredField.sizeInBytes = memberSize;
-    offset += memberSize;
-
-    // Each structure-type member that is a matrix or array-of-matrices must be
-    // decorated with
-    // * A MatrixStride decoration, and
-    // * one of the RowMajor or ColMajor Decorations.
-    if (const auto *arrayType = astContext.getAsConstantArrayType(fieldType)) {
-      // We have an array of matrices as a field, we need to decorate
-      // MatrixStride on the field. So skip possible arrays here.
-      fieldType = arrayType->getElementType();
+      loweredField.offset = offset;
+      return loweredField;
     }
 
-    // Non-floating point matrices are represented as arrays of vectors, and
-    // therefore ColMajor and RowMajor decorations should not be applied to
-    // them.
-    QualType elemType = {};
-    if (isMxNMatrix(fieldType, &elemType) && elemType->isFloatingType()) {
-      memberAlignment = memberSize = stride = 0;
-      std::tie(memberAlignment, memberSize) = alignmentCalc.getAlignmentAndSize(
-          fieldType, rule, /*isRowMajor*/ llvm::None, &stride);
+    return loweredField;
+  };
 
-      loweredField.matrixStride = stride;
-      loweredField.isRowMajor = isRowMajorMatrix(spvOptions, fieldType);
-    }
+  // First, check to see if any of the structure members had 'register(c#)'
+  // location semantics. If so, members that do not have the 'register(c#)'
+  // assignment should be allocated after the *highest explicit address*.
+  // Example:
+  // float x : register(c10);   // Offset = 160 (10 * 16)
+  // float y;                   // Offset = 164 (160 + 4)
+  // float z: register(c1);     // Offset = 16  (1  * 16)
+  //
+  // This step is only here to simplify the struct layout generation.
+  std::vector<const HybridStructType::FieldInfo *> sortedFields =
+      sortFields(fields);
 
-    loweredFields.push_back(loweredField);
+  // The resulting vector of fields with proper layout information.
+  // Second, build each field, and determine their actual offset in the
+  // structure.
+  llvm::SmallVector<StructType::FieldInfo, 4> loweredFields;
+  llvm::DenseMap<const HybridStructType::FieldInfo *, uint32_t> fieldToIndexMap;
+
+  for (size_t i = 0; i < sortedFields.size(); i++) {
+    const StructType::FieldInfo *previousField =
+        i > 0 ? &loweredFields.back() : nullptr;
+    const HybridStructType::FieldInfo *currentField = sortedFields[i];
+    const size_t field_index = loweredFields.size();
+
+    loweredFields.emplace_back(fieldVisitor(previousField, currentField));
+    fieldToIndexMap[sortedFields[i]] = field_index;
   }
 
   // Re-order the sorted fields back to their original order.
+  llvm::SmallVector<StructType::FieldInfo, 4> result;
   for (const auto &field : fields)
     result.push_back(loweredFields[fieldToIndexMap[&field]]);
-
   return result;
 }
 

tools/clang/lib/SPIRV/LowerTypeVisitor.h

@@ -90,6 +90,12 @@ private:
   populateLayoutInformation(llvm::ArrayRef<HybridStructType::FieldInfo> fields,
                             SpirvLayoutRule rule);
 
+  /// Create a clang::StructType::FieldInfo from HybridStructType::FieldInfo.
+  /// This function only considers the field as standalone.
+  /// Offset and layout constraint from the parent struct are not considered.
+  StructType::FieldInfo lowerField(const HybridStructType::FieldInfo *field,
+                                   SpirvLayoutRule rule);
+
 private:
   ASTContext &astContext;                /// AST context
   SpirvContext &spvContext;              /// SPIR-V context