[SPIRV] Handle bit fields with initializer lists (microsoft#6301)

The spir-v backend currently treats each bitfield as a separate member
in the spir-v type when processing an initializer list. This commit
changes that so that it matches the behaviour of the DXIL backend.

In particular, if `{bf1, bf2}` is used to initialize a struct with two
bitfield, then `bf1` and bf2` will be merged into a single field when
constructing the struct.

Fixes microsoft#5688

Author: s-perron

tools/clang/include/clang/SPIRV/AstTypeProbe.h

@@ -13,6 +13,7 @@
 #include "dxc/Support/SPIRVOptions.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Type.h"
+#include "clang/SPIRV/SpirvType.h"
 #include "clang/Sema/Sema.h"
 
 namespace clang {
@@ -383,6 +384,28 @@ bool isStructureContainingAnyKindOfBuffer(QualType type);
 /// types, or it's an array of scalar, vector, or matrix of numeric types.
 bool isScalarOrNonStructAggregateOfNumericalTypes(QualType type);
 
+/// Calls `operation` on for each field in the base and derives class defined by
+/// `recordType`. The `operation` will receive the AST type linked to the field,
+/// the SPIRV type linked to the field, and the index of the field in the final
+/// SPIR-V representation. This index of the field can vary from the AST
+/// field-index because bitfields are merged into a single field in the SPIR-V
+/// representation.
+///
+/// If `includeMerged` is true, `operation` will be called on the same spir-v
+/// field for each field it represents. For example, if a spir-v field holds the
+/// values for 3 bit-fields, `operation` will be called 3 times with the same
+/// `spirvFieldIndex`. The `bitfield` information in `field` will be different.
+///
+/// If false, `operation` will be called once on the first field in the merged
+/// field.
+///
+/// If the operation returns false, we stop processing fields.
+void forEachSpirvField(
+    const RecordType *recordType, const StructType *spirvType,
+    std::function<bool(size_t spirvFieldIndex, const QualType &fieldType,
+                       const StructType::FieldInfo &field)>
+        operation,
+    bool includeMerged = false);
 } // namespace spirv
 } // namespace clang
 

tools/clang/lib/SPIRV/AstTypeProbe.cpp

@@ -1548,5 +1548,46 @@ bool isScalarOrNonStructAggregateOfNumericalTypes(QualType type) {
   return false;
 }
 
+void forEachSpirvField(
+    const RecordType *recordType, const StructType *spirvType,
+    std::function<bool(size_t, const QualType &, const StructType::FieldInfo &)>
+        operation,
+    bool includeMerged) {
+  const auto *cxxDecl = recordType->getAsCXXRecordDecl();
+  const auto *recordDecl = recordType->getDecl();
+
+  // Iterate through the base class (one field per base class).
+  // Bases cannot be melded into 1 field like bitfields, simple iteration.
+  uint32_t lastConvertedIndex = 0;
+  size_t astFieldIndex = 0;
+  for (const auto &base : cxxDecl->bases()) {
+    const auto &type = base.getType();
+    const auto &spirvField = spirvType->getFields()[astFieldIndex];
+    if (!operation(spirvField.fieldIndex, type, spirvField)) {
+      return;
+    }
+    lastConvertedIndex = spirvField.fieldIndex;
+    ++astFieldIndex;
+  }
+
+  // Iterate through the derived class fields. Field could be merged.
+  for (const auto *field : recordDecl->fields()) {
+    const auto &spirvField = spirvType->getFields()[astFieldIndex];
+    const uint32_t currentFieldIndex = spirvField.fieldIndex;
+    if (!includeMerged && astFieldIndex > 0 &&
+        currentFieldIndex == lastConvertedIndex) {
+      ++astFieldIndex;
+      continue;
+    }
+
+    const auto &type = field->getType();
+    if (!operation(currentFieldIndex, type, spirvField)) {
+      return;
+    }
+    lastConvertedIndex = currentFieldIndex;
+    ++astFieldIndex;
+  }
+}
+
 } // namespace spirv
 } // namespace clang

tools/clang/lib/SPIRV/InitListHandler.cpp

@@ -17,12 +17,13 @@
 #include <algorithm>
 #include <iterator>
 
+#include "LowerTypeVisitor.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace clang {
 namespace spirv {
 
-InitListHandler::InitListHandler(const ASTContext &ctx, SpirvEmitter &emitter)
+InitListHandler::InitListHandler(ASTContext &ctx, SpirvEmitter &emitter)
     : astContext(ctx), theEmitter(emitter),
       spvBuilder(emitter.getSpirvBuilder()),
       diags(emitter.getDiagnosticsEngine()) {}
@@ -399,39 +400,54 @@ InitListHandler::createInitForStructType(QualType type, SourceLocation srcLoc,
     tryToSplitStruct();
   }
 
-  llvm::SmallVector<SpirvInstruction *, 4> fields;
-
-  // Initialize base classes first.
-  llvm::SmallVector<SpirvInstruction *, 4> base_fields;
-  const RecordDecl *structDecl = type->getAsStructureType()->getDecl();
-  if (auto *cxxStructDecl = dyn_cast<CXXRecordDecl>(structDecl)) {
-    for (CXXBaseSpecifier base : cxxStructDecl->bases()) {
-      QualType baseType = base.getType();
-      const RecordType *baseStructType = baseType->getAsStructureType();
-      if (baseStructType == nullptr) {
-        continue;
-      }
-      const RecordDecl *baseStructDecl = baseStructType->getDecl();
-      for (const auto *field : baseStructDecl->fields()) {
-        base_fields.push_back(
-            createInitForType(field->getType(), field->getLocation(), range));
-        if (!base_fields.back())
-          return nullptr;
-      }
-      fields.push_back(spvBuilder.createCompositeConstruct(
-          baseType, base_fields, srcLoc, range));
-      base_fields.clear();
-    }
-  }
+  const RecordType *recordType = type->getAs<RecordType>();
+  assert(recordType);
 
-  for (const auto *field : structDecl->fields()) {
-    fields.push_back(
-        createInitForType(field->getType(), field->getLocation(), range));
-    if (!fields.back())
-      return nullptr;
-  }
+  LowerTypeVisitor lowerTypeVisitor(astContext, theEmitter.getSpirvContext(),
+                                    theEmitter.getSpirvOptions());
+  const SpirvType *spirvType =
+      lowerTypeVisitor.lowerType(type, SpirvLayoutRule::Void, false, srcLoc);
 
-  // TODO: use OpConstantComposite when all components are constants
+  llvm::SmallVector<SpirvInstruction *, 4> fields;
+  const StructType *structType = dyn_cast<StructType>(spirvType);
+  assert(structType != nullptr);
+  forEachSpirvField(
+      recordType, structType,
+      [this, &fields, srcLoc, range](size_t spirvFieldIndex,
+                                     const QualType &fieldType,
+                                     const StructType::FieldInfo &fieldInfo) {
+        SpirvInstruction *init = createInitForType(fieldType, srcLoc, range);
+
+        // For non bit-fields, `init` will be the value for the component.
+        if (!fieldInfo.bitfield.hasValue()) {
+          assert(fields.size() == fieldInfo.fieldIndex);
+          fields.push_back(init);
+          return true;
+        }
+
+        // For a bit fields we need to insert it into the container.
+        // The first time we see this bit field, init is used as the value.
+        // This assumes that 0 is the first offset in the bitfield.
+        if (fields.size() <= fieldInfo.fieldIndex) {
+          assert(fieldInfo.bitfield->offsetInBits == 0);
+          fields.push_back(init);
+          return true;
+        }
+
+        // For the remaining bitfields, we need to insert them into the existing
+        // container, which is the last element in `fields`.
+        assert(fields.size() == fieldInfo.fieldIndex + 1);
+        SpirvInstruction *offset = spvBuilder.getConstantInt(
+            astContext.UnsignedIntTy,
+            llvm::APInt(32, fieldInfo.bitfield->offsetInBits));
+        SpirvInstruction *count = spvBuilder.getConstantInt(
+            astContext.UnsignedIntTy,
+            llvm::APInt(32, fieldInfo.bitfield->sizeInBits));
+        fields.back() = spvBuilder.createBitFieldInsert(
+            fieldType, fields.back(), init, offset, count, srcLoc);
+        return true;
+      },
+      true);
   return spvBuilder.createCompositeConstruct(type, fields, srcLoc, range);
 }
 

tools/clang/lib/SPIRV/InitListHandler.h

@@ -80,7 +80,7 @@ class InitListHandler {
   /// Constructs an InitListHandler which uses the given emitter for normal
   /// translation tasks. It will reuse the ModuleBuilder embedded in the given
   /// emitter.
-  explicit InitListHandler(const ASTContext &, SpirvEmitter &);
+  InitListHandler(ASTContext &ctx, SpirvEmitter &emitter);
 
   /// Processes the given InitListExpr and returns the <result-id> for the final
   /// SPIR-V value.
@@ -143,7 +143,7 @@ class InitListHandler {
                                                    SourceLocation);
 
 private:
-  const ASTContext &astContext;
+  ASTContext &astContext;
   SpirvEmitter &theEmitter;
   SpirvBuilder &spvBuilder;
   DiagnosticsEngine &diags;

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -564,64 +564,6 @@ const StructType *lowerStructType(const SpirvCodeGenOptions &spirvOptions,
   return output;
 }
 
-// Calls `operation` on for each field in the base and derives class defined by
-// `recordType`. The `operation` will receive the AST type linked to the field,
-// the SPIRV type linked to the field, and the index of the field in the final
-// SPIR-V representation. This index of the field can vary from the AST
-// field-index because bitfields are merged into a single field in the SPIR-V
-// representation.
-//
-// If `includeMerged` is true, `operation` will be called on the same spir-v
-// field for each field it represents. For example, if a spir-v field holds the
-// values for 3 bit-fields, `operation` will be called 3 times with the same
-// `spirvFieldIndex`. The `bitfield` information in `field` will be different.
-//
-// If false, `operation` will be called once on the first field in the merged
-// field.
-//
-// If the operation returns false, we stop processing fields.
-void forEachSpirvField(
-    const RecordType *recordType, const StructType *spirvType,
-    std::function<bool(size_t spirvFieldIndex, const QualType &fieldType,
-                       const StructType::FieldInfo &field)>
-        operation,
-    bool includeMerged = false) {
-  const auto *cxxDecl = recordType->getAsCXXRecordDecl();
-  const auto *recordDecl = recordType->getDecl();
-
-  // Iterate through the base class (one field per base class).
-  // Bases cannot be melded into 1 field like bitfields, simple iteration.
-  uint32_t lastConvertedIndex = 0;
-  size_t astFieldIndex = 0;
-  for (const auto &base : cxxDecl->bases()) {
-    const auto &type = base.getType();
-    const auto &spirvField = spirvType->getFields()[astFieldIndex];
-    if (!operation(spirvField.fieldIndex, type, spirvField)) {
-      return;
-    }
-    lastConvertedIndex = spirvField.fieldIndex;
-    ++astFieldIndex;
-  }
-
-  // Iterate through the derived class fields. Field could be merged.
-  for (const auto *field : recordDecl->fields()) {
-    const auto &spirvField = spirvType->getFields()[astFieldIndex];
-    const uint32_t currentFieldIndex = spirvField.fieldIndex;
-    if (!includeMerged && astFieldIndex > 0 &&
-        currentFieldIndex == lastConvertedIndex) {
-      ++astFieldIndex;
-      continue;
-    }
-
-    const auto &type = field->getType();
-    if (!operation(currentFieldIndex, type, spirvField)) {
-      return;
-    }
-    lastConvertedIndex = currentFieldIndex;
-    ++astFieldIndex;
-  }
-}
-
 } // namespace
 
 SpirvEmitter::SpirvEmitter(CompilerInstance &ci)

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -63,6 +63,7 @@ class SpirvEmitter : public ASTConsumer {
 
   ASTContext &getASTContext() { return astContext; }
   SpirvBuilder &getSpirvBuilder() { return spvBuilder; }
+  SpirvContext &getSpirvContext() { return spvContext; }
   DiagnosticsEngine &getDiagnosticsEngine() { return diags; }
   CompilerInstance &getCompilerInstance() { return theCompilerInstance; }
   SpirvCodeGenOptions &getSpirvOptions() { return spirvOptions; }

tools/clang/test/CodeGenSPIRV/var.init.struct.hlsl

@@ -41,6 +41,13 @@ struct W {
   float4 color;
 };
 
+struct BitFields {
+  uint R:8;
+  uint G:8;
+  uint B:8;
+  uint A:8;
+};
+
 void main() {
 // CHECK-LABEL: %bb_entry = OpLabel
 
@@ -103,4 +110,11 @@ void main() {
 // CHECK-NEXT: [[float4_zero:%[0-9]+]] = OpConvertSToF %v4float [[int4_zero]]
 // CHECK-NEXT:             {{%[0-9]+}} = OpCompositeConstruct %W [[float4_zero]]
     W w = { (0).xxxx };
+
+// CHECK: [[v1:%[0-9]+]] = OpBitFieldInsert %uint %uint_3 %uint_2 %uint_8 %uint_8
+// CHECK: [[v2:%[0-9]+]] = OpBitFieldInsert %uint [[v1]] %uint_1 %uint_16 %uint_8
+// CHECK: [[v3:%[0-9]+]] = OpBitFieldInsert %uint [[v2]] %uint_0 %uint_24 %uint_8
+// CHECK: [[bf:%[0-9]+]] = OpCompositeConstruct %BitFields [[v3]]
+// CHECK:                  OpStore %bf [[bf]]
+    BitFields bf = {3, 2, 1, 0};
 }