[llvm-branch-commits] [llvm] [AMDGPU] Add support for sub-dword types for address space 13 (PR #213223)

[Tim Gymnich via llvm-branch-commits]

================
@@ -3493,23 +3521,118 @@ static bool lowerLoadStoreVGPR(LegalizerHelper &Helper, MachineInstr &MI) {
   // integer types rather than plain scalars.
   const LLT I32 = LLT::integer(32);
 
-  // Only whole-dword, non-extending/non-truncating accesses are implemented.
-  // Reject anything else with a diagnostic instead of failing to legalize
-  // (sub-dword support lands in a later change).
-  if (!isVGPRLoadStoreSizeSupported(MMO.getMemoryType().getSizeInBits(),
-                                    ValSize)) {
+  // Whole-dword and naturally aligned 8-/16-bit accesses are implemented.
+  // Reject anything else with a diagnostic instead of failing to legalize.
+  if (!isVGPRLoadStoreSupported(MemSize, ValSize, MMO.getAlign())) {
     const Function &F = B.getMF().getFunction();
     F.getContext().diagnose(DiagnosticInfoUnsupported(
         F,
         "unsupported access of VGPR 'as memory' address space (13); only "
-        "whole-dword loads and stores are implemented",
+        "whole-dword and naturally aligned 8-/16-bit loads and stores are "
+        "implemented",
         MI.getDebugLoc()));
     if (!IsStore)
       B.buildUndef(ValReg);
     MI.eraseFromParent();
     return true;
   }
 
+  // Handle bytes and aligned shorts. These become a bit-field extract out of
+  // the containing dword (loads), or a read-modify-write of it (stores); see
+  // AMDGPULowerIdxOps.
+  if (MemSize < 32) {
+    assert(MemSize == 8 || MemSize == 16);
+    assert(MemSize <= ValSize && ValSize <= 32);
+
+    // Determine the bit-offset, optimizing the case where the LSBs are
+    // known constant.
+    Register BaseReg = PtrReg;
+    int64_t Offset = 0;
+    if (auto *PtrAdd = getOpcodeDef<GPtrAdd>(PtrReg, MRI)) {
+      if (auto MaybeOff =
+              getIConstantVRegValWithLookThrough(PtrAdd->getOffsetReg(), MRI)) {
+        BaseReg = PtrAdd->getBaseReg();
+        Offset = MaybeOff->Value.getSExtValue();
+      }
+    }
+
+    bool HaveConstantBitOffset = false;
+    int64_t ConstantBitOffsetVal = 0;
+    if (Offset == 0 && MMO.getAlign() >= Align(4)) {
+      HaveConstantBitOffset = true;
+      ConstantBitOffsetVal = 0;
+    } else {
+      auto &VT = *Helper.getValueTracking();
+      KnownBits BaseKB = VT.getKnownBits(BaseReg).trunc(2);
+      if (BaseKB.isConstant()) {
+        Offset += BaseKB.getConstant().getZExtValue();
+        HaveConstantBitOffset = true;
+        ConstantBitOffsetVal = (Offset & 3) * 8;
+      }
+    }
+
+    // Setup common registers.
+    const auto PtrAsInt = B.buildPtrToInt(I32, PtrReg);
+    auto Two = B.buildConstant(I32, 2);
+    const auto Index = B.buildLShr(I32, PtrAsInt, Two);
+
+    const auto BitWidthReg = B.buildConstant(I32, MemSize);
----------------
tgymnich wrote:

try to avoid `auto` here and below

https://github.com/llvm/llvm-project/pull/213223