[POC Only] Demonstration of approach to undo exit value rewriting

[reames]

Summary

Not for review.

Just to demonstrate what one example approach might look like.
target datalayout = "n8:16:32:64"
define i32 @test2(i32 %n) {
entry:

br label %loop

loop:

%iv = phi i32 [0, %entry], [%iv.next, %loop]
%iv.next = add i32 %iv, 1
%cmp = icmp ult i32 %iv, %n
br i1 %cmp, label %loop, label %exit

exit:

ret i32 %n ;; rewrites this to %iv

}

The hard part is figuring out the right heuristic here.

[danilaml]

Thanks for submitting this. It's good to see different approaches.

lib/Transforms/Scalar/ReuseLoopExitValues.cpp

110

Is the "loop context" really needed for expressions? Aren't dominance checks done later sufficient to know whether it's legal to replace one value with another or not?

135

I'm not sure why this dominance check is used. We are searching for the candidate to replace U.get(), Isn't it enough for DefI to dominate U.get()? Currently, this check prevents from replacing uses in phis which are common.

144

Won't this function always return U.get()? Any original ExitValue would always dominate it's rewritten form, if I'm not mistaken.

163

I had to add some type checks here (or create bitcasts) to prevent asserts.

[danilaml]

lib/Transforms/Scalar/ReuseLoopExitValues.cpp

110

Nevermind. I relized that it's likely unavoidable for loops consisting of several BBs.

[lebedev.ri]

What would it take to get something like this moving?
It seems the need for this is again being encounterd in-the-wild as per https://reviews.llvm.org/D73501#1905171

[reames]

In D66450#1911219, @lebedev.ri wrote:

What would it take to get something like this moving?
It seems the need for this is again being encounterd in-the-wild as per https://reviews.llvm.org/D73501#1905171

Someone to drive it. A more solid implementation is needed, and then we need to check for unexpected performance interactions.

include/llvm/Transforms/Scalar/ReuseLoopExitValues.h

+//===- IndVarSimplify.h - Induction Variable Simplification -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides the interface for the Induction Variable
+// Simplification pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCALAR_REUSELOOPEXITVALUES_H
+#define LLVM_TRANSFORMS_SCALAR_REUSELOOPEXITVALUES_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class Loop;
+class LPMUpdater;
+
+class ReuseLoopExitValuesPass : public PassInfoMixin<ReuseLoopExitValuesPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCALAR_REUSELOOPEXITVALUES_H

lib/Passes/PassBuilder.cpp

 #include "llvm/Transforms/Scalar/MemCpyOptimizer.h"
 #include "llvm/Transforms/Scalar/MergeICmps.h"
 #include "llvm/Transforms/Scalar/MergedLoadStoreMotion.h"
 #include "llvm/Transforms/Scalar/NaryReassociate.h"
 #include "llvm/Transforms/Scalar/NewGVN.h"
 #include "llvm/Transforms/Scalar/PartiallyInlineLibCalls.h"
 #include "llvm/Transforms/Scalar/Reassociate.h"
 #include "llvm/Transforms/Scalar/RewriteStatepointsForGC.h"
+#include "llvm/Transforms/Scalar/ReuseLoopExitValues.h"
 #include "llvm/Transforms/Scalar/SCCP.h"
 #include "llvm/Transforms/Scalar/SROA.h"
 #include "llvm/Transforms/Scalar/Scalarizer.h"
 #include "llvm/Transforms/Scalar/SimpleLoopUnswitch.h"
 #include "llvm/Transforms/Scalar/SimplifyCFG.h"
 #include "llvm/Transforms/Scalar/Sink.h"
 #include "llvm/Transforms/Scalar/SpeculateAroundPHIs.h"
 #include "llvm/Transforms/Scalar/SpeculativeExecution.h"
if (EnableHotColdSplit && !LTOPreLink)
     MPM.addPass(HotColdSplittingPass());
 
   // LoopSink pass sinks instructions hoisted by LICM, which serves as a
   // canonicalization pass that enables other optimizations. As a result,
   // LoopSink pass needs to be a very late IR pass to avoid undoing LICM
   // result too early.
   OptimizePM.addPass(LoopSinkPass());
 
