[InstCombine] Fold integer stores using llvm.assume

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

@@ -1535,6 +1535,16 @@ Instruction *InstCombinerImpl::visitStoreInst(StoreInst &SI) {
     if (Value *V = simplifyNonNullOperand(Ptr, /*HasDereferenceable=*/true))
       return replaceOperand(SI, 1, V);
 
+  // If the stored value is constant at this point (e.g., constrained by a
+  // guaranteed-to-execute llvm.assume), fold the store.
+  if (Val->getType()->isIntegerTy() && !isa<Constant>(Val) &&
+      !isa<PoisonValue>(Val)) {
+    KnownBits Known = computeKnownBits(Val, &SI);
+    if (Known.isConstant())
+      return replaceOperand(
+          SI, 0, ConstantInt::get(SI.getContext(), Known.getConstant()));
+  }
+
   return nullptr;
 }
 

llvm/test/Transforms/InstCombine/store-assume-fold.ll

@@ -0,0 +1,135 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+@a = external global i1
+@b = external global i32
+@c = external global i8
+@v = external global <2 x i8>
+
+declare void @llvm.assume(i1)
+
+define void @assume_store_i1(i1 %x) {
+; CHECK-LABEL: define void @assume_store_i1(
+; CHECK-SAME: i1 [[X:%.*]]) {
+; CHECK-NEXT:    store i1 true, ptr @a, align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[X]])
+; CHECK-NEXT:    ret void
+;
+  store i1 %x, ptr @a, align 1
+  call void @llvm.assume(i1 %x)
+  ret void
+}
+
+define void @assume_store_i1_not(i1 %x) {
+; CHECK-LABEL: define void @assume_store_i1_not(
+; CHECK-SAME: i1 [[X:%.*]]) {
+; CHECK-NEXT:    store i1 false, ptr @a, align 1
+; CHECK-NEXT:    [[NOT:%.*]] = xor i1 [[X]], true
+; CHECK-NEXT:    call void @llvm.assume(i1 [[NOT]])
+; CHECK-NEXT:    ret void
+;
+  store i1 %x, ptr @a, align 1
+  %not = xor i1 %x, true
+  call void @llvm.assume(i1 %not)
+  ret void
+}
+
+define i1 @assume_store_i1_xor(ptr %G) {
+; CHECK-LABEL: define i1 @assume_store_i1_xor(
+; CHECK-SAME: ptr [[G:%.*]]) {
+; CHECK-NEXT:    [[L:%.*]] = load i1, ptr [[G]], align 1
+; CHECK-NEXT:    [[XOR:%.*]] = xor i1 [[L]], true
+; CHECK-NEXT:    store i1 true, ptr @a, align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[XOR]])
+; CHECK-NEXT:    ret i1 [[XOR]]
+;
+  %L = load i1, ptr %G, align 1
+  %xor = xor i1 %L, true
+  store i1 %xor, ptr @a, align 1
+  call void @llvm.assume(i1 %xor)
+  ret i1 %xor
+}
+
+define void @assume_store_i32_eq(i32 %x) {
+; CHECK-LABEL: define void @assume_store_i32_eq(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    store i32 10, ptr @b, align 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 10
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    ret void
+;
+  store i32 %x, ptr @b, align 4
+  %cmp = icmp eq i32 %x, 10
+  call void @llvm.assume(i1 %cmp)
+  ret void
+}
+
+define void @unreachable_implies_false(i8 %x) {
+; CHECK-LABEL: define void @unreachable_implies_false(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:  [[RET:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i8 [[X]], 6
+; CHECK-NEXT:    store i1 false, ptr @a, align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP0]])
+; CHECK-NEXT:    ret void
+;
+; Original code before `simplifycfg` from #134992:
+;   %cmp = icmp ugt i8 %x, 5
+;   store i1 %cmp, ptr @a, align 1
+;   br i1 %cmp, label %ub, label %ret
+; ret:
+;   ret void
+; ub:
+;   unreachable
+ret:
+  %cmp = icmp ugt i8 %x, 5
+  store i1 %cmp, ptr @a, align 1
+  %0 = xor i1 %cmp, true
+  call void @llvm.assume(i1 %0)
+  ret void
+}
+
+; Negative test: assume does not make the store value constant.
+define void @assume_store_i8_nonconstant(i8 %x) {
+; CHECK-LABEL: define void @assume_store_i8_nonconstant(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    store i8 [[X]], ptr @c, align 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i8 [[X]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    ret void
+;
+  store i8 %x, ptr @c, align 1
+  %cmp = icmp ne i8 %x, 0
+  call void @llvm.assume(i1 %cmp)
+  ret void
+}
+
+; Negative test: volatile stores must not be folded.
+define void @assume_store_i1_volatile(i1 %x) {
+; CHECK-LABEL: define void @assume_store_i1_volatile(
+; CHECK-SAME: i1 [[X:%.*]]) {
+; CHECK-NEXT:    store volatile i1 [[X]], ptr @a, align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[X]])
+; CHECK-NEXT:    ret void
+;
+  store volatile i1 %x, ptr @a, align 1
+  call void @llvm.assume(i1 %x)
+  ret void
+}
+
+; Negative test: non-scalar stores are not folded.
+define void @assume_store_vec_noninteger(<2 x i8> %x) {
+; CHECK-LABEL: define void @assume_store_vec_noninteger(
+; CHECK-SAME: <2 x i8> [[X:%.*]]) {
+; CHECK-NEXT:    store <2 x i8> [[X]], ptr @v, align 2
+; CHECK-NEXT:    [[LANE0:%.*]] = extractelement <2 x i8> [[X]], i64 0
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[LANE0]], 7
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    ret void
+;
+  store <2 x i8> %x, ptr @v, align 2
+  %lane0 = extractelement <2 x i8> %x, i32 0
+  %cmp = icmp eq i8 %lane0, 7
+  call void @llvm.assume(i1 %cmp)
+  ret void
+}