[InferAlignment][NFC] Unify Load/Store handling in tryToImproveAlign

[ParkHanbum]

Removes code duplication in tryToImproveAlign by unifying load and
store instruction handling with getLoadStore helper functions.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-ir

Author: hanbeom (ParkHanbum)

Changes

This patch do some misc refactoring in InferAlignment.
-. Removes code duplication in tryToImproveAlign by unifying load and
store instruction handling with getLoadStore helper functions.
-. Removed the code to get the alignment via KnownBits as it was
redundant with the implementation of getOrEnforceKnownAlignment.
-. Adds needEnforceAlignment that use ValueToBasePtr cache to
avoid redundant calculations for previously processed pointers in
tryToImproveAlign.

---

Full diff: https://github.com/llvm/llvm-project/pull/112699.diff

2 Files Affected:

• (modified) llvm/include/llvm/IR/Instructions.h (+10)
• (modified) llvm/lib/Transforms/Scalar/InferAlignment.cpp (+42-21)

diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index e89739a5552662..3772c57ef83b04 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -4952,6 +4952,16 @@ inline Align getLoadStoreAlignment(const Value *I) {
   return cast<StoreInst>(I)->getAlign();
 }
 
+/// A helper function that set the alignment of load or store instruction.
+inline void setLoadStoreAlignment(Value *I, Align NewAlign) {
+  assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
+         "Expected Load or Store instruction");
+  if (auto *LI = dyn_cast<LoadInst>(I))
+    LI->setAlignment(NewAlign);
+  else
+    cast<StoreInst>(I)->setAlignment(NewAlign);
+}
+
 /// A helper function that returns the address space of the pointer operand of
 /// load or store instruction.
 inline unsigned getLoadStoreAddressSpace(const Value *I) {
diff --git a/llvm/lib/Transforms/Scalar/InferAlignment.cpp b/llvm/lib/Transforms/Scalar/InferAlignment.cpp
index 6e0c206bd19805..ac8e6ae943c9f6 100644
--- a/llvm/lib/Transforms/Scalar/InferAlignment.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAlignment.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/InferAlignment.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Instructions.h"
@@ -22,31 +23,47 @@
 
 using namespace llvm;
 
+DenseMap<Value *, Value *> ValueToBasePtr;
+
 static bool tryToImproveAlign(
     const DataLayout &DL, Instruction *I,
     function_ref<Align(Value *PtrOp, Align OldAlign, Align PrefAlign)> Fn) {
-  if (auto *LI = dyn_cast<LoadInst>(I)) {
-    Value *PtrOp = LI->getPointerOperand();
-    Align OldAlign = LI->getAlign();
-    Align NewAlign = Fn(PtrOp, OldAlign, DL.getPrefTypeAlign(LI->getType()));
-    if (NewAlign > OldAlign) {
-      LI->setAlignment(NewAlign);
-      return true;
-    }
-  } else if (auto *SI = dyn_cast<StoreInst>(I)) {
-    Value *PtrOp = SI->getPointerOperand();
-    Value *ValOp = SI->getValueOperand();
-    Align OldAlign = SI->getAlign();
-    Align NewAlign = Fn(PtrOp, OldAlign, DL.getPrefTypeAlign(ValOp->getType()));
+
+  if (auto *PtrOp = getLoadStorePointerOperand(I)) {
+    Align OldAlign = getLoadStoreAlignment(I);
+    Align PrefAlign = DL.getPrefTypeAlign(getLoadStoreType(I));
+
+    Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign);
     if (NewAlign > OldAlign) {
-      SI->setAlignment(NewAlign);
+      setLoadStoreAlignment(I, NewAlign);
       return true;
     }
   }
+
   // TODO: Also handle memory intrinsics.
   return false;
 }
 
+static bool needEnforceAlignment(Value *PtrOp, Instruction *I, Align PrefAlign,
+                                 const DataLayout &DL) {
+  auto it = ValueToBasePtr.find(PtrOp);
+  if (it != ValueToBasePtr.end()) {
+    Value *V = it->second;
+    Align CurrentAlign;
+    if (auto Alloca = dyn_cast<AllocaInst>(V))
+      CurrentAlign = Alloca->getAlign();
+    if (auto GO = dyn_cast<GlobalObject>(V))
+      CurrentAlign = GO->getPointerAlignment(DL);
+
+    if (PrefAlign <= CurrentAlign) {
+      setLoadStoreAlignment(I, CurrentAlign);
+      return false;
+    }
+  }
+
+  return true;
+}
+
 bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
   const DataLayout &DL = F.getDataLayout();
   bool Changed = false;
@@ -57,9 +74,15 @@ bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
     for (Instruction &I : BB) {
       Changed |= tryToImproveAlign(
           DL, &I, [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
-            if (PrefAlign > OldAlign)
-              return std::max(OldAlign,
-                              tryEnforceAlignment(PtrOp, PrefAlign, DL));
+            if (needEnforceAlignment(PtrOp, &I, PrefAlign, DL) &&
+                PrefAlign > OldAlign) {
+              Align NewAlign = tryEnforceAlignment(PtrOp, PrefAlign, DL);
+              if (NewAlign > OldAlign) {
+                ValueToBasePtr[PtrOp] = PtrOp->stripPointerCasts();
+                return NewAlign;
+              }
+            }
+
             return OldAlign;
           });
     }
@@ -70,10 +93,8 @@ bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
     for (Instruction &I : BB) {
       Changed |= tryToImproveAlign(
           DL, &I, [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
-            KnownBits Known = computeKnownBits(PtrOp, DL, 0, &AC, &I, &DT);
-            unsigned TrailZ = std::min(Known.countMinTrailingZeros(),
-                                       +Value::MaxAlignmentExponent);
-            return Align(1ull << std::min(Known.getBitWidth() - 1, TrailZ));
+            return getOrEnforceKnownAlignment(PtrOp, MaybeAlign(), DL, &I, &AC,
+                                              &DT);
           });
     }
   }

[nikic]

The first commit looks fine to me, I don't understand the purpose of the other two. The alignment enforcement isn't expensive enough to make caching useful, I think?

[nikic]

You can't use globals this way in LLVM, it's not thread safe and will mix data across functions.

[ParkHanbum]

@nikic I assumed it would be lower than the cost of stripPointerCasts; if not, I'll delete that commit.
I'll also delete the second commit if it's have no effective.

[nikic]

@nikic I assumed it would be lower than the cost of stripPointerCasts; if not, I'll delete that commit. I'll also delete the second commit if it's have no effective.

stripPointerCasts() is very cheap.

[ParkHanbum]

@nikic I made a wrong assumption, sorry

[nikic]

LGTM