[CodeGenPrepare] Propagate assume-constrained constants before erasing llvm.assume

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -817,6 +817,12 @@ bool CodeGenPrepare::eliminateAssumptions(Function &F) {
       if (auto *Assume = dyn_cast<AssumeInst>(I)) {
         MadeChange = true;
         Value *Operand = Assume->getOperand(0);
+        Value *V;
+        ConstantInt *C;
+        if (match(Operand, m_c_ICmp(m_Value(V), m_ConstantInt(C))) &&
+            cast<ICmpInst>(Operand)->getPredicate() == ICmpInst::ICMP_EQ)
+          replaceDominatedUsesWith(V, C, getDT(), Assume);
+
         Assume->eraseFromParent();
 
         resetIteratorIfInvalidatedWhileCalling(&BB, [&]() {

llvm/test/CodeGen/X86/masked-load-store-assume.ll

@@ -0,0 +1,126 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
+
+define void @tgt_i8(ptr %p, i64 %x, i64 %C) {
+; CHECK-LABEL: tgt_i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movb %sil, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  %is.low8.mask = icmp eq i64 %C, 255
+  call void @llvm.assume(i1 %is.low8.mask)
+
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  ret void
+}
+
+define void @tgt_i16(ptr %p, i64 %x, i64 %C) {
+; CHECK-LABEL: tgt_i16:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movw %si, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  ret void
+}
+
+define void @tgt_i32(ptr %p, i64 %x, i64 %C) {
+; CHECK-LABEL: tgt_i32:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movl %esi, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  %is.low32.mask = icmp eq i64 %C, 4294967295
+  call void @llvm.assume(i1 %is.low32.mask)
+
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  ret void
+}
+
+define void @tgt_i16_not_const(ptr %p, i64 %x, i64 %C) {
+; CHECK-LABEL: tgt_i16_not_const:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movq (%rdi), %rax
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    andq %rdx, %rsi
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    movq %rsi, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  %is.some.mask = icmp uge i64 %C, 1
+  call void @llvm.assume(i1 %is.some.mask)
+
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  ret void
+}
+
+define void @tgt_i16_not_dominating(ptr %p, i64 %x, i64 %C, i1 %c) {
+; CHECK-LABEL: tgt_i16_not_dominating:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movq (%rdi), %rax
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    andq %rdx, %rsi
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    movq %rsi, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  br i1 %c, label %assume.bb, label %merge
+
+assume.bb:
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+  br label %merge
+
+merge:
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  ret void
+}
+
+define void @tgt_i16_use_before_assume(ptr %p, i64 %x, i64 %C) {
+; CHECK-LABEL: tgt_i16_use_before_assume:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    movq (%rdi), %rax
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    andq %rdx, %rsi
+; CHECK-NEXT:    xorq %rax, %rsi
+; CHECK-NEXT:    movq %rsi, (%rdi)
+; CHECK-NEXT:    retq
+entry:
+  %notC = xor i64 %C, -1
+  %old = load i64, ptr %p, align 8
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  store i64 %merged, ptr %p, align 8
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+  ret void
+}

llvm/test/Transforms/CodeGenPrepare/X86/propagate-assumed-constants.ll

@@ -0,0 +1,154 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='require<profile-summary>,function(codegenprepare)' -S -mtriple=x86_64-linux < %s | FileCheck %s
+
+define i64 @propagate_equality(i64 %x, i64 %old, i64 %C) {
+; CHECK-LABEL: @propagate_equality(
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 65535, -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], 65535
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  ret i64 %merged
+}
+
+define i64 @commuted_equality(i64 %x, i64 %old, i64 %C) {
+; CHECK-LABEL: @commuted_equality(
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 255, -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], 255
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  %is.low8.mask = icmp eq i64 255, %C
+  call void @llvm.assume(i1 %is.low8.mask)
+
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  ret i64 %merged
+}
+
+; A non-equality constraint must not be turned into a constant.
+define i64 @non_equality(i64 %x, i64 %old, i64 %C) {
+; CHECK-LABEL: @non_equality(
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 [[C:%.*]], -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], [[C]]
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  %is.some.mask = icmp uge i64 %C, 1
+  call void @llvm.assume(i1 %is.some.mask)
+
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  ret i64 %merged
+}
+
+; Equality to another non-constant value must not be folded.
+define i64 @equality_to_value(i64 %x, i64 %old, i64 %C, i64 %D) {
+; CHECK-LABEL: @equality_to_value(
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 [[C:%.*]], -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], [[C]]
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  %is.mask = icmp eq i64 %C, %D
+  call void @llvm.assume(i1 %is.mask)
+
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  ret i64 %merged
+}
+
+; The assume is in a block that does not dominate the use, so the constant must
+; not be propagated there.
+define i64 @not_dominating(i64 %x, i64 %old, i64 %C, i1 %c) {
+; CHECK-LABEL: @not_dominating(
+; CHECK-NEXT:  merge:
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 [[C:%.*]], -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], [[C]]
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  br i1 %c, label %assume.bb, label %merge
+
+assume.bb:
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+  br label %merge
+
+merge:
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+  ret i64 %merged
+}
+
+; Uses of the assumed value before the assume must not be replaced.
+define i64 @use_before_assume(i64 %x, i64 %old, i64 %C) {
+; CHECK-LABEL: @use_before_assume(
+; CHECK-NEXT:    [[NOTC:%.*]] = xor i64 [[C:%.*]], -1
+; CHECK-NEXT:    [[NEW_BYTES:%.*]] = and i64 [[X:%.*]], [[C]]
+; CHECK-NEXT:    [[OLD_REST:%.*]] = and i64 [[OLD:%.*]], [[NOTC]]
+; CHECK-NEXT:    [[MERGED:%.*]] = or disjoint i64 [[OLD_REST]], [[NEW_BYTES]]
+; CHECK-NEXT:    ret i64 [[MERGED]]
+;
+  %notC = xor i64 %C, -1
+  %new.bytes = and i64 %x, %C
+  %old.rest = and i64 %old, %notC
+  %merged = or disjoint i64 %old.rest, %new.bytes
+
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+  ret i64 %merged
+}
+
+; The propagated constant should also be usable when the assume directly feeds a
+; conditional branch that dominates further uses.
+define i64 @assume_feeding_branch(i64 %C, i1 %c) {
+; CHECK-LABEL: @assume_feeding_branch(
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[TRUE:%.*]], label [[FALSE:%.*]]
+; CHECK:       true:
+; CHECK-NEXT:    [[R1:%.*]] = add i64 65535, 1
+; CHECK-NEXT:    br label [[MERGE:%.*]]
+; CHECK:       false:
+; CHECK-NEXT:    [[R2:%.*]] = add i64 65535, 2
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       merge:
+; CHECK-NEXT:    [[RES:%.*]] = phi i64 [ [[R1]], [[TRUE]] ], [ [[R2]], [[FALSE]] ]
+; CHECK-NEXT:    ret i64 [[RES]]
+;
+  %is.low16.mask = icmp eq i64 %C, 65535
+  call void @llvm.assume(i1 %is.low16.mask)
+  br i1 %c, label %true, label %false
+
+true:
+  %r1 = add i64 %C, 1
+  br label %merge
+
+false:
+  %r2 = add i64 %C, 2
+  br label %merge
+
+merge:
+  %res = phi i64 [%r1, %true], [%r2, %false]
+  ret i64 %res
+}
+