[LoopUnrollAnalyzer] Use constant folding API for loads

Use ConstantFoldLoadFromConst() instead of a partial re-implementation.
This makes the code slightly more generic by not depending on the
exact structure of the constant.

Author: nikic

llvm/lib/Analysis/LoopUnrollAnalyzer.cpp

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/LoopUnrollAnalyzer.h"
+#include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
@@ -105,45 +106,20 @@ bool UnrolledInstAnalyzer::visitLoad(LoadInst &I) {
   auto AddressIt = SimplifiedAddresses.find(AddrOp);
   if (AddressIt == SimplifiedAddresses.end())
     return false;
-  const APInt &SimplifiedAddrOp = AddressIt->second.Offset;
 
   auto *GV = dyn_cast<GlobalVariable>(AddressIt->second.Base);
   // We're only interested in loads that can be completely folded to a
   // constant.
   if (!GV || !GV->hasDefinitiveInitializer() || !GV->isConstant())
     return false;
 
-  ConstantDataSequential *CDS =
-      dyn_cast<ConstantDataSequential>(GV->getInitializer());
-  if (!CDS)
+  Constant *Res =
+      ConstantFoldLoadFromConst(GV->getInitializer(), I.getType(),
+                                AddressIt->second.Offset, I.getDataLayout());
+  if (!Res)
     return false;
 
-  // We might have a vector load from an array. FIXME: for now we just bail
-  // out in this case, but we should be able to resolve and simplify such
-  // loads.
-  if (CDS->getElementType() != I.getType())
-    return false;
-
-  unsigned ElemSize = CDS->getElementType()->getPrimitiveSizeInBits() / 8U;
-  if (SimplifiedAddrOp.getActiveBits() > 64)
-    return false;
-  int64_t SimplifiedAddrOpV = SimplifiedAddrOp.getSExtValue();
-  if (SimplifiedAddrOpV < 0) {
-    // FIXME: For now we conservatively ignore out of bound accesses, but
-    // we're allowed to perform the optimization in this case.
-    return false;
-  }
-  uint64_t Index = static_cast<uint64_t>(SimplifiedAddrOpV) / ElemSize;
-  if (Index >= CDS->getNumElements()) {
-    // FIXME: For now we conservatively ignore out of bound accesses, but
-    // we're allowed to perform the optimization in this case.
-    return false;
-  }
-
-  Constant *CV = CDS->getElementAsConstant(Index);
-  assert(CV && "Constant expected.");
-  SimplifiedValues[&I] = CV;
-
+  SimplifiedValues[&I] = Res;
   return true;
 }
 

llvm/test/Transforms/LoopUnroll/full-unroll-heuristics-load-const.ll

@@ -0,0 +1,39 @@
+; RUN: opt < %s -S -passes=loop-unroll -unroll-max-iteration-count-to-analyze=1000 -unroll-threshold=10 -unroll-max-percent-threshold-boost=100 | FileCheck %s -check-prefix=TEST1
+; RUN: opt < %s -S -passes=loop-unroll -unroll-max-iteration-count-to-analyze=1000 -unroll-threshold=20 -unroll-max-percent-threshold-boost=200 | FileCheck %s -check-prefix=TEST2
+; RUN: opt < %s -S -passes=loop-unroll -unroll-max-iteration-count-to-analyze=1000 -unroll-threshold=20 -unroll-max-percent-threshold-boost=100 | FileCheck %s -check-prefix=TEST3
+
+; This test is a copy of full-unroll-heuristics.ll but with the constant
+; wrapped in an extra struct. This should not hinder the analysis.
+
+; If the absolute threshold is too low, we should not unroll:
+; TEST1: %array_const_idx = getelementptr inbounds { [9 x i32] }, ptr @known_constant, i64 0, i32 0, i64 %iv
+
+; Otherwise, we should:
+; TEST2-NOT: %array_const_idx = getelementptr inbounds { [9 x i32] }, ptr @known_constant, i64 0, i32 0, i64 %iv
+
+; If we do not boost threshold, the unroll will not happen:
+; TEST3: %array_const_idx = getelementptr inbounds { [9 x i32] }, ptr @known_constant, i64 0, i32 0, i64 %iv
+
+@known_constant = internal unnamed_addr constant { [9 x i32] } { [9 x i32] [i32 0, i32 -1, i32 0, i32 -1, i32 5, i32 -1, i32 0, i32 -1, i32 0] }, align 16
+
+define i32 @foo(ptr noalias nocapture readonly %src) {
+entry:
+  br label %loop
+
+loop:                                                ; preds = %loop, %entry
+  %iv = phi i64 [ 0, %entry ], [ %inc, %loop ]
+  %r  = phi i32 [ 0, %entry ], [ %add, %loop ]
+  %arrayidx = getelementptr inbounds i32, ptr %src, i64 %iv
+  %src_element = load i32, ptr %arrayidx, align 4
+  %array_const_idx = getelementptr inbounds { [9 x i32] }, ptr @known_constant, i64 0, i32 0, i64 %iv
+  %const_array_element = load i32, ptr %array_const_idx, align 4
+  %mul = mul nsw i32 %src_element, %const_array_element
+  %add = add nsw i32 %mul, %r
+  %inc = add nuw nsw i64 %iv, 1
+  %exitcond86.i = icmp eq i64 %inc, 9
+  br i1 %exitcond86.i, label %loop.end, label %loop
+
+loop.end:                                            ; preds = %loop
+  %r.lcssa = phi i32 [ %r, %loop ]
+  ret i32 %r.lcssa
+}