[GlobalOpt] Add range metadata to loads from constant global variables

llvm/lib/Transforms/IPO/GlobalOpt.cpp

@@ -45,6 +45,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
@@ -2498,6 +2499,236 @@ OptimizeGlobalAliases(Module &M,
  return Changed;
 }
 
+struct AccessPattern {
+ Type *Ty;
+
+ APInt Stride;
+ APInt Offset;
+};
+
+template <> struct DenseMapInfo<AccessPattern> {
+ static inline AccessPattern getEmptyKey() {
+ return {(Type *)1, APInt(), APInt()};
+ }
+ static inline AccessPattern getTombstoneKey() {
+ return {(Type *)2, APInt(), APInt()};
+ }
+ static unsigned getHashValue(const AccessPattern &AP) {
+ return hash_combine(AP.Ty, AP.Stride, AP.Offset);
+ }
+ static bool isEqual(const AccessPattern &LHS, const AccessPattern &RHS) {
+ return LHS.Ty == RHS.Ty && LHS.Stride == RHS.Stride &&
+ LHS.Offset == RHS.Offset;
+ }
+};
+
+// return (gcd, x, y) such that a*x + b*y = gcd
+std::tuple<APInt, APInt, APInt> ExtendedSignedGCD(APInt a, APInt b) {
+ unsigned BW = a.getBitWidth();
+ APInt x = APInt(BW, 1);
+ APInt y = APInt(BW, 0);
+ APInt x1 = APInt(BW, 0);
+ APInt y1 = APInt(BW, 1);
+
+ while (b != 0) {
+ APInt q = APInt(BW, 0);
+ APInt r = APInt(BW, 0);
+ APInt::sdivrem(a, b, q, r);
+ a = std::move(b);
+ b = std::move(r);
+
+ std::swap(x, x1);
+ std::swap(y, y1);
+ x1 -= q * x;
+ y1 -= q * y;
+ }
+ return {a, x, y};
+}
+
+// Build if possible a new pair of Stride and Offset that are part of the
+// original but are also aligned.
+std::optional<std::pair<APInt, APInt>>
+AlignStrideAndOffset(const APInt &Stride, const APInt &Offset,
+ const APInt &Align) {
+ // Here Offset * Align is added only to make sure Missing is positive or zero
+ APInt Missing = ((Offset * Align) - Offset).urem(Align);
+
+ // fast path for common case,
+ if (Missing == 0)
+ return {
+ {(Stride * Align).udiv(APIntOps::GreatestCommonDivisor(Stride, Align)),
+ Offset}};
+
+ auto [GCD, X, Y] = ExtendedSignedGCD(Stride, Align);
+ assert(APIntOps::GreatestCommonDivisor(Stride, Align) == GCD);
+ assert((X * Stride + Y * Align) == GCD);
+
+ if (Missing.urem(GCD) != 0) {
+ // The new Stride + Offset cannot be created because there is no elements in
+ // the original that would be properly aligned
+ return std::nullopt;
+ }
+
+ APInt StrideAlign = Stride * Align;
+ // X could be negative, so we need to use sdiv
+ // Here + Offset * Align is added only to make sure Missing is positive
+ APInt NumStride =
+ (((Missing * X).sdiv(GCD)) + (StrideAlign * Align)).urem(Align);
+
+ APInt NewStride = StrideAlign.udiv(GCD);
+ APInt NewOffset = (Offset + (NumStride * Stride)).urem(NewStride);
+ return {{std::move(NewStride), std::move(NewOffset)}};
+}
+
+static bool addRangeMetadata(Module &M) {
+ const DataLayout &DL = M.getDataLayout();
+ bool Changed = false;
+
+ for (GlobalValue &Global : M.global_values()) {
+
+ auto *GV = dyn_cast<GlobalVariable>(&Global);
+ if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
+ continue;
+
+ // To be able to go to the next GlobalVariable with a return
+ [&] {
+ unsigned IndexBW = DL.getIndexTypeSizeInBits(GV->getType());
+
+ struct PointerInfo {
+ Value *Ptr;
+
+ // Zero denotes not set
+ APInt Stride;
+ APInt Offset;
+ };
+
+ // GEPs only take one pointer operand, the one we will come from, so we
+ // dont need to do uniqueing during the DFS
+ SmallVector<PointerInfo> Stack;
+
+ // All loads of the global that this code can analyze grouped by access
+ // pattern. Loads with the same access pattern can access the same offsets
+ // in the global, so they can be treated the same.
+ SmallDenseMap<AccessPattern, SmallVector<LoadInst *>> LoadsByAccess;
+
+ Stack.push_back({GV, APInt(IndexBW, 0), APInt(IndexBW, 0)});
+
+ while (!Stack.empty()) {
+ PointerInfo Curr = Stack.pop_back_val();
+
+ if (!isa<GlobalVariable>(Curr.Ptr)) {
+ if (auto *LI = dyn_cast<LoadInst>(Curr.Ptr)) {
+
+ if (!LI->getType()->isIntegerTy())
+ continue;
+
+ if (LI->hasMetadata(LLVMContext::MD_range))
+ continue;
+
+ // This is an access at a fixed offset, I expect this is handled
+ // elsewhere so we skip it.
+ if (Curr.Stride == 0)
+ continue;
+
+ // This case is very rare, but what it means is that we
+ // dont know at compile-time what offsets into the Global arrays are
+ // safe to access with this load. So we give-up.
+ if (LI->getAlign() > GV->getAlign().valueOrOne())
+ continue;
+
+ auto NewStrideAndOffset =
+ AlignStrideAndOffset(Curr.Stride, Curr.Offset,
+ APInt(IndexBW, LI->getAlign().value()));
+
+ if (!NewStrideAndOffset) {
+ // This load cannot access an offset with the correct alignment
+ LI->replaceAllUsesWith(PoisonValue::get(LI->getType()));
+ continue;
+ }
+
+ AccessPattern AP{LI->getType(), NewStrideAndOffset->first,
+ NewStrideAndOffset->second};
+ assert(AP.Stride != 0);
+ LoadsByAccess[AP].push_back(LI);
+ continue;
+ }
+ auto *GEP = dyn_cast<GetElementPtrInst>(Curr.Ptr);
+ if (!GEP)
+ continue;
+
+ SmallMapVector<Value *, APInt, 4> VarOffsets;
+ if (!GEP->collectOffset(DL, IndexBW, VarOffsets, Curr.Offset))
+ continue;
+
+ for (auto [V, Scale] : VarOffsets) {
+
+ // Commented out because I dont understand why we would need this
+ // But it was part of getStrideAndModOffsetOfGEP
+ // // Only keep a power of two factor for non-inbounds
+ // if (!GEP->isInBounds())
+ // Scale =
+ // APInt::getOneBitSet(Scale.getBitWidth(),
+ // Scale.countr_zero());
+
+ if (Curr.Stride == 0)
+ Curr.Stride = Scale;
+ else
+ Curr.Stride = APIntOps::GreatestCommonDivisor(Curr.Stride, Scale);
+ }
+ }
+
+ for (User *U : Curr.Ptr->users()) {
+ if (isa<LoadInst, GetElementPtrInst>(U)) {
+ Curr.Ptr = U;
+ Stack.push_back(Curr);
+ }
+ }
+ }
+
+ for (auto [AP, Loads] : LoadsByAccess) {
+ {
+ APInt SMin = APInt::getSignedMaxValue(AP.Ty->getIntegerBitWidth());
+ APInt SMax = APInt::getSignedMinValue(AP.Ty->getIntegerBitWidth());
+
+ APInt LastValidOffset =
+ APInt(IndexBW, DL.getTypeAllocSize(GV->getValueType()) -
+ DL.getTypeStoreSize(AP.Ty));
+ for (APInt Offset = AP.Offset; Offset.ule(LastValidOffset);
+ Offset += AP.Stride) {
+ assert(Offset.isAligned(Loads[0]->getAlign()));
+ Constant *Cst = ConstantFoldLoadFromConstPtr(GV, AP.Ty, Offset, DL);
+
+ if (!isa_and_nonnull<ConstantInt>(Cst))
+ // Lambda captures of a struct binding is only available starting
+ // in C++20, so we skip to the next element with goto
+ goto NextGroup;
+
+ // MD_range is order agnostics
+ SMin = APIntOps::smin(SMin, Cst->getUniqueInteger());
+ SMax = APIntOps::smax(SMax, Cst->getUniqueInteger());
+ }
+
+ MDBuilder MDHelper(M.getContext());
+
+ Changed = true;
+ if (SMin == SMax) {
+ for (LoadInst *LI : Loads)
+ LI->replaceAllUsesWith(ConstantInt::get(AP.Ty, SMin));
+ } else {
+ // The Range is allowed to wrap
+ MDNode *RNode = MDHelper.createRange(SMin, SMax + 1);
+ for (LoadInst *LI : Loads)
+ LI->setMetadata(LLVMContext::MD_range, RNode);
+ }
+ }
+ NextGroup:
+ (void)0; // Label expect statements
+ }
+ }();
+ }
+ return Changed;
+}
+
 static Function *
 FindAtExitLibFunc(Module &M,
  function_ref<TargetLibraryInfo &(Function &)> GetTLI,
@@ -2887,6 +3118,10 @@ optimizeGlobalsInModule(Module &M, const DataLayout &DL,
  Changed |= LocalChange;
  }
 
+ // Add range metadata to loads from constant global variables based on the
+ // values that could be loaded from the variable
+ Changed |= addRangeMetadata(M);
+
  // TODO: Move all global ctors functions to the end of the module for code
  // layout.