[AggressiveInstCombine] Add `AssumptionCache` to aggressive instcombine

[anton-afanasyev]

Summary

Add support for @llvm.assume() to TruncInstCombine allowing
optimizations based on these intrinsics while computing known bits.

[spatel]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

290

IIUC, there's a subtle difference apart from adding the assumption cache in these calls - we are also now passing in the context instruction and domtree.

I'm not sure if it is possible to expose that in a regression test, but should we make that a preliminary patch? So add the wrapper calls while passing a null AC as step 1?

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

290

Ok, did it this way: https://reviews.llvm.org/D109288

[anton-afanasyev]

Rebase against wrapper calls

[anton-afanasyev]

Rebase

[spatel]

LGTM

[anton-afanasyev]

Closed by commit rGd1f9b216776d

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
 #include "AggressiveInstCombineInternal.h"
 #include "llvm-c/Initialization.h"
 #include "llvm-c/Transforms/AggressiveInstCombine.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
if (MadeChange)
     for (BasicBlock &BB : F)
       SimplifyInstructionsInBlock(&BB);
 
   return MadeChange;
 }
 
 /// This is the entry point for all transforms. Pass manager differences are
 /// handled in the callers of this function.
-static bool runImpl(Function &F, TargetLibraryInfo &TLI, DominatorTree &DT) {
+static bool runImpl(Function &F, AssumptionCache &AC, TargetLibraryInfo &TLI,
+                    DominatorTree &DT) {
   bool MadeChange = false;
   const DataLayout &DL = F.getParent()->getDataLayout();
-  TruncInstCombine TIC(TLI, DL, DT);
+  TruncInstCombine TIC(AC, TLI, DL, DT);
   MadeChange |= TIC.run(F);
   MadeChange |= foldUnusualPatterns(F, DT);
   return MadeChange;
 }
 
 void AggressiveInstCombinerLegacyPass::getAnalysisUsage(
     AnalysisUsage &AU) const {
   AU.setPreservesCFG();
+  AU.addRequired<AssumptionCacheTracker>();
   AU.addRequired<DominatorTreeWrapperPass>();
   AU.addRequired<TargetLibraryInfoWrapperPass>();
   AU.addPreserved<AAResultsWrapperPass>();
   AU.addPreserved<BasicAAWrapperPass>();
   AU.addPreserved<DominatorTreeWrapperPass>();
   AU.addPreserved<GlobalsAAWrapperPass>();
 }
 
 bool AggressiveInstCombinerLegacyPass::runOnFunction(Function &F) {
+  auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  return runImpl(F, TLI, DT);
+  return runImpl(F, AC, TLI, DT);
 }
 
 PreservedAnalyses AggressiveInstCombinePass::run(Function &F,
                                                  FunctionAnalysisManager &AM) {
+  auto &AC = AM.getResult<AssumptionAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
-  if (!runImpl(F, TLI, DT)) {
+  if (!runImpl(F, AC, TLI, DT)) {
     // No changes, all analyses are preserved.
     return PreservedAnalyses::all();
   }
   // Mark all the analyses that instcombine updates as preserved.
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }
 
 char AggressiveInstCombinerLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(AggressiveInstCombinerLegacyPass,
                       "aggressive-instcombine",
                       "Combine pattern based expressions", false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(AggressiveInstCombinerLegacyPass, "aggressive-instcombine",
                     "Combine pattern based expressions", false, false)
 
 // Initialization Routines
 void llvm::initializeAggressiveInstCombine(PassRegistry &Registry) {
   initializeAggressiveInstCombinerLegacyPassPass(Registry);

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombineInternal.h

 //
 // The motivation for this optimization is that evaluating and expression using
 // smaller bit-width is preferable, especially for vectorization where we can
 // fit more values in one vectorized instruction. In addition, this optimization
 // may decrease the number of cast instructions, but will not increase it.
 //===----------------------------------------------------------------------===//
 
 namespace llvm {
-  class DataLayout;
-  class DominatorTree;
-  class Function;
-  class Instruction;
-  class TargetLibraryInfo;
-  class TruncInst;
-  class Type;
-  class Value;
+class AssumptionCache;
 
 class TruncInstCombine {
+  AssumptionCache ∾
   TargetLibraryInfo &TLI;
   const DataLayout &DL;
   const DominatorTree &DT;
 
   /// List of all TruncInst instructions to be processed.
   SmallVector<TruncInst *, 4> Worklist;
 
   /// Current processed TruncInst instruction.
class TruncInstCombine {
   };
   /// An ordered map representing expression dag post-dominated by current
   /// processed TruncInst. It maps each instruction in the dag to its Info
   /// structure. The map is ordered such that each instruction appears before
   /// all other instructions in the dag that uses it.
   MapVector<Instruction *, Info> InstInfoMap;
 
 public:
-  TruncInstCombine(TargetLibraryInfo &TLI, const DataLayout &DL,
-                   const DominatorTree &DT)
-      : TLI(TLI), DL(DL), DT(DT), CurrentTruncInst(nullptr) {}
+  TruncInstCombine(AssumptionCache &AC, TargetLibraryInfo &TLI,
+                   const DataLayout &DL, const DominatorTree &DT)
+      : AC(AC), TLI(TLI), DL(DL), DT(DT), CurrentTruncInst(nullptr) {}
 
   /// Perform TruncInst pattern optimization on given function.
   bool run(Function &F);
 
 private:
   /// Build expression dag dominated by the /p CurrentTruncInst and append it to
   /// the InstInfoMap container.
   ///
private:
   /// Check if it is eligible to be reduced to a smaller type.
   ///
   /// \return the scalar version of the new type to be used for the reduced
   ///         expression dag, or nullptr if the expression dag is not eligible
   ///         to be reduced.
   Type *getBestTruncatedType();
 
   KnownBits computeKnownBits(const Value *V) const {
-    return llvm::computeKnownBits(V, DL, /*Depth=*/0, /*AC=*/nullptr,
-                                  /*CtxI=*/nullptr, &DT);
+    return llvm::computeKnownBits(V, DL, /*Depth=*/0, &AC,
+                                  /*CtxI=*/cast<Instruction>(CurrentTruncInst),
+                                  &DT);
   }
 
   unsigned ComputeNumSignBits(const Value *V) const {
-    return llvm::ComputeNumSignBits(V, DL, /*Depth=*/0, /*AC=*/nullptr,
-                                    /*CtxI=*/nullptr, &DT);
+    return llvm::ComputeNumSignBits(
+        V, DL, /*Depth=*/0, &AC, /*CtxI=*/cast<Instruction>(CurrentTruncInst),
+        &DT);
   }
 
   /// Given a \p V value and a \p SclTy scalar type return the generated reduced
   /// value of \p V based on the type \p SclTy.
   ///
   /// \param V value to be reduced.
   /// \param SclTy scalar version of new type to reduce to.
   /// \return the new reduced value.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "AggressiveInstCombineInternal.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
-#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/Support/KnownBits.h"
 
 using namespace llvm;
 
Type *TruncInstCombine::getBestTruncatedType() {
   // bits of the value-to-be-shifted are zeros.
   // For `ashr` adjust MinBitWidth so that all potentially truncated
   // bits of the value-to-be-shifted are sign bits (all zeros or ones)
   // and even one (first) untruncated bit is sign bit.
   // Exit early if MinBitWidth is not less than original bitwidth.
   for (auto &Itr : InstInfoMap) {
     Instruction *I = Itr.first;
     if (I->isShift()) {
       KnownBits KnownRHS = computeKnownBits(I->getOperand(1));
       unsigned MinBitWidth = KnownRHS.getMaxValue()
                                  .uadd_sat(APInt(OrigBitWidth, 1))
                                  .getLimitedValue(OrigBitWidth);
       if (MinBitWidth == OrigBitWidth)
         return nullptr;
       if (I->getOpcode() == Instruction::LShr) {
         KnownBits KnownLHS = computeKnownBits(I->getOperand(0));
         MinBitWidth =

llvm/test/Transforms/AggressiveInstCombine/trunc_assume.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -aggressive-instcombine -S | FileCheck %s
 
 define i16 @trunc_shl(i16 %x, i16 %y) {
 ; CHECK-LABEL: @trunc_shl(
 ; CHECK-NEXT:    [[CMP0:%.*]] = icmp ult i16 [[Y:%.*]], 16
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP0]])
-; CHECK-NEXT:    [[ZEXTX:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = zext i16 [[Y]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = shl i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = shl i16 [[X:%.*]], [[Y]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %cmp0 = icmp ult i16 %y, 16
   call void @llvm.assume(i1 %cmp0)
 
   %zextx = zext i16 %x to i32
   %zexty = zext i16 %y to i32
 
   %i0 = shl i32 %zextx, %zexty
   %r = trunc i32 %i0 to i16
   ret i16 %r
 }
 
 define i16 @trunc_lshr(i16 %x, i16 %y) {
 ; CHECK-LABEL: @trunc_lshr(
 ; CHECK-NEXT:    [[CMP0:%.*]] = icmp ult i16 [[X:%.*]], 0
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp ult i16 [[Y:%.*]], 16
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP0]])
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
-; CHECK-NEXT:    [[ZEXTX:%.*]] = zext i16 [[X]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = zext i16 [[Y]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = lshr i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = lshr i16 [[X]], [[Y]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %cmp0 = icmp ult i16 %x, 65536
   %cmp1 = icmp ult i16 %y, 16
   call void @llvm.assume(i1 %cmp0)
   call void @llvm.assume(i1 %cmp1)
 
   %zextx = zext i16 %x to i32
   %zexty = zext i16 %y to i32
 
   %i0 = lshr i32 %zextx, %zexty
   %r = trunc i32 %i0 to i16
   ret i16 %r
 }
 
 define i16 @trunc_ashr(i16 %x, i16 %y) {
 ; CHECK-LABEL: @trunc_ashr(
 ; CHECK-NEXT:    [[CMP0:%.*]] = icmp slt i16 [[X:%.*]], 32767
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp sge i16 [[X]], -32768
 ; CHECK-NEXT:    [[CMP2:%.*]] = icmp ult i16 [[Y:%.*]], 16
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP0]])
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP2]])
-; CHECK-NEXT:    [[ZEXTX:%.*]] = sext i16 [[X]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = sext i16 [[Y]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = ashr i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = ashr i16 [[X]], [[Y]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %cmp0 = icmp slt i16 %x, 32767
   %cmp1 = icmp sge i16 %x, -32768
   %cmp2 = icmp ult i16 %y, 16
   call void @llvm.assume(i1 %cmp0)
   call void @llvm.assume(i1 %cmp1)
   call void @llvm.assume(i1 %cmp2)