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

[Matt Arsenault 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) {
----------------
arsenm wrote:

If you somehow extracted only a subset of the element, that implies there's a cast or extract somewhere that would separately fold into the load. It wouldn't be directly handled as part of this demanded element reduction 

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