[DropUnnecessaryAssumes] Add pass for dropping assumes (#159403)

This adds a new pass for dropping assumes that are unlikely to be useful
for further optimization.

It works by discarding any assumes whose affected values are one-use
(which implies that they are only used by the assume, i.e. ephemeral).

This pass currently runs at the start of the module optimization
pipeline, that is post-inline and post-link. Before that point, it is
more likely for previously "useless" assumes to become useful again,
e.g. because an additional user of the value is introduced after
inlining + CSE.

Author: nikic

clang/test/CodeGen/inline-asm-x86-flag-output.c

@@ -389,7 +389,7 @@ int test_assume_boolean_flag(long nr, volatile long *addr) {
       : "=@cca"(x), "=@ccae"(y), "=m"(*(volatile long *)(addr))
       : "r"(nr)
       : "cc");
-  if (x)
+  if (x && y)
     return 0;
   return 1;
 }

llvm/include/llvm/Transforms/Scalar/DropUnnecessaryAssumes.h

@@ -0,0 +1,27 @@
+//===------------------------------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Pass that drops assumes that are unlikely to be useful.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_SCALAR_DROPUNNECESSARYASSUMES_H
+#define LLVM_TRANSFORMS_SCALAR_DROPUNNECESSARYASSUMES_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+struct DropUnnecessaryAssumesPass
+    : public PassInfoMixin<DropUnnecessaryAssumesPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_SCALAR_DROPUNNECESSARYASSUMES_H

llvm/lib/Passes/PassBuilder.cpp

@@ -273,6 +273,7 @@
 #include "llvm/Transforms/Scalar/DFAJumpThreading.h"
 #include "llvm/Transforms/Scalar/DeadStoreElimination.h"
 #include "llvm/Transforms/Scalar/DivRemPairs.h"
+#include "llvm/Transforms/Scalar/DropUnnecessaryAssumes.h"
 #include "llvm/Transforms/Scalar/EarlyCSE.h"
 #include "llvm/Transforms/Scalar/FlattenCFG.h"
 #include "llvm/Transforms/Scalar/Float2Int.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -92,6 +92,7 @@
 #include "llvm/Transforms/Scalar/DFAJumpThreading.h"
 #include "llvm/Transforms/Scalar/DeadStoreElimination.h"
 #include "llvm/Transforms/Scalar/DivRemPairs.h"
+#include "llvm/Transforms/Scalar/DropUnnecessaryAssumes.h"
 #include "llvm/Transforms/Scalar/EarlyCSE.h"
 #include "llvm/Transforms/Scalar/Float2Int.h"
 #include "llvm/Transforms/Scalar/GVN.h"
@@ -1498,6 +1499,13 @@ PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
   invokeOptimizerEarlyEPCallbacks(MPM, Level, LTOPhase);
 
   FunctionPassManager OptimizePM;
+
+  // Only drop unnecessary assumes post-inline and post-link, as otherwise
+  // additional uses of the affected value may be introduced through inlining
+  // and CSE.
+  if (!isLTOPreLink(LTOPhase))
+    OptimizePM.addPass(DropUnnecessaryAssumesPass());
+
   // Scheduling LoopVersioningLICM when inlining is over, because after that
   // we may see more accurate aliasing. Reason to run this late is that too
   // early versioning may prevent further inlining due to increase of code

llvm/lib/Passes/PassRegistry.def

@@ -425,6 +425,7 @@ FUNCTION_PASS("dot-post-dom", PostDomPrinter())
 FUNCTION_PASS("dot-post-dom-only", PostDomOnlyPrinter())
 FUNCTION_PASS("dse", DSEPass())
 FUNCTION_PASS("dwarf-eh-prepare", DwarfEHPreparePass(TM))
+FUNCTION_PASS("drop-unnecessary-assumes", DropUnnecessaryAssumesPass())
 FUNCTION_PASS("expand-large-div-rem", ExpandLargeDivRemPass(TM))
 FUNCTION_PASS("expand-memcmp", ExpandMemCmpPass(TM))
 FUNCTION_PASS("expand-reductions", ExpandReductionsPass())

llvm/lib/Transforms/Scalar/CMakeLists.txt

