[AssumeBundles] Detection of Empty bundles

Summary: Prevent InstCombine from removing llvm.assume for which the arguement is true when they have operand bundles with usefull information.

Reviewers: jdoerfert, nikic, lebedev.ri

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D76147

include/llvm/IR/KnowledgeRetention.h

@@ -115,6 +115,16 @@ inline RetainedKnowledge getKnowledgeFromUseInAssume(const Use *U) {
                                          U->getOperandNo());
 }
 
+/// Return true iff the operand bundles of the provided llvm.assume doesn't
+/// contain any valuable information. This is true when:
+///  - The operand bundle is empty
+///  - The operand bundle only contains information about dropped values or
+///    constant folded values.
+///
+/// the argument to the call of llvm.assume may still be useful even if the
+/// function returned true.
+bool isAssumeWithEmptyBundle(CallInst &Assume);
+
 //===----------------------------------------------------------------------===//
 // Utilities for testing
 //===----------------------------------------------------------------------===//

lib/IR/KnowledgeRetention.cpp

@@ -286,6 +286,16 @@ RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
   return Result;
 }
 
+bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
+  IntrinsicInst &Assume = cast<IntrinsicInst>(CI);
+  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
+         "this function is intended to be used on llvm.assume");
+  return none_of(Assume.bundle_op_infos(),
+                 [](const CallBase::BundleOpInfo &BOI) {
+                   return BOI.Tag->getKey() != "ignore";
+                 });
+}
+
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   for (Instruction &I : instructions(F))

lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -48,6 +48,7 @@
 #include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/IntrinsicsPowerPC.h"
+#include "llvm/IR/KnowledgeRetention.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PatternMatch.h"
@@ -4093,7 +4094,7 @@ Instruction *InstCombiner::visitCallInst(CallInst &CI) {
     // then this one is redundant, and should be removed.
     KnownBits Known(1);
     computeKnownBits(IIOperand, Known, 0, II);
-    if (Known.isAllOnes())
+    if (Known.isAllOnes() && isAssumeWithEmptyBundle(*II))
       return eraseInstFromFunction(*II);
 
     // Update the cache of affected values for this assumption (we might be

lib/Transforms/Utils/Local.cpp

@@ -59,6 +59,7 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/KnowledgeRetention.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
@@ -411,7 +412,8 @@ bool llvm::wouldInstructionBeTriviallyDead(Instruction *I,
     // true are operationally no-ops.  In the future we can consider more
     // sophisticated tradeoffs for guards considering potential for check
     // widening, but for now we keep things simple.
-    if (II->getIntrinsicID() == Intrinsic::assume ||
+    if ((II->getIntrinsicID() == Intrinsic::assume &&
+         isAssumeWithEmptyBundle(*II)) ||
         II->getIntrinsicID() == Intrinsic::experimental_guard) {
       if (ConstantInt *Cond = dyn_cast<ConstantInt>(II->getArgOperand(0)))
         return !Cond->isZero();

test/Transforms/InstCombine/assume.ll

@@ -231,6 +231,27 @@ declare void @escape(i32* %a)
 
 ; Canonicalize a nonnull assumption on a load into metadata form.
 
+define i32 @bundle1(i32* %P) {
+; CHECK-LABEL: @bundle1(
+; CHECK-NEXT:    tail call void @llvm.assume(i1 true) [ "nonnull"(i32* [[P:%.*]]) ]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32* [[P]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+  tail call void @llvm.assume(i1 true) ["nonnull"(i32* %P)]
+  %load = load i32, i32* %P
+  ret i32 %load
+}
+
+define i32 @bundle2(i32* %P) {
+; CHECK-LABEL: @bundle2(
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32* [[P:%.*]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+  tail call void @llvm.assume(i1 true) ["ignore"(i32* undef)]
+  %load = load i32, i32* %P
+  ret i32 %load
+}
+
 define i1 @nonnull1(i32** %a) {
 ; CHECK-LABEL: @nonnull1(
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32*, i32** [[A:%.*]], align 8, !nonnull !6