Iterate over the spir-v fields to handle bitfields (microsoft#6746)

The code that implements `RWByteAddressBuffer::Store` will iterate over
all of the fields in a struct to write each element in the struct.
However, it does not use the "Spir-V fields", which accounts for
multiple fields being packed into the same bitfield. This is fixed by
using the `forEachSpirvField` function to make sure that the bitfield
are correctly handled.

Fixes microsoft#6483

Author: s-perron

tools/clang/lib/SPIRV/RawBufferMethods.cpp

@@ -8,6 +8,7 @@
 
 #include "RawBufferMethods.h"
 #include "AlignmentSizeCalculator.h"
+#include "LowerTypeVisitor.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/CharUnits.h"
 #include "clang/AST/RecordLayout.h"
@@ -284,44 +285,48 @@ SpirvInstruction *RawBufferHandler::processTemplatedLoadFromBuffer(
   // aligned like their field with the largest alignment.
   // As a result, there might exist some padding after some struct members.
   if (const auto *structType = targetType->getAs<RecordType>()) {
-    const auto *decl = structType->getDecl();
+    LowerTypeVisitor lowerTypeVisitor(astContext, theEmitter.getSpirvContext(),
+                                      theEmitter.getSpirvOptions(), spvBuilder);
+    auto *decl = targetType->getAsTagDecl();
+    assert(decl && "Expected all structs to be tag decls.");
+    const StructType *spvType = dyn_cast<StructType>(lowerTypeVisitor.lowerType(
+        targetType, theEmitter.getSpirvOptions().sBufferLayoutRule, llvm::None,
+        decl->getLocation()));
     llvm::SmallVector<SpirvInstruction *, 4> loadedElems;
-    uint32_t fieldOffsetInBytes = 0;
-    uint32_t structAlignment = 0, structSize = 0, stride = 0;
-    std::tie(structAlignment, structSize) =
-        AlignmentSizeCalculator(astContext, theEmitter.getSpirvOptions())
-            .getAlignmentAndSize(targetType,
-                                 theEmitter.getSpirvOptions().sBufferLayoutRule,
-                                 llvm::None, &stride);
-    for (const auto *field : decl->fields()) {
-      AlignmentSizeCalculator alignmentCalc(astContext,
-                                            theEmitter.getSpirvOptions());
-      uint32_t fieldSize = 0, fieldAlignment = 0;
-      std::tie(fieldAlignment, fieldSize) = alignmentCalc.getAlignmentAndSize(
-          field->getType(), theEmitter.getSpirvOptions().sBufferLayoutRule,
-          /*isRowMajor*/ llvm::None, &stride);
-      fieldOffsetInBytes = roundToPow2(fieldOffsetInBytes, fieldAlignment);
-      auto *byteOffset = address.getByteAddress();
-      if (fieldOffsetInBytes != 0) {
-        byteOffset = spvBuilder.createBinaryOp(
-            spv::Op::OpIAdd, astContext.UnsignedIntTy, byteOffset,
-            spvBuilder.getConstantInt(astContext.UnsignedIntTy,
-                                      llvm::APInt(32, fieldOffsetInBytes)),
-            loc, range);
-      }
-
-      loadedElems.push_back(processTemplatedLoadFromBuffer(
-          buffer, byteOffset, field->getType(), range));
+    forEachSpirvField(
+        structType, spvType,
+        [this, &buffer, &address, range,
+         &loadedElems](size_t spirvFieldIndex, const QualType &fieldType,
+                       const auto &field) {
+          auto *baseOffset = address.getByteAddress();
+          if (field.offset.hasValue() && field.offset.getValue() != 0) {
+            const auto loc = buffer->getSourceLocation();
+            SpirvConstant *offset = spvBuilder.getConstantInt(
+                astContext.UnsignedIntTy,
+                llvm::APInt(32, field.offset.getValue()));
+            baseOffset = spvBuilder.createBinaryOp(
+                spv::Op::OpIAdd, astContext.UnsignedIntTy, baseOffset, offset,
+                loc, range);
+          }
 
-      fieldOffsetInBytes += fieldSize;
-    }
+          loadedElems.push_back(processTemplatedLoadFromBuffer(
+              buffer, baseOffset, fieldType, range));
+          return true;
+        });
 
     // After we're done with loading the entire struct, we need to update the
     // byteAddress (in case we are loading an array of structs).
     //
     // struct size = 34 bytes (34 / 8) = 4 full words (34 % 8) = 2 > 0,
     // therefore need to move to the next aligned address So the starting byte
     // offset after loading the entire struct is: 8 * (4 + 1) = 40
+    uint32_t structAlignment = 0, structSize = 0, stride = 0;
+    std::tie(structAlignment, structSize) =
+        AlignmentSizeCalculator(astContext, theEmitter.getSpirvOptions())
+            .getAlignmentAndSize(targetType,
+                                 theEmitter.getSpirvOptions().sBufferLayoutRule,
+                                 llvm::None, &stride);
+
     assert(structAlignment != 0);
     SpirvInstruction *structWidth = spvBuilder.getConstantInt(
         astContext.UnsignedIntTy,
@@ -577,7 +582,7 @@ void RawBufferHandler::processTemplatedStoreToBuffer(SpirvInstruction *value,
       return;
     default:
       theEmitter.emitError(
-          "templated load of ByteAddressBuffer is only implemented for "
+          "templated store of ByteAddressBuffer is only implemented for "
           "16, 32, and 64-bit types",
           loc);
       return;
@@ -604,40 +609,36 @@ void RawBufferHandler::processTemplatedStoreToBuffer(SpirvInstruction *value,
   // aligned like their field with the largest alignment.
   // As a result, there might exist some padding after some struct members.
   if (const auto *structType = valueType->getAs<RecordType>()) {
-    const auto *decl = structType->getDecl();
-    uint32_t fieldOffsetInBytes = 0;
-    uint32_t structAlignment = 0, structSize = 0, stride = 0;
-    std::tie(structAlignment, structSize) =
-        AlignmentSizeCalculator(astContext, theEmitter.getSpirvOptions())
-            .getAlignmentAndSize(valueType,
-                                 theEmitter.getSpirvOptions().sBufferLayoutRule,
-                                 llvm::None, &stride);
-    uint32_t fieldIndex = 0;
-    for (const auto *field : decl->fields()) {
-      AlignmentSizeCalculator alignmentCalc(astContext,
-                                            theEmitter.getSpirvOptions());
-      uint32_t fieldSize = 0, fieldAlignment = 0;
-      std::tie(fieldAlignment, fieldSize) = alignmentCalc.getAlignmentAndSize(
-          field->getType(), theEmitter.getSpirvOptions().sBufferLayoutRule,
-          /*isRowMajor*/ llvm::None, &stride);
-      fieldOffsetInBytes = roundToPow2(fieldOffsetInBytes, fieldAlignment);
-      auto *byteOffset = address.getByteAddress();
-      if (fieldOffsetInBytes != 0) {
-        byteOffset = spvBuilder.createBinaryOp(
-            spv::Op::OpIAdd, astContext.UnsignedIntTy, byteOffset,
-            spvBuilder.getConstantInt(astContext.UnsignedIntTy,
-                                      llvm::APInt(32, fieldOffsetInBytes)),
-            loc, range);
-      }
+    LowerTypeVisitor lowerTypeVisitor(astContext, theEmitter.getSpirvContext(),
+                                      theEmitter.getSpirvOptions(), spvBuilder);
+    auto *decl = valueType->getAsTagDecl();
+    assert(decl && "Expected all structs to be tag decls.");
+    const StructType *spvType = dyn_cast<StructType>(lowerTypeVisitor.lowerType(
+        valueType, theEmitter.getSpirvOptions().sBufferLayoutRule, llvm::None,
+        decl->getLocation()));
+    assert(spvType);
+    forEachSpirvField(
+        structType, spvType,
+        [this, &address, loc, range, buffer, value](size_t spirvFieldIndex,
+                                                    const QualType &fieldType,
+                                                    const auto &field) {
+          auto *baseOffset = address.getByteAddress();
+          if (field.offset.hasValue() && field.offset.getValue() != 0) {
+            SpirvConstant *offset = spvBuilder.getConstantInt(
+                astContext.UnsignedIntTy,
+                llvm::APInt(32, field.offset.getValue()));
+            baseOffset = spvBuilder.createBinaryOp(
+                spv::Op::OpIAdd, astContext.UnsignedIntTy, baseOffset, offset,
+                loc, range);
+          }
 
-      processTemplatedStoreToBuffer(
-          spvBuilder.createCompositeExtract(field->getType(), value,
-                                            {fieldIndex}, loc, range),
-          buffer, byteOffset, field->getType(), range);
-
-      fieldOffsetInBytes += fieldSize;
-      ++fieldIndex;
-    }
+          processTemplatedStoreToBuffer(
+              spvBuilder.createCompositeExtract(
+                  fieldType, value, {static_cast<uint32_t>(spirvFieldIndex)},
+                  loc, range),
+              buffer, baseOffset, fieldType, range);
+          return true;
+        });
 
     // After we're done with storing the entire struct, we need to update the
     // byteAddress (in case we are storing an array of structs).
@@ -647,6 +648,13 @@ void RawBufferHandler::processTemplatedStoreToBuffer(SpirvInstruction *value,
     // (34 % 8) = 2 > 0, therefore need to move to the next aligned address
     // So the starting byte offset after loading the entire struct is:
     // 8 * (4 + 1) = 40
+    uint32_t structAlignment = 0, structSize = 0, stride = 0;
+    std::tie(structAlignment, structSize) =
+        AlignmentSizeCalculator(astContext, theEmitter.getSpirvOptions())
+            .getAlignmentAndSize(valueType,
+                                 theEmitter.getSpirvOptions().sBufferLayoutRule,
+                                 llvm::None, &stride);
+
     assert(structAlignment != 0);
     auto *structWidth = spvBuilder.getConstantInt(
         astContext.UnsignedIntTy,

tools/clang/test/CodeGenSPIRV/method.byte-address-buffer.load.hlsl

@@ -2,6 +2,11 @@
 
 ByteAddressBuffer myBuffer;
 
+struct S {
+  uint32_t x : 8;
+  uint32_t y : 8;
+};
+
 [numthreads(1, 1, 1)]
 void main() {
   uint addr = 0;
@@ -50,4 +55,11 @@ void main() {
 // CHECK-NEXT: [[load4_word3:%[0-9]+]] = OpLoad %uint [[load_ptr6]]
 // CHECK-NEXT: {{%[0-9]+}} = OpCompositeConstruct %v4uint [[load4_word0]] [[load4_word1]] [[load4_word2]] [[load4_word3]]
   uint4 word4 = myBuffer.Load4(addr);
+
+// CHECK: [[idx:%[0-9]+]] = OpShiftRightLogical %uint %uint_0 %uint_2
+// CHECK: [[ac:%[0-9]+]] = OpAccessChain %_ptr_Uniform_uint %myBuffer %uint_0 [[idx]]
+// CHECK: [[bitfield:%[0-9]+]] = OpLoad %uint [[ac]]
+// CHECK: [[s:%[0-9]+]] = OpCompositeConstruct %S [[bitfield]]
+// CHECK: OpStore %s [[s]]
+  S s = myBuffer.Load<S>(0);
 }

tools/clang/test/CodeGenSPIRV/method.byte-address-buffer.store.hlsl

@@ -2,6 +2,12 @@
 
 RWByteAddressBuffer outBuffer;
 
+struct S
+{
+    uint32_t x:8;
+    uint32_t y:8;
+};
+
 [numthreads(1, 1, 1)]
 void main() {
   uint addr = 0;
@@ -67,4 +73,13 @@ void main() {
 // CHECK-NEXT: [[outBufPtr3:%[0-9]+]] = OpAccessChain %_ptr_Uniform_uint %outBuffer %uint_0 [[baseAddr_plus3]]
 // CHECK-NEXT: OpStore [[outBufPtr3]] [[word3]]
   outBuffer.Store4(addr, words4);
+
+// CHECK: [[s:%[0-9]+]] = OpLoad %S %s
+// CHECK: [[bitfield:%[0-9]+]] = OpCompositeExtract %uint [[s]] 0
+// CHECK: [[idx:%[0-9]+]] = OpShiftRightLogical %uint %uint_0 %uint_2
+// CHECK: [[ac:%[0-9]+]] = OpAccessChain %_ptr_Uniform_uint %outBuffer %uint_0 [[idx]]
+// CHECK: OpStore [[ac]] [[bitfield]]
+  S s  = (S)0;
+  s.x = 5;
+  outBuffer.Store(0, s);
 }