[SPIR-V] Correctly use halfword-aligned byte address (microsoft#5594)

All loads and stores for ByteAddressBuffer were being aligned to the
nearest word. This is incorrect behavior for 16-bit types. In order to
implement this, I had to move some of the logic for calculating offsets
from compile time to runtime. For loads and stores of 32 or 64-bit types
with a literal byte address, the optimizer will be able to reduce this
down to nearly the same code as was being output before.

For 16-bit types, stores are slightly more convoluted (we need to load
the old value, shift and mask it, and then combine it with the new
value, in order to avoid clobbering any existing value). The optimizer
isn't currently able to optimize away unnecessary shifts and masks, but
my intuition is that it should be possible, in case we want to optimize
this more.

Fixes microsoft#5475

Author: cassiebeckley

tools/clang/lib/SPIRV/RawBufferMethods.cpp

@@ -23,100 +23,95 @@ class RawBufferHandler {
       : theEmitter(emitter), astContext(emitter.getASTContext()),
         spvBuilder(emitter.getSpirvBuilder()) {}
 
-  /// \brief Performs (RW)ByteAddressBuffer.Load<T>(address).
+  /// \brief Performs (RW)ByteAddressBuffer.Load<T>(byteAddress).
   /// (RW)ByteAddressBuffers are represented as structs with only one member
   /// which is a runtime array in SPIR-V. This method works by loading one or
   /// more uints, and performing necessary casts and composite constructions
-  /// to build the 'targetType'. The 'offset' parameter can be used for finer
-  /// grained load of bitwidths smaller than 32-bits. The layout rule for the
-  /// result will be `Void` because the value will be built and used internally
-  /// only. It does not have to match `buffer`.
+  /// to build the 'targetType'. The layout rule for the result will be `Void`
+  /// because the value will be built and used internally only. It does not have
+  /// to match `buffer`.
   ///
   /// Example:
-  /// targetType = uint16_t, address=0, offset=0
-  ///                 --> Load the first 16-bit uint starting at address 0.
-  /// targetType = uint16_t, address=0, offset=16
-  ///                 --> Load the second 16-bit uint starting at address 0.
-  SpirvInstruction *processTemplatedLoadFromBuffer(SpirvInstruction *buffer,
-                                                   SpirvInstruction *&index,
-                                                   const QualType targetType,
-                                                   uint32_t &bitOffset,
-                                                   SourceRange range = {});
+  /// targetType = uint16_t, byteAddress=0
+  ///                 --> Load the first 16-bit uint starting at byte address 0.
+  SpirvInstruction *processTemplatedLoadFromBuffer(
+      SpirvInstruction *buffer, SpirvInstruction *byteAddress,
+      const QualType targetType, SourceRange range = {});
 
   /// \brief Performs RWByteAddressBuffer.Store<T>(address, value).
   /// RWByteAddressBuffers are represented in SPIR-V as structs with only one
   /// member which is a runtime array of uints. This method works by decomposing
   /// the given |value| to reach numeric/bool types. Then performs necessary
   /// casts to uints and stores them in the underlying runtime array.
-  /// The |bitOffset| parameter can be used for finer-grained bit-offset
-  /// control.
   ///
   /// Example:
-  /// targetType = uint16_t, address=0, offset=0
+  /// targetType = uint16_t, address=0
   ///                 --> Store to the first 16-bit uint starting at address 0.
-  /// targetType = uint16_t, address=0, offset=16
-  ///                 --> Store to the second 16-bit uint starting at address 0.
   void processTemplatedStoreToBuffer(SpirvInstruction *value,
                                      SpirvInstruction *buffer,
-                                     SpirvInstruction *&index,
+                                     SpirvInstruction *&byteAddress,
                                      const QualType valueType,
-                                     uint32_t &bitOffset,
                                      SourceRange range = {});
 
 private:
-  SpirvInstruction *load16BitsAtBitOffset0(SpirvInstruction *buffer,
-                                           SpirvInstruction *&index,
-                                           QualType target16BitType,
-                                           uint32_t &bitOffset,
-                                           SourceRange range = {});
-
-  SpirvInstruction *load32BitsAtBitOffset0(SpirvInstruction *buffer,
-                                           SpirvInstruction *&index,
-                                           QualType target32BitType,
-                                           uint32_t &bitOffset,
-                                           SourceRange range = {});
-
-  SpirvInstruction *load64BitsAtBitOffset0(SpirvInstruction *buffer,
-                                           SpirvInstruction *&index,
-                                           QualType target64BitType,
-                                           uint32_t &bitOffset,
-                                           SourceRange range = {});
-
-  SpirvInstruction *load16BitsAtBitOffset16(SpirvInstruction *buffer,
-                                            SpirvInstruction *&index,
-                                            QualType target16BitType,
-                                            uint32_t &bitOffset,
-                                            SourceRange range = {});
+  class BufferAddress {
+  public:
+    BufferAddress(SpirvInstruction *&byteAddress, SpirvEmitter &emitter)
+        : byteAddress(byteAddress), wordIndex(),
+          spvBuilder(emitter.getSpirvBuilder()),
+          astContext(emitter.getASTContext()) {}
+    SpirvInstruction *getByteAddress();
+    SpirvInstruction *getWordIndex(SourceLocation loc, SourceRange range);
+
+    void incrementByteAddress(SpirvInstruction *width, SourceLocation loc,
+                              SourceRange range);
+    void incrementByteAddress(uint32_t width, SourceLocation loc,
+                              SourceRange range);
+
+    void incrementWordIndex(SourceLocation loc, SourceRange range);
+
+  private:
+    SpirvInstruction *byteAddress;
+    llvm::Optional<SpirvInstruction *> wordIndex;
+
+    SpirvBuilder &spvBuilder;
+    ASTContext &astContext;
+  };
 
-private:
-  void store16BitsAtBitOffset0(SpirvInstruction *value,
-                               SpirvInstruction *buffer,
-                               SpirvInstruction *&index,
-                               const QualType valueType,
+  SpirvInstruction *processTemplatedLoadFromBuffer(SpirvInstruction *buffer,
+                                                   BufferAddress &address,
+                                                   const QualType targetType,
+                                                   SourceRange range = {});
+  void processTemplatedStoreToBuffer(SpirvInstruction *value,
+                                     SpirvInstruction *buffer,
+                                     BufferAddress &address,
+                                     const QualType valueType,
+                                     SourceRange range = {});
+
+  SpirvInstruction *load16Bits(SpirvInstruction *buffer, BufferAddress &address,
+                               QualType target16BitType,
                                SourceRange range = {});
 
-  void store32BitsAtBitOffset0(SpirvInstruction *value,
-                               SpirvInstruction *buffer,
-                               SpirvInstruction *&index,
-                               const QualType valueType,
+  SpirvInstruction *load32Bits(SpirvInstruction *buffer, BufferAddress &address,
+                               QualType target32BitType,
                                SourceRange range = {});
 
-  void store64BitsAtBitOffset0(SpirvInstruction *value,
-                               SpirvInstruction *buffer,
-                               SpirvInstruction *&index,
-                               const QualType valueType,
+  SpirvInstruction *load64Bits(SpirvInstruction *buffer, BufferAddress &address,
+                               QualType target64BitType,
                                SourceRange range = {});
 
-  void store16BitsAtBitOffset16(SpirvInstruction *value,
-                                SpirvInstruction *buffer,
-                                SpirvInstruction *&index,
-                                const QualType valueType,
-                                SourceRange range = {});
+private:
+  void store16Bits(SpirvInstruction *value, SpirvInstruction *buffer,
+                   BufferAddress &address, const QualType valueType,
+                   SourceRange range = {});
+
+  void store32Bits(SpirvInstruction *value, SpirvInstruction *buffer,
+                   BufferAddress &address, const QualType valueType,
+                   SourceRange range = {});
 
-  void storeArrayOfScalars(std::deque<SpirvInstruction *> values,
-                           SpirvInstruction *buffer, SpirvInstruction *&index,
-                           const QualType valueType, uint32_t &bitOffset,
-                           SourceLocation, SourceRange range = {});
+  void store64Bits(SpirvInstruction *value, SpirvInstruction *buffer,
+                   BufferAddress &address, const QualType valueType,
+                   SourceRange range = {});
 
   /// \brief Serializes the given values into their components until a scalar or
   /// a struct has been reached. Returns the most basic type it reaches.

tools/clang/lib/SPIRV/RawBufferMethods.h

@@ -4361,33 +4361,35 @@ SpirvInstruction *SpirvEmitter::processByteAddressBufferLoadStore(
                      BuiltinType::UInt)
                : !expr->getType()->isSpecificBuiltinType(BuiltinType::UInt));
 
-  // Do a OpShiftRightLogical by 2 (divide by 4 to get aligned memory
-  // access). The AST always casts the address to unsinged integer, so shift
-  // by unsinged integer 2.
-  auto *constUint2 =
-      spvBuilder.getConstantInt(astContext.UnsignedIntTy, llvm::APInt(32, 2));
   const auto range = expr->getSourceRange();
-  SpirvInstruction *address = spvBuilder.createBinaryOp(
-      spv::Op::OpShiftRightLogical, addressType, byteAddress, constUint2,
-      expr->getExprLoc(), range);
 
   if (isTemplatedLoadOrStore) {
     // Templated load. Need to (potentially) perform more
     // loads/casts/composite-constructs.
-    uint32_t bitOffset = 0;
     if (doStore) {
       auto *values = doExpr(expr->getArg(1));
       RawBufferHandler(*this).processTemplatedStoreToBuffer(
-          values, objectInfo, address, expr->getArg(1)->getType(), bitOffset,
-          range);
+          values, objectInfo, byteAddress, expr->getArg(1)->getType(), range);
       return nullptr;
     } else {
       RawBufferHandler rawBufferHandler(*this);
       return rawBufferHandler.processTemplatedLoadFromBuffer(
-          objectInfo, address, expr->getType(), bitOffset, range);
+          objectInfo, byteAddress, expr->getType(), range);
     }
   }
 
+  // Do a OpShiftRightLogical by 2 (divide by 4 to get aligned memory
+  // access). The AST always casts the address to unsigned integer, so shift
+  // by unsigned integer 2.
+  auto *constUint2 =
+      spvBuilder.getConstantInt(astContext.UnsignedIntTy, llvm::APInt(32, 2));
+  SpirvInstruction *address = spvBuilder.createBinaryOp(
+      spv::Op::OpShiftRightLogical, addressType, byteAddress, constUint2,
+      expr->getExprLoc(), range);
+
+  // We might be able to reduce duplication by handling this with
+  // processTemplatedLoadFromBuffer
+
   // Perform access chain into the RWByteAddressBuffer.
   // First index must be zero (member 0 of the struct is a
   // runtimeArray). The second index passed to OpAccessChain should be