; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+lzcnt,+bmi,+popcnt | FileCheck %s

; Test that loads are folded into the memory form of LZCNT/TZCNT/POPCNT when
; the loaded value is also used in an icmp-eq-zero/select sequence whose
; condition is satisfied by the EFLAGS already produced by the bit-count
; instruction (LZCNT/TZCNT set CF=1, POPCNT sets ZF=1 when input is zero).
;
; Before the fix, ISel chose the register form (rr) because the loaded value
; had two uses: the bit-count instruction and the compare.  The peephole
; optimizer later removed the redundant TEST but missed folding the load into
; the now single-use bit-count instruction.
; See https://github.com/llvm/llvm-project/issues/164782

; --- lzcnt: CF=1 when input is zero, cmovb takes the fallback ---

define i32 @lzcnt32_load_cmov(ptr %p0, i32 %a1) {
; CHECK-LABEL: lzcnt32_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    lzcntl (%rdi), %eax
; CHECK-NEXT:    cmovbl %esi, %eax
; CHECK-NEXT:    retq
  %v = load i32, ptr %p0, align 4
  %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
  %iszero = icmp eq i32 %v, 0
  %res = select i1 %iszero, i32 %a1, i32 %cnt
  ret i32 %res
}

define i64 @lzcnt64_load_cmov(ptr %p0, i64 %a1) {
; CHECK-LABEL: lzcnt64_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    lzcntq (%rdi), %rax
; CHECK-NEXT:    cmovbq %rsi, %rax
; CHECK-NEXT:    retq
  %v = load i64, ptr %p0, align 8
  %cnt = tail call range(i64 0, 65) i64 @llvm.ctlz.i64(i64 %v, i1 true)
  %iszero = icmp eq i64 %v, 0
  %res = select i1 %iszero, i64 %a1, i64 %cnt
  ret i64 %res
}

; --- tzcnt: CF=1 when input is zero, cmovb takes the fallback ---

define i32 @tzcnt32_load_cmov(ptr %p0, i32 %a1) {
; CHECK-LABEL: tzcnt32_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    tzcntl (%rdi), %eax
; CHECK-NEXT:    cmovbl %esi, %eax
; CHECK-NEXT:    retq
  %v = load i32, ptr %p0, align 4
  %cnt = tail call range(i32 0, 33) i32 @llvm.cttz.i32(i32 %v, i1 true)
  %iszero = icmp eq i32 %v, 0
  %res = select i1 %iszero, i32 %a1, i32 %cnt
  ret i32 %res
}

define i64 @tzcnt64_load_cmov(ptr %p0, i64 %a1) {
; CHECK-LABEL: tzcnt64_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    tzcntq (%rdi), %rax
; CHECK-NEXT:    cmovbq %rsi, %rax
; CHECK-NEXT:    retq
  %v = load i64, ptr %p0, align 8
  %cnt = tail call range(i64 0, 65) i64 @llvm.cttz.i64(i64 %v, i1 true)
  %iszero = icmp eq i64 %v, 0
  %res = select i1 %iszero, i64 %a1, i64 %cnt
  ret i64 %res
}

; --- popcnt: ZF=1 when input is zero, cmove takes the fallback ---

define i32 @popcnt32_load_cmov(ptr %p0, i32 %a1) {
; CHECK-LABEL: popcnt32_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    popcntl (%rdi), %eax
; CHECK-NEXT:    cmovel %esi, %eax
; CHECK-NEXT:    retq
  %v = load i32, ptr %p0, align 4
  %cnt = tail call i32 @llvm.ctpop.i32(i32 %v)
  %iszero = icmp eq i32 %v, 0
  %res = select i1 %iszero, i32 %a1, i32 %cnt
  ret i32 %res
}

