[SPIR-V] HLSL2021: initial bitfield implementation (#4831)

This commit adds the logic in the SPIR-V backend to generate proper
bitfields. Bitfield are packed using a first-fit method, linearly
packing them, but not mixing types. Goal is to follow C/C++ rules.

Bitfield merging was initially guessed from the offset stored in the FieldInfo.
This offset is not always available and has a very specific meaning.
When the struct is a function local variable, the layout rule is Void,
meaning we shouldn't assume any kind of byte offset, but rely on
construct index.
This commit adds a fieldIndex member to the FieldInfo struct, and this
field is used to determine if 2 fields are merged.

When doing a buffer texture load, the struct must be extracted from a
vector type, and rebuilt. This commit adds support for bitfield extraction for such
types. Fixing this helped me see scalar assignment were also failling in
some cases. Addressing bitfield extraction/insertion issues on with
commit.

Signed-off-by: Nathan Gauër <[email protected]>
Co-authored-by: Cassandra Beckley <[email protected]>

Author: Keenuts

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -174,20 +174,21 @@ class SpirvBuilder {
                                           SourceLocation loc,
                                           SourceRange range = {});
 
-  /// \brief Creates a load instruction loading the value of the given
-  /// <result-type> from the given pointer. Returns the instruction pointer for
-  /// the loaded value.
-  SpirvLoad *createLoad(QualType resultType, SpirvInstruction *pointer,
-                        SourceLocation loc, SourceRange range = {});
+  /// \brief Creates a load sequence loading the value of the given
+  /// <result-type> from the given pointer (load + optional extraction,
+  /// ex:bitfield). Returns the instruction pointer for the loaded value.
+  SpirvInstruction *createLoad(QualType resultType, SpirvInstruction *pointer,
+                               SourceLocation loc, SourceRange range = {});
   SpirvLoad *createLoad(const SpirvType *resultType, SpirvInstruction *pointer,
                         SourceLocation loc, SourceRange range = {});
 
   /// \brief Creates an OpCopyObject instruction from the given pointer.
   SpirvCopyObject *createCopyObject(QualType resultType,
                                     SpirvInstruction *pointer, SourceLocation);
 
-  /// \brief Creates a store instruction storing the given value into the given
+  /// \brief Creates a store sequence storing the given value into the given
   /// address. Returns the instruction pointer for the store instruction.
+  /// This function handles storing to bitfields.
   SpirvStore *createStore(SpirvInstruction *address, SpirvInstruction *value,
                    SourceLocation loc, SourceRange range = {});
 

tools/clang/include/clang/SPIRV/SpirvInstruction.h

@@ -14,6 +14,7 @@
 #include "clang/AST/APValue.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/SPIRV/SpirvType.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Optional.h"
@@ -203,6 +204,7 @@ class SpirvInstruction {
 
   void setRValue(bool rvalue = true) { isRValue_ = rvalue; }
   bool isRValue() const { return isRValue_; }
+  bool isLValue() const { return !isRValue_; }
 
   void setRelaxedPrecision() { isRelaxedPrecision_ = true; }
   bool isRelaxedPrecision() const { return isRelaxedPrecision_; }
@@ -213,6 +215,9 @@ class SpirvInstruction {
   void setPrecise(bool p = true) { isPrecise_ = p; }
   bool isPrecise() const { return isPrecise_; }
 
+  void setBitfieldInfo(const BitfieldInfo &info) { bitfieldInfo = info; }
+  llvm::Optional<BitfieldInfo> getBitfieldInfo() const { return bitfieldInfo; }
+
   /// Legalization-specific code
   ///
   /// Note: the following two functions are currently needed in order to support
@@ -255,6 +260,7 @@ class SpirvInstruction {
   bool isRelaxedPrecision_;
   bool isNonUniform_;
   bool isPrecise_;
+  llvm::Optional<BitfieldInfo> bitfieldInfo;
 };
 
 /// \brief OpCapability instruction

tools/clang/include/clang/SPIRV/SpirvType.h

@@ -30,6 +30,13 @@ enum class StructInterfaceType : uint32_t {
   UniformBuffer = 2,
 };
 
+struct BitfieldInfo {
+  // Offset of the bitfield, in bits, from the basetype start.
+  uint32_t offsetInBits;
+  // Size of the bitfield, in bits.
+  uint32_t sizeInBits;
+};
+
 class SpirvType {
 public:
   enum Kind {
@@ -290,14 +297,16 @@ class StructType : public SpirvType {
 public:
   struct FieldInfo {
   public:
-    FieldInfo(const SpirvType *type_, llvm::StringRef name_ = "",
+    FieldInfo(const SpirvType *type_, uint32_t fieldIndex_,
+              llvm::StringRef name_ = "",
               llvm::Optional<uint32_t> offset_ = llvm::None,
               llvm::Optional<uint32_t> matrixStride_ = llvm::None,
               llvm::Optional<bool> isRowMajor_ = llvm::None,
               bool relaxedPrecision = false, bool precise = false)
-        : type(type_), name(name_), offset(offset_), sizeInBytes(llvm::None),
-          matrixStride(matrixStride_), isRowMajor(isRowMajor_),
-          isRelaxedPrecision(relaxedPrecision), isPrecise(precise) {
+        : type(type_), fieldIndex(fieldIndex_), name(name_), offset(offset_),
+          sizeInBytes(llvm::None), matrixStride(matrixStride_),
+          isRowMajor(isRowMajor_), isRelaxedPrecision(relaxedPrecision),
+          isPrecise(precise) {
       // A StructType may not contain any hybrid types.
       assert(!isa<HybridType>(type_));
     }
@@ -306,6 +315,10 @@ class StructType : public SpirvType {
 
     // The field's type.
     const SpirvType *type;
+    // The index of this field in the composite construct.
+    // When the struct contains bitfields, StructType index and construct index
+    // can diverge as we merge bitfields together.
+    uint32_t fieldIndex;
     // The field's name.
     std::string name;
     // The integer offset in bytes for this field.
@@ -320,6 +333,8 @@ class StructType : public SpirvType {
     bool isRelaxedPrecision;
     // Whether this field is marked as 'precise'.
     bool isPrecise;
+    // Information about the bitfield (if applicable).
+    llvm::Optional<BitfieldInfo> bitfield;
   };
 
   StructType(
@@ -467,9 +482,11 @@ class HybridStructType : public HybridType {
               clang::VKOffsetAttr *offset = nullptr,
               hlsl::ConstantPacking *packOffset = nullptr,
               const hlsl::RegisterAssignment *regC = nullptr,
-              bool precise = false)
+              bool precise = false,
+              llvm::Optional<BitfieldInfo> bitfield = llvm::None)
         : astType(astType_), name(name_), vkOffsetAttr(offset),
-          packOffsetAttr(packOffset), registerC(regC), isPrecise(precise) {}
+          packOffsetAttr(packOffset), registerC(regC), isPrecise(precise),
+          bitfield(std::move(bitfield)) {}
 
     // The field's type.
     QualType astType;
@@ -483,6 +500,9 @@ class HybridStructType : public HybridType {
     const hlsl::RegisterAssignment *registerC;
     // Whether this field is marked as 'precise'.
     bool isPrecise;
+    // Whether this field is a bitfield or not. If set to false, bitfield width
+    // value is undefined.
+    llvm::Optional<BitfieldInfo> bitfield;
   };
 
   HybridStructType(

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -4292,8 +4292,9 @@ void DeclResultIdMapper::storeOutStageVarsToStorage(
     }
     auto *ptrToOutputStageVar = spvBuilder.createAccessChain(
         outputControlPointType, found->second, {ctrlPointID}, /*loc=*/{});
-    auto *load = spvBuilder.createLoad(outputControlPointType,
-                                       ptrToOutputStageVar, /*loc=*/{});
+    auto *load =
+        spvBuilder.createLoad(outputControlPointType, ptrToOutputStageVar,
+                              /*loc=*/{});
     spvBuilder.createStore(ptr, load, /*loc=*/{});
     return;
   }

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -23,6 +23,8 @@
 #include "clang/SPIRV/String.h"
 // clang-format on
 
+#include <functional>
+
 namespace clang {
 namespace spirv {
 
@@ -2326,6 +2328,17 @@ EmitTypeHandler::getOrCreateConstantComposite(SpirvConstantComposite *inst) {
   return inst->getResultId();
 }
 
+static inline bool
+isFieldMergeWithPrevious(const StructType::FieldInfo &previous,
+                         const StructType::FieldInfo &field) {
+  if (previous.fieldIndex == field.fieldIndex) {
+    // Right now, the only reason for those indices to be shared is if both
+    // are merged bitfields.
+    assert(previous.bitfield.hasValue() && field.bitfield.hasValue());
+  }
+  return previous.fieldIndex == field.fieldIndex;
+}
+
 uint32_t EmitTypeHandler::emitType(const SpirvType *type) {
   // First get the decorations that would apply to this type.
   bool alreadyExists = false;
@@ -2447,24 +2460,32 @@ uint32_t EmitTypeHandler::emitType(const SpirvType *type) {
   }
   // Structure types
   else if (const auto *structType = dyn_cast<StructType>(type)) {
-    llvm::ArrayRef<StructType::FieldInfo> fields = structType->getFields();
-    size_t numFields = fields.size();
+    std::vector<std::reference_wrapper<const StructType::FieldInfo>>
+        fieldsToGenerate;
+    {
+      llvm::ArrayRef<StructType::FieldInfo> fields = structType->getFields();
+      for (size_t i = 0; i < fields.size(); ++i) {
+        if (i > 0 && isFieldMergeWithPrevious(fields[i - 1], fields[i]))
+          continue;
+        fieldsToGenerate.push_back(std::ref(fields[i]));
+      }
+    }
 
     // Emit OpMemberName for the struct members.
-    for (size_t i = 0; i < numFields; ++i)
-      emitNameForType(fields[i].name, id, i);
+    for (size_t i = 0; i < fieldsToGenerate.size(); ++i)
+      emitNameForType(fieldsToGenerate[i].get().name, id, i);
 
     llvm::SmallVector<uint32_t, 4> fieldTypeIds;
-    for (auto &field : fields) {
-      fieldTypeIds.push_back(emitType(field.type));
-    }
+    for (auto &field : fieldsToGenerate)
+      fieldTypeIds.push_back(emitType(field.get().type));
 
-    for (size_t i = 0; i < numFields; ++i) {
-      auto &field = fields[i];
+    for (size_t i = 0; i < fieldsToGenerate.size(); ++i) {
+      const auto &field = fieldsToGenerate[i].get();
       // Offset decorations
-      if (field.offset.hasValue())
+      if (field.offset.hasValue()) {
         emitDecoration(id, spv::Decoration::Offset, {field.offset.getValue()},
                        i);
+      }
 
       // MatrixStride decorations
       if (field.matrixStride.hasValue())

tools/clang/lib/SPIRV/LiteralTypeVisitor.cpp

@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "LiteralTypeVisitor.h"
+#include "LowerTypeVisitor.h"
 #include "clang/SPIRV/AstTypeProbe.h"
 #include "clang/SPIRV/SpirvFunction.h"
 
@@ -389,6 +390,11 @@ bool LiteralTypeVisitor::updateTypeForCompositeMembers(
       const auto *decl = structType->getDecl();
       size_t i = 0;
       for (const auto *field : decl->fields()) {
+        // If the field is a bitfield, it might be squashed later when building
+        // the SPIR-V type depending on context. This means indices starting
+        // from this bitfield are not guaranteed, and we shouldn't touch them.
+        if (field->isBitField())
+          break;
         tryToUpdateInstLitType(constituents[i], field->getType());
         ++i;
       }