[llvm] [VectorCombine][AMDGPU] Shrink demanded vector loads (PR #202501)

[via llvm-commits]

================
@@ -2157,6 +2160,231 @@ bool VectorCombine::scalarizeLoadExtract(LoadInst *LI, VectorType *VecTy,
   return true;
 }
 
+/// Try to shrink vector loads feeding extractelement instructions when some
+/// of the loaded lanes are not demanded.
+bool VectorCombine::shrinkLoadForExtracts(Instruction &I) {
+  if (!DB)
+    return false;
+
+  auto *OldLoad = dyn_cast<LoadInst>(&I);
+  if (!OldLoad || !OldLoad->isSimple())
+    return false;
+
+  auto *OldLoadTy = dyn_cast<FixedVectorType>(OldLoad->getType());
+  if (!OldLoadTy)
+    return false;
+
+  Type *OldEltTy = OldLoadTy->getElementType();
+  bool IsIntegerElement = OldEltTy->isIntegerTy();
+  bool IsFloatingPointElement = OldEltTy->isFloatingPointTy();
+  if (!IsIntegerElement && !IsFloatingPointElement)
+    return false;
+
+  constexpr unsigned ByteWidth = 8;
+  TypeSize OldEltSizeInBits = DL->getTypeSizeInBits(OldEltTy);
+  if (OldEltSizeInBits.isScalable())
+    return false;
+  unsigned OldEltSize = OldEltSizeInBits.getFixedValue();
+  if (OldEltSize < ByteWidth || OldEltSize % ByteWidth != 0)
+    return false;
+
+  unsigned NumElts = OldLoadTy->getNumElements();
+  SmallVector<ExtractElementInst *, 8> Extracts;
+  SmallVector<APInt, 8> ExtractDemandedBits;
+
+  for (User *U : OldLoad->users()) {
+    auto *Ext = dyn_cast<ExtractElementInst>(U);
+    if (!Ext || Ext->use_empty())
+      return false;
+
+    auto *Idx = dyn_cast<ConstantInt>(Ext->getIndexOperand());
+    if (!Idx || Idx->getValue().uge(NumElts))
+      return false;
+
+    Extracts.push_back(Ext);
+    ExtractDemandedBits.push_back(DB->getDemandedBits(Ext));
+  }
+
+  TypeSize OldLoadSize = DL->getTypeStoreSize(OldLoadTy);
+  if (OldLoadSize.isScalable())
+    return false;
+
+  struct ChunkRun {
+    unsigned Start;
+    unsigned Length;
+  };
+  SmallVector<ChunkRun, 4> ChunkRuns;
+  SmallVector<Type *, 4> NewLoadTys;
+  SmallVector<uint64_t, 4> NewLoadOffsets;
+  Type *ChunkTy = nullptr;
+  unsigned ChunkBits = 0;
+  unsigned ChunksPerElt = 0;
+  unsigned ChunkBytes = 0;
+
+  // Try progressively smaller byte-sized power-of-two chunks. This keeps the
+  // transform target independent while allowing TTI to opt in to the chunk/load
+  // widths that are profitable for a target.
+  auto TryBuildShrinkLoadPlan = [&](unsigned CandidateChunkBits,
+                                    Type *CandidateChunkTy) {
+    if (OldEltSize % CandidateChunkBits != 0)
+      return false;
+    // Sub-element chunk numbering maps directly to byte offsets below, so only
+    // use it for little-endian targets. Whole-element shrink is endian-neutral.
+    if (CandidateChunkBits != OldEltSize && !DL->isLittleEndian())
+      return false;
+
+    unsigned CandidateChunksPerElt = OldEltSize / CandidateChunkBits;
+    unsigned NumChunks = NumElts * CandidateChunksPerElt;
+    APInt DemandedChunks = APInt::getZero(NumChunks);
+
+    for (auto [Ext, DemandedBits] : zip(Extracts, ExtractDemandedBits)) {
+      auto *Idx = cast<ConstantInt>(Ext->getIndexOperand());
+      unsigned EltIdx = Idx->getZExtValue();
+      for (unsigned Chunk = 0; Chunk != CandidateChunksPerElt; ++Chunk) {
+        if (!DemandedBits
+                 .extractBits(CandidateChunkBits, Chunk * CandidateChunkBits)
+                 .isZero())
+          DemandedChunks.setBit(EltIdx * CandidateChunksPerElt + Chunk);
+      }
+    }
+
+    if (DemandedChunks.isAllOnes())
+      return false;
+
+    SmallVector<ChunkRun, 4> CandidateChunkRuns;
+    for (unsigned I = 0; I != NumChunks;) {
+      if (!DemandedChunks[I]) {
+        ++I;
+        continue;
+      }
+
+      unsigned Start = I;
+      do {
+        ++I;
+      } while (I != NumChunks && DemandedChunks[I]);
+      CandidateChunkRuns.push_back({Start, I - Start});
+    }
+
+    if (CandidateChunkRuns.empty())
+      return false;
+
+    unsigned CandidateChunkBytes = CandidateChunkBits / ByteWidth;
+    SmallVector<Type *, 4> CandidateLoadTys;
+    SmallVector<uint64_t, 4> CandidateLoadOffsets;
+    TypeSize NewSize = TypeSize::getFixed(0);
+    for (const ChunkRun &Run : CandidateChunkRuns) {
+      Type *NewLoadTy = FixedVectorType::get(CandidateChunkTy, Run.Length);
+      CandidateLoadTys.push_back(NewLoadTy);
+      CandidateLoadOffsets.push_back(Run.Start * CandidateChunkBytes);
+      NewSize += DL->getTypeStoreSize(NewLoadTy);
+    }
+
+    if (NewSize >= OldLoadSize)
+      return false;
+
+    if (!TTI.shouldReduceLoadWidth(OldLoad, CandidateLoadTys,
+                                   CandidateLoadOffsets, CostKind))
+      return false;
+
+    ChunkBits = CandidateChunkBits;
+    ChunksPerElt = CandidateChunksPerElt;
+    ChunkBytes = CandidateChunkBytes;
+    ChunkTy = CandidateChunkTy;
+    ChunkRuns = std::move(CandidateChunkRuns);
+    NewLoadTys = std::move(CandidateLoadTys);
+    NewLoadOffsets = std::move(CandidateLoadOffsets);
+    return true;
+  };
+
+  bool FoundPlan = false;
+  if (IsIntegerElement) {
+    for (unsigned CandidateChunkBits = 1u << Log2_32(OldEltSize);
+         CandidateChunkBits >= ByteWidth; CandidateChunkBits >>= 1) {
+      if (TryBuildShrinkLoadPlan(
+              CandidateChunkBits,
+              IntegerType::get(I.getContext(), CandidateChunkBits))) {
+        FoundPlan = true;
+        break;
+      }
+      if (CandidateChunkBits == ByteWidth)
+        break;
+    }
+  } else {
+    // Floating-point elements cannot use DemandedBits to look inside the
+    // element, but they can still shrink at whole-element granularity.
+    FoundPlan = TryBuildShrinkLoadPlan(OldEltSize, OldEltTy);
+  }
+
+  if (!FoundPlan)
+    return false;
+
+  LLVM_DEBUG(dbgs() << "Found a vector load with undemanded elements: " << I
+                    << "\n");
+
+  Builder.SetInsertPoint(OldLoad);
+  Builder.SetCurrentDebugLocation(OldLoad->getDebugLoc());
+
+  SmallVector<LoadInst *, 4> NewLoads;
+  SmallVector<unsigned, 4> RunStarts;
+  for (auto [Idx, Run] : enumerate(ChunkRuns)) {
+    unsigned Start = Run.Start;
+    Type *NewLoadTy = NewLoadTys[Idx];
+    uint64_t NewLoadOffset = NewLoadOffsets[Idx];
+    Value *NewPtr =
+        Builder.CreateGEP(Builder.getInt8Ty(), OldLoad->getPointerOperand(),
+                          Builder.getInt64(NewLoadOffset));
+    auto *NewLoad = cast<LoadInst>(Builder.CreateAlignedLoad(
+        NewLoadTy, NewPtr,
+        commonAlignment(OldLoad->getAlign(), NewLoadOffset)));
+    copyMetadataForLoad(*NewLoad, *OldLoad);
+    AAMDNodes OldAAMD = OldLoad->getAAMetadata();
+    NewLoad->setAAMetadata(
+        OldAAMD.adjustForAccess(NewLoadOffset, NewLoadTy, *DL));
+    NewLoads.push_back(NewLoad);
+    RunStarts.push_back(Start);
+  }
+
+  auto ExtractSubElt = [&](unsigned SubElt, Instruction *InsertPt) -> Value * {
+    for (auto [I, NewLoad] : enumerate(NewLoads)) {
+      unsigned Start = RunStarts[I];
+      unsigned Length =
+          cast<FixedVectorType>(NewLoad->getType())->getNumElements();
+      if (SubElt < Start || SubElt >= Start + Length)
+        continue;
+
+      Builder.SetInsertPoint(InsertPt);
+      return Builder.CreateExtractElement(NewLoad,
+                                          Builder.getInt32(SubElt - Start));
+    }
+    return Constant::getNullValue(ChunkTy);
+  };
+
+  for (ExtractElementInst *Ext : Extracts) {
----------------
tianhbai wrote:

I rechecked the backend path, and I think this concern may mean my current `VectorCombine` approach is the wrong layer for this case.

The original IR shape is:

```llvm
%v = load <6 x i64>, ptr addrspace(4) %p, align 8
```

Only 9 of the 12 loaded dwords are used. My current VectorCombine prototype handles that by rewriting it roughly into three `<3 x i32>` loads, so it does change the load element type from `i64` to `i32` in the middle end.

After SelectionDAG legalization, the backend already has the shape that seems closer to the real issue:

```text
v4i32 load
extract_vector_elt 0
extract_vector_elt 1
extract_vector_elt 2
```

for each 128-bit chunk. AMDGPU then selects `GLOBAL_LOAD_DWORDX4` / `global_load_b128`, with `sub3` unused.

I also checked #122671. It looks related, but it only handles the single demanded element case in `SimplifyDemandedVectorElts` by scalarizing the vector load. This case needs the multi-extract case to shrink `v4i32` -> `v3i32`, so I do not think it is covered by that PR.

Given that, do you think this case should not be handled in VectorCombine? Would you prefer an AMDGPU backend fix instead, either as a SelectionDAG combine for the legalized `v4i32` load with constant extract users, or as a post-isel AMDGPU shrink from `GLOBAL_LOAD_DWORDX4` to `GLOBAL_LOAD_DWORDX3` when `sub3` is unused?

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