define i64 @popcnt64_load_cmov(ptr %p0, i64 %a1) {
; CHECK-LABEL: popcnt64_load_cmov:
; CHECK:       # %bb.0:
; CHECK-NEXT:    popcntq (%rdi), %rax
; CHECK-NEXT:    cmoveq %rsi, %rax
; CHECK-NEXT:    retq
  %v = load i64, ptr %p0, align 8
  %cnt = tail call i64 @llvm.ctpop.i64(i64 %v)
  %iszero = icmp eq i64 %v, 0
  %res = select i1 %iszero, i64 %a1, i64 %cnt
  ret i64 %res
}

; --- negative: icmp against non-zero constant; load must NOT be folded ---

define i32 @lzcnt32_no_fold_nonzero_cmp(ptr %p0, i32 %a1) {
; CHECK-LABEL: lzcnt32_no_fold_nonzero_cmp:
; CHECK:       # %bb.0:
; CHECK-NEXT:    movl (%rdi), %ecx
; CHECK-NEXT:    lzcntl %ecx, %eax
; CHECK-NEXT:    cmpl $1, %ecx
; CHECK-NEXT:    cmovel %esi, %eax
; CHECK-NEXT:    retq
  %v = load i32, ptr %p0, align 4
  %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
  %isone = icmp eq i32 %v, 1
  %res = select i1 %isone, i32 %a1, i32 %cnt
  ret i32 %res
}

define i32 @lzcnt32_no_fold_var_cmp(ptr %p0, i32 %a1, i32 %a2) {
; CHECK-LABEL: lzcnt32_no_fold_var_cmp:
; CHECK:       # %bb.0:
; CHECK-NEXT:    movl (%rdi), %ecx
; CHECK-NEXT:    lzcntl %ecx, %eax
; CHECK-NEXT:    cmpl %edx, %ecx
; CHECK-NEXT:    cmovel %esi, %eax
; CHECK-NEXT:    retq
  %v = load i32, ptr %p0, align 4
  %cnt = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 %v, i1 true)
  %eq = icmp eq i32 %v, %a2
  %res = select i1 %eq, i32 %a1, i32 %cnt
  ret i32 %res
}

define i32 @nofold_load_with_intervening_store_1(ptr %p1, ptr %p2) {
; CHECK-LABEL: nofold_load_with_intervening_store_1:
; CHECK:       # %bb.0:
; CHECK-NEXT:    movq (%rdi), %rax
; CHECK-NEXT:    movl $4294967295, %ecx # imm = 0xFFFFFFFF
; CHECK-NEXT:    movq %rcx, (%rdi)
; CHECK-NEXT:    cmpq $2, %rax
; CHECK-NEXT:    setge (%rsi)
; CHECK-NEXT:    jl .LBB8_1
; CHECK-NEXT:  # %bb.2: # %taken
; CHECK-NEXT:    movq $0, (%rdi)
; CHECK-NEXT:    movl $1, %eax
; CHECK-NEXT:    retq
; CHECK-NEXT:  .LBB8_1: # %ret
; CHECK-NEXT:    xorl %eax, %eax
; CHECK-NEXT:    retq
  %v   = load i64, ptr %p1, align 8
  store i64 4294967295, ptr %p1, align 8
  %cmp = icmp sgt i64 %v, 1
  %ext = zext i1 %cmp to i8
  store i8 %ext, ptr %p2, align 1
  br i1 %cmp, label %taken, label %ret
ret:
  ret i32 0
taken:
  store i64 0, ptr %p1
  ret i32 1
}

define i64 @nofold_load_with_intervening_store_2(ptr %p, i64 %fallback) {
; CHECK-LABEL: nofold_load_with_intervening_store_2:
; CHECK:       # %bb.0:
; CHECK-NEXT:    movq (%rdi), %rax
; CHECK-NEXT:    movq $-1, (%rdi)
; CHECK-NEXT:    popcntq %rax, %rax
; CHECK-NEXT:    cmoveq %rsi, %rax
; CHECK-NEXT:    retq
  %v = load i64, ptr %p
  store i64 -1, ptr %p
  %cnt = tail call i64 @llvm.ctpop.i64(i64 %v)
  %iszero = icmp eq i64 %v, 0
  %res = select i1 %iszero, i64 %fallback, i64 %cnt
  ret i64 %res
}