-  // And finally clean up LCSSA form before generating code.
+  // And finally clean up LCSSA form before generating code, and reuse
+  // expressions which compute values as close to their defs as possible.
   OptimizePM.addPass(InstSimplifyPass());
 
+  // Reuse expressions which compute values as close to their defs as possible;
+  // this reverses the loop exit value canonical form we use at the IR level to
+  // use values computed in loops along exits.
+  OptimizePM.addPass(ReuseLoopExitValuesPass());
+
   // This hoists/decomposes div/rem ops. It should run after other sink/hoist
   // passes to avoid re-sinking, but before SimplifyCFG because it can allow
   // flattening of blocks.
   OptimizePM.addPass(DivRemPairsPass());
 
   // LoopSink (and other loop passes since the last simplifyCFG) might have
   // resulted in single-entry-single-exit or empty blocks. Clean up the CFG.
   OptimizePM.addPass(SimplifyCFGPass());

lib/Passes/PassRegistry.def

 FUNCTION_PASS("view-cfg-only", CFGOnlyViewerPass())
 FUNCTION_PASS("transform-warning", WarnMissedTransformationsPass())
 FUNCTION_PASS("asan", AddressSanitizerPass(false, false, false))
 FUNCTION_PASS("kasan", AddressSanitizerPass(true, false, false))
 FUNCTION_PASS("msan", MemorySanitizerPass({}))
 FUNCTION_PASS("kmsan", MemorySanitizerPass({0, false, /*Kernel=*/true}))
 FUNCTION_PASS("tsan", ThreadSanitizerPass())
 FUNCTION_PASS("sancov-func", SanitizerCoveragePass())
