[InstCombine] Fold integer stores using llvm.assume

Author: ken-matsui

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

@@ -9,7 +9,7 @@ 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 [[X]], ptr @a, align 1
+; CHECK-NEXT:    store i1 true, ptr @a, align 1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X]])
 ; CHECK-NEXT:    ret void
 ;
@@ -21,7 +21,7 @@ define void @assume_store_i1(i1 %x) {
 define void @assume_store_i1_not(i1 %x) {
 ; CHECK-LABEL: define void @assume_store_i1_not(
 ; CHECK-SAME: i1 [[X:%.*]]) {
-; CHECK-NEXT:    store i1 [[X]], ptr @a, align 1
+; 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
@@ -37,7 +37,7 @@ define i1 @assume_store_i1_xor(ptr %G) {
 ; CHECK-SAME: ptr [[G:%.*]]) {
 ; CHECK-NEXT:    [[L:%.*]] = load i1, ptr [[G]], align 1
 ; CHECK-NEXT:    [[XOR:%.*]] = xor i1 [[L]], true
-; CHECK-NEXT:    store i1 [[XOR]], ptr @a, align 1
+; CHECK-NEXT:    store i1 true, ptr @a, align 1
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[XOR]])
 ; CHECK-NEXT:    ret i1 [[XOR]]
 ;
@@ -51,7 +51,7 @@ define i1 @assume_store_i1_xor(ptr %G) {
 define void @assume_store_i32_eq(i32 %x) {
 ; CHECK-LABEL: define void @assume_store_i32_eq(
 ; CHECK-SAME: i32 [[X:%.*]]) {
-; CHECK-NEXT:    store i32 [[X]], ptr @b, align 4
+; 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
@@ -66,9 +66,8 @@ define void @unreachable_implies_false(i8 %x) {
 ; CHECK-LABEL: define void @unreachable_implies_false(
 ; CHECK-SAME: i8 [[X:%.*]]) {
 ; CHECK-NEXT:  [[RET:.*:]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[X]], 5
-; CHECK-NEXT:    store i1 [[CMP]], ptr @a, align 1
-; CHECK-NEXT:    [[TMP0:%.*]] = xor i1 [[CMP]], true
+; 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
 ;