[MemoryBuiltins] Add getBaseObjectSize() (NFCI) (#155911)

getObjectSize() is based on ObjectSizeOffsetVisitor, which has become
very expensive over time. The implementation is geared towards computing
as-good-as-possible results for the objectsize intrinsics and similar.
However, we also use it in BasicAA, which is very hot, and really only
cares about the base cases like alloca/malloc/global, not any of the
analysis for GEPs, phis, or loads.

Add a new getBaseObjectSize() API for this use case, which only handles
the non-recursive cases. As a bonus, this API can easily return a
TypeSize and thus support scalable vectors. For now, I'm explicitly
discarding the scalable sizes in BasicAA just to avoid unnecessary
behavior changes during this refactor.

Author: nikic

llvm/include/llvm/Analysis/MemoryBuiltins.h

@@ -181,6 +181,14 @@ LLVM_ABI bool getObjectSize(const Value *Ptr, uint64_t &Size,
                             const DataLayout &DL, const TargetLibraryInfo *TLI,
                             ObjectSizeOpts Opts = {});
 
+/// Like getObjectSize(), but only returns the size of base objects (like
+/// allocas, global variables and allocator calls) and std::nullopt otherwise.
+/// Requires ExactSizeFromOffset mode.
+LLVM_ABI std::optional<TypeSize> getBaseObjectSize(const Value *Ptr,
+                                                   const DataLayout &DL,
+                                                   const TargetLibraryInfo *TLI,
+                                                   ObjectSizeOpts Opts = {});
+
 /// Try to turn a call to \@llvm.objectsize into an integer value of the given
 /// Type. Returns null on failure. If MustSucceed is true, this function will
 /// not return null, and may return conservative values governed by the second

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -103,12 +103,15 @@ static std::optional<TypeSize> getObjectSize(const Value *V,
                                              const TargetLibraryInfo &TLI,
                                              bool NullIsValidLoc,
                                              bool RoundToAlign = false) {
-  uint64_t Size;
   ObjectSizeOpts Opts;
   Opts.RoundToAlign = RoundToAlign;
   Opts.NullIsUnknownSize = NullIsValidLoc;
-  if (getObjectSize(V, Size, DL, &TLI, Opts))
-    return TypeSize::getFixed(Size);
+  if (std::optional<TypeSize> Size = getBaseObjectSize(V, DL, &TLI, Opts)) {
+    // FIXME: Remove this check, only exists to preserve previous behavior.
+    if (Size->isScalable())
+      return std::nullopt;
+    return Size;
+  }
   return std::nullopt;
 }
 

llvm/lib/Analysis/MemoryBuiltins.cpp

@@ -589,6 +589,59 @@ bool llvm::getObjectSize(const Value *Ptr, uint64_t &Size, const DataLayout &DL,
   return true;
 }
 
+std::optional<TypeSize> llvm::getBaseObjectSize(const Value *Ptr,
+                                                const DataLayout &DL,
+                                                const TargetLibraryInfo *TLI,
+                                                ObjectSizeOpts Opts) {
+  assert(Opts.EvalMode == ObjectSizeOpts::Mode::ExactSizeFromOffset &&
+         "Other modes are currently not supported");
+
+  auto Align = [&](TypeSize Size, MaybeAlign Alignment) {
+    if (Opts.RoundToAlign && Alignment && !Size.isScalable())
+      return TypeSize::getFixed(alignTo(Size.getFixedValue(), *Alignment));
+    return Size;
+  };
+
+  if (isa<UndefValue>(Ptr))
+    return TypeSize::getZero();
+
+  if (isa<ConstantPointerNull>(Ptr)) {
+    if (Opts.NullIsUnknownSize || Ptr->getType()->getPointerAddressSpace())
+      return std::nullopt;
+    return TypeSize::getZero();
+  }
+
+  if (auto *GV = dyn_cast<GlobalVariable>(Ptr)) {
+    if (!GV->getValueType()->isSized() || GV->hasExternalWeakLinkage() ||
+        !GV->hasInitializer() || GV->isInterposable())
+      return std::nullopt;
+    return Align(DL.getTypeAllocSize(GV->getValueType()), GV->getAlign());
+  }
+
+  if (auto *A = dyn_cast<Argument>(Ptr)) {
+    Type *MemoryTy = A->getPointeeInMemoryValueType();
+    if (!MemoryTy || !MemoryTy->isSized())
+      return std::nullopt;
+    return Align(DL.getTypeAllocSize(MemoryTy), A->getParamAlign());
+  }
+
+  if (auto *AI = dyn_cast<AllocaInst>(Ptr)) {
+    if (std::optional<TypeSize> Size = AI->getAllocationSize(DL))
+      return Align(*Size, AI->getAlign());
+    return std::nullopt;
+  }
+
+  if (auto *CB = dyn_cast<CallBase>(Ptr)) {
+    if (std::optional<APInt> Size = getAllocSize(CB, TLI)) {
+      if (std::optional<uint64_t> ZExtSize = Size->tryZExtValue())
+        return TypeSize::getFixed(*ZExtSize);
+    }
+    return std::nullopt;
+  }
+
+  return std::nullopt;
+}
+
 Value *llvm::lowerObjectSizeCall(IntrinsicInst *ObjectSize,
                                  const DataLayout &DL,
                                  const TargetLibraryInfo *TLI,