+FUNCTION_PASS("reuse-loop-exit-values", ReuseLoopExitValuesPass());
 #undef FUNCTION_PASS
 
 #ifndef FUNCTION_PASS_WITH_PARAMS
 #define FUNCTION_PASS_WITH_PARAMS(NAME, CREATE_PASS, PARSER)
 #endif
 FUNCTION_PASS_WITH_PARAMS("unroll",
                            [](LoopUnrollOptions Opts) {
                              return LoopUnrollPass(Opts);

lib/Transforms/Scalar/CMakeLists.txt

add_llvm_library(LLVMScalarOpts
   MergedLoadStoreMotion.cpp
   NaryReassociate.cpp
   NewGVN.cpp
   PartiallyInlineLibCalls.cpp
   PlaceSafepoints.cpp
   Reassociate.cpp
   Reg2Mem.cpp
   RewriteStatepointsForGC.cpp
+  ReuseLoopExitValues.cpp
   SCCP.cpp
   SROA.cpp
   Scalar.cpp
   Scalarizer.cpp
   SeparateConstOffsetFromGEP.cpp
   SimpleLoopUnswitch.cpp
   SimplifyCFGPass.cpp
   Sink.cpp

lib/Transforms/Scalar/ReuseLoopExitValues.cpp

+//===- IndVarSimplify.cpp - Induction Variable Elimination ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// TODO: naming, comments, excess includes, etc...
+// TODO: probably should rename to something more generic since there is
+// nothing specifical to loop exit values in the code
+// TODO: need to generalize to values computable along one (dominating) exit
+// 
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/ReuseLoopExitValues.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/None.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/OrderedInstructions.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/ScalarEvolutionExpander.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/ConstantRange.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Use.h"
+#include "llvm/IR/User.h"
+#include "llvm/IR/Value.h"
+#include "llvm/IR/ValueHandle.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compiler.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
+#include <cassert>
+#include <cstdint>
+#include <utility>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "reuse-loop-exit-values"
+
+namespace {
+bool run(Function &F,
+         DominatorTree &DT, LoopInfo &LI, ScalarEvolution &SE) {
+  // Construct a mapping of SCEV @ Context to possible llvm::Value*s which
+  // compute the given expression.  Note that we don't worry about which
+  // expressions are available at which program points, we just build a master
+  // list.
+  // TODO: need to remove duplicates w/a set
+  DenseMap<std::pair<const SCEV*, Loop*>,
+           SmallVector<Value *,1>> AvailableExprs;
+  auto AddValueAtContext = [&](Value *V, Loop* L) {
+    if (!SE.isSCEVable(V->getType()))
+      return;
+    const SCEV *S = SE.getSCEV(V);
+    if (isa<SCEVCouldNotCompute>(S))
+      return;
+    dbgs() << "Adding " << *S << "[" << L << "]";
+    dbgs() << *V << "\n";
+    
+    AvailableExprs[{S, L}].push_back(V);
+    while (L) {
+      L = L->getParentLoop();
+      S = SE.getSCEVAtScope(S, L);
+      if (isa<SCEVCouldNotCompute>(S))
+        break;
+    dbgs() << "Adding " << *S << "[" << L << "]";
+    dbgs() << *V << "\n";
+      AvailableExprs[{S, L}].push_back(V);
+    }
+  };
+  for (Instruction &I : instructions(F)) {
+    Loop *L = LI.getLoopFor(I.getParent());
+    AddValueAtContext(&I, L);
+    for (Value *Op : I.operands())
+      AddValueAtContext(Op, L);
+  }
+
+  // TODO: There is an important missing heurstic here.  We do not *always*
+  // want to reverse this if it results in longer live ranges *within* a loop
+  // which could otherwise be folded.  Thought required?
+  OrderedInstructions OI(&DT);
+  auto FindNearbyDef = [&](SmallVectorImpl<Value*>& Defs, Use &U) {
+    // Find a def which is nearby (i.e. within the block containing the most
+    // dominating definitions).  We don't care about *order wtthin that block*,
+    // and are simply trying to reduce cross block live ranges.
+    Instruction *UserI = cast<Instruction>(U.getUser());
+    Value *Closest = U.get();
+    for (Value *Def : Defs) {
+      Instruction *DefI = dyn_cast<Instruction>(Def);
+      if (!DefI)
+        // All values are equally distant
+        continue;
+      if (!OI.dominates(DefI, UserI))
+        // Not a valid def for *this* use or in the use BB (ignored for
+        // efficiency reasons)
+        continue;
+      Instruction *ClosestI = dyn_cast<Instruction>(Closest);
+      if (!ClosestI ||
+          OI.dominates(ClosestI, DefI))
+        Closest = Def;
+    }
+    return Closest;
+  };
+
+  bool Changed = false;
+  for (Instruction &I : instructions(F)) {
+    Loop *L = LI.getLoopFor(I.getParent());
+    for (Use &U : I.operands()) {
+      if (!SE.isSCEVable(U.get()->getType()))
+        continue;
+      const SCEV *S = SE.getSCEV(U.get());
+      if (isa<SCEVCouldNotCompute>(S))
+        continue;
+      auto &Defs = AvailableExprs[{S, L}];
+      assert(!Defs.empty());
+      dbgs() << "Considering " << *U.get() << " in " << I << "\n"
+             << "   with defs: \n";
+      for (Value *Def : Defs)
+        dbgs() << "  -- " << *Def << "\n";
+      Value *NewDef = FindNearbyDef(Defs, U);
+      if (NewDef == U.get())
+        continue;
+
+      // TODO: Kill instructions which become dead after this loop
+      U.set(NewDef);
+      Changed = true;
+    }
+  }
+  return Changed;
+}
+}
+
+PreservedAnalyses
+ReuseLoopExitValuesPass::run(Function &F, FunctionAnalysisManager &AM) {
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
+  auto &LI = AM.getResult<LoopAnalysis>(F);
+  auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
+
+  dbgs() << "---------------\n";
+  dbgs() << F.getName() << "\n";
+  dbgs() << "---------------\n";
+
+  if (!::run(F, DT, LI, SE))
+    return PreservedAnalyses::all();
+
+  auto PA = PreservedAnalyses::none();
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+