@@ -11,6 +11,7 @@ add_llvm_component_library(LLVMScalarOpts
   DeadStoreElimination.cpp
   DFAJumpThreading.cpp
   DivRemPairs.cpp
+  DropUnnecessaryAssumes.cpp
   EarlyCSE.cpp
   FlattenCFGPass.cpp
   Float2Int.cpp

llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp

@@ -0,0 +1,62 @@
+//===------------------------------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/DropUnnecessaryAssumes.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/Transforms/Utils/Local.h"
+
+using namespace llvm;
+using namespace llvm::PatternMatch;
+
+PreservedAnalyses
+DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
+  AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
+  bool Changed = false;
+
+  for (AssumptionCache::ResultElem &Elem : AC.assumptions()) {
+    auto *Assume = cast_or_null<AssumeInst>(Elem.Assume);
+    if (!Assume)
+      continue;
+
+    // TODO: Handle assumes with operand bundles.
+    if (Assume->hasOperandBundles())
+      continue;
+
+    Value *Cond = Assume->getArgOperand(0);
+    // Don't drop type tests, which have special semantics.
+    if (match(Cond, m_Intrinsic<Intrinsic::type_test>()))
+      continue;
+
+    SmallPtrSet<Value *, 8> Affected;
+    findValuesAffectedByCondition(Cond, /*IsAssume=*/true,
+                                  [&](Value *A) { Affected.insert(A); });
+
+    // If all the affected uses have only one use (part of the assume), then
+    // the assume does not provide useful information. Note that additional
+    // users may appear as a result of inlining and CSE, so we should only
+    // make this assumption late in the optimization pipeline.
+    // TODO: Handle dead cyclic usages.
+    // TODO: Handle multiple dead assumes on the same value.
+    if (!all_of(Affected, match_fn(m_OneUse(m_Value()))))
+      continue;
+
+    Assume->eraseFromParent();
+    RecursivelyDeleteTriviallyDeadInstructions(Cond);
+    Changed = true;
+  }
+
+  if (Changed) {
+    PreservedAnalyses PA;
+    PA.preserveSet<CFGAnalyses>();
+    return PA;
+  }
+  return PreservedAnalyses::all();
+}

llvm/test/Other/new-pm-defaults.ll

@@ -244,6 +244,7 @@
 ; CHECK-O-NEXT: Running pass: ReversePostOrderFunctionAttrsPass
 ; CHECK-O-NEXT: Running pass: RecomputeGlobalsAAPass
 ; CHECK-EP-OPTIMIZER-EARLY: Running pass: NoOpModulePass
+; CHECK-DEFAULT-NEXT: Running pass: DropUnnecessaryAssumesPass
 ; CHECK-O-NEXT: Running pass: Float2IntPass
 ; CHECK-O-NEXT: Running pass: LowerConstantIntrinsicsPass on foo
 ; CHECK-MATRIX: Running pass: LowerMatrixIntrinsicsPass on f

llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

@@ -167,6 +167,7 @@
 ; CHECK-POSTLINK-O-NEXT: Running pass: ReversePostOrderFunctionAttrsPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: RecomputeGlobalsAAPass
 ; CHECK-POST-EP-OPT-EARLY-NEXT: Running pass: NoOpModulePass
+; CHECK-POSTLINK-O-NEXT: Running pass: DropUnnecessaryAssumesPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: Float2IntPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: LowerConstantIntrinsicsPass
 ; CHECK-POSTLINK-O3-NEXT: Running pass: ControlHeightReductionPass

llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll

@@ -150,6 +150,7 @@
 ; CHECK-O-NEXT: Running pass: EliminateAvailableExternallyPass
 ; CHECK-O-NEXT: Running pass: ReversePostOrderFunctionAttrsPass
 ; CHECK-O-NEXT: Running pass: RecomputeGlobalsAAPass
+; CHECK-O-NEXT: Running pass: DropUnnecessaryAssumesPass
 ; CHECK-O-NEXT: Running pass: Float2IntPass
 ; CHECK-O-NEXT: Running pass: LowerConstantIntrinsicsPass
 ; CHECK-O3-NEXT: Running pass: ControlHeightReductionPass