Add caching to reduce finding its base

Author: ParkHanbum

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,6 +23,8 @@
 
 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) {
@@ -36,10 +39,31 @@ static bool tryToImproveAlign(
       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;
@@ -50,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;
           });
     }