[CGP] Also freeze ctlz/cttz operand when despeculating

[nikic]

Summary

D125887 changed the ctlz/cttz despeculation transform to insert a freeze for the introduced branch on zero. While this does fix the "branch on poison" issue, we may still get in trouble if we pick a different value for the branch and for the ctz argument (i.e. non-zero for the branch, but zero for the ctz). To avoid this, we should use the same frozen value in both positions.

This does cause a regression in RISCV codegen by introducing an additional sext. The DAG looks like this:

SelectionDAG has 22 nodes:
  t0: ch = EntryToken
      t2: i64,ch = CopyFromReg t0, Register:i64 %3
    t4: i64 = AssertSext t2, ValueType:ch:i32
  t23: i64 = freeze t4
        t9: ch = CopyToReg t0, Register:i64 %0, t23
        t16: ch = CopyToReg t0, Register:i64 %4, Constant:i64<32>
      t18: ch = TokenFactor t9, t16
          t25: i64 = sign_extend_inreg t23, ValueType:ch:i32
        t24: i64 = setcc t25, Constant:i64<0>, seteq:ch
      t28: i64 = and t24, Constant:i64<1>
    t19: ch = brcond t18, t28, BasicBlock:ch<cond.end 0x8311f68>
  t21: ch = br t19, BasicBlock:ch<cond.false 0x8311e80>

I don't see a really obvious way to improve this, as we can't push the freeze past the AssertSext (which may produce poison).

[nikic]

ping

[spatel]

https://alive2.llvm.org/ce/z/iwd4ba
This is needed for correctness, so LGTM.

It might be good to add tests for other targets. I didn't step through this to see why, but the code for AArch and ARM actually looks better when we re-use the frozen value.

rbit	w9, w0
mov	w8, #32
clz	w9, w9
cmp	w0, #0
csel	w0, w8, w9, eq

vs.

rbit	w8, w0
clz	w0, w8

[nikic]

In D126638#3566645, @spatel wrote:

https://alive2.llvm.org/ce/z/iwd4ba
This is needed for correctness, so LGTM.

It might be good to add tests for other targets. I didn't step through this to see why, but the code for AArch and ARM actually looks better when we re-use the frozen value.

rbit	w9, w0
mov	w8, #32
clz	w9, w9
cmp	w0, #0
csel	w0, w8, w9, eq

vs.

rbit	w8, w0
clz	w0, w8

For what input is that? As far as I can tell, AArch64 never uses despeculation, and ARM uses a libcall legalization where it makes no difference. I'm trying this one:

define i32 @ctlz_i32(i32 %a) nounwind {
  %r = call i32 @llvm.ctlz.i32(i32 %a, i1 false)
  ret i32 %r
}

declare i32 @llvm.ctlz.i32(i32, i1)

[spatel]

In D126638#3569452, @nikic wrote:

For what input is that? As far as I can tell, AArch64 never uses despeculation, and ARM uses a libcall legalization where it makes no difference. I'm trying this one:

I was experimenting with tests based on the IR in the regression test in this patch, so just changing the frozen value alone after doing the despeculation:

define i32 @tgt(i32 %A) {
entry:
  %f = freeze i32 %A
  %i = icmp eq i32 %f, 0
  br i1 %i, label %end, label %false
false:
  %z = call i32 @llvm.ctlz.i32(i32 %A, i1 true)
  br label %end
end:
  %ctz = phi i32 [ 32, %entry ], [ %z, %false ]
  ret i32 %ctz
}

define i32 @tgt2(i32 %A) {
entry:
  %f = freeze i32 %A
  %i = icmp eq i32 %f, 0
  br i1 %i, label %end, label %false
false:
  %z = call i32 @llvm.ctlz.i32(i32 %f, i1 true)
  br label %end
end:
  %ctz = phi i32 [ 32, %entry ], [ %z, %false ]
  ret i32 %ctz
}

And then sending that through llc with:
$ llc -o - ctlz.ll -mtriple=aarch64
$ llc -o - ctlz.ll -mtriple=armv7

I used 'v7' with arm; I do see the libcalls with base 'arm' codegen.

[nikic]

In D126638#3570001, @spatel wrote:

In D126638#3569452, @nikic wrote:

For what input is that? As far as I can tell, AArch64 never uses despeculation, and ARM uses a libcall legalization where it makes no difference. I'm trying this one:

I was experimenting with tests based on the IR in the regression test in this patch, so just changing the frozen value alone after doing the despeculation:
[...]
And then sending that through llc with:
$ llc -o - ctlz.ll -mtriple=aarch64
$ llc -o - ctlz.ll -mtriple=armv7

I used 'v7' with arm; I do see the libcalls with base 'arm' codegen.

Okay, I see. I don't think adding those tests would be super meaningful, because both aarch64/armv7 disable ctz despeculation via TargetLowering hooks, so I don't think we'd particularly care how the pattern would optimize.

llvm/lib/CodeGen/CodeGenPrepare.cpp

static bool despeculateCountZeros(IntrinsicInst *CountZeros,
 
   // Only handle legal scalar cases. Anything else requires too much work.
   Type *Ty = CountZeros->getType();
   unsigned SizeInBits = Ty->getScalarSizeInBits();
   if (Ty->isVectorTy() || SizeInBits > DL->getLargestLegalIntTypeSizeInBits())
     return false;
 
   // Bail if the value is never zero.
-  Value *Op = CountZeros->getOperand(0);
+  Use &Op = CountZeros->getOperandUse(0);
   if (isKnownNonZero(Op, *DL))
     return false;
 
   // The intrinsic will be sunk behind a compare against zero and branch.
   BasicBlock *StartBlock = CountZeros->getParent();
   BasicBlock *CallBlock = StartBlock->splitBasicBlock(CountZeros, "cond.false");
 
   // Create another block after the count zero intrinsic. A PHI will be added
   // in this block to select the result of the intrinsic or the bit-width
   // constant if the input to the intrinsic is zero.
   BasicBlock::iterator SplitPt = ++(BasicBlock::iterator(CountZeros));
   BasicBlock *EndBlock = CallBlock->splitBasicBlock(SplitPt, "cond.end");
 
   // Set up a builder to create a compare, conditional branch, and PHI.
   IRBuilder<> Builder(CountZeros->getContext());
   Builder.SetInsertPoint(StartBlock->getTerminator());
   Builder.SetCurrentDebugLocation(CountZeros->getDebugLoc());
 
   // Replace the unconditional branch that was created by the first split with
   // a compare against zero and a conditional branch.
   Value *Zero = Constant::getNullValue(Ty);
-  // Avoid introducing branch on poison.
+  // Avoid introducing branch on poison. This also replaces the ctz operand.
   if (!isGuaranteedNotToBeUndefOrPoison(Op))
     Op = Builder.CreateFreeze(Op, Op->getName() + ".fr");
   Value *Cmp = Builder.CreateICmpEQ(Op, Zero, "cmpz");
   Builder.CreateCondBr(Cmp, EndBlock, CallBlock);
   StartBlock->getTerminator()->eraseFromParent();
 
   // Create a PHI in the end block to select either the output of the intrinsic
   // or the bit width of the operand.

llvm/test/CodeGen/RISCV/rv64zbb.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV64I
 ; RUN: llc -mtriple=riscv64 -mattr=+zbb -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV64ZBB
 
 declare i32 @llvm.ctlz.i32(i32, i1)
 
 define signext i32 @ctlz_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: ctlz_i32:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi sp, sp, -16
 ; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
-; RV64I-NEXT:    beqz a0, .LBB0_2
+; RV64I-NEXT:    sext.w a1, a0
+; RV64I-NEXT:    beqz a1, .LBB0_2
 ; RV64I-NEXT:  # %bb.1: # %cond.false
 ; RV64I-NEXT:    srliw a1, a0, 1
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 2
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 4
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 8
; RV64ZBB-NEXT:    ret
   ret i32 %1
 }
 
 define signext i32 @log2_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: log2_i32:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi sp, sp, -16
 ; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
-; RV64I-NEXT:    beqz a0, .LBB1_2
+; RV64I-NEXT:    sext.w a1, a0
+; RV64I-NEXT:    beqz a1, .LBB1_2
 ; RV64I-NEXT:  # %bb.1: # %cond.false
 ; RV64I-NEXT:    srliw a1, a0, 1
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 2
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 4
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    srliw a1, a0, 8

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi sp, sp, -16
 ; RV64I-NEXT:    .cfi_def_cfa_offset 16
 ; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
 ; RV64I-NEXT:    .cfi_offset ra, -8
 ; RV64I-NEXT:    srliw a0, a0, 1
 ; RV64I-NEXT:    beqz a0, .LBB4_2
 ; RV64I-NEXT:  # %bb.1: # %cond.false
-; RV64I-NEXT:    srli a1, a0, 1
+; RV64I-NEXT:    srliw a1, a0, 1
 ; RV64I-NEXT:    or a0, a0, a1
-; RV64I-NEXT:    srli a1, a0, 2
+; RV64I-NEXT:    srliw a1, a0, 2
 ; RV64I-NEXT:    or a0, a0, a1
-; RV64I-NEXT:    srli a1, a0, 4
+; RV64I-NEXT:    srliw a1, a0, 4
 ; RV64I-NEXT:    or a0, a0, a1
-; RV64I-NEXT:    srli a1, a0, 8
+; RV64I-NEXT:    srliw a1, a0, 8
 ; RV64I-NEXT:    or a0, a0, a1
-; RV64I-NEXT:    slli a1, a0, 33
-; RV64I-NEXT:    srli a1, a1, 49
+; RV64I-NEXT:    srliw a1, a0, 16
 ; RV64I-NEXT:    or a0, a0, a1
 ; RV64I-NEXT:    not a0, a0
 ; RV64I-NEXT:    srli a1, a0, 1
 ; RV64I-NEXT:    lui a2, 349525
 ; RV64I-NEXT:    addiw a2, a2, 1365
 ; RV64I-NEXT:    and a1, a1, a2
 ; RV64I-NEXT:    subw a0, a0, a1
 ; RV64I-NEXT:    lui a1, 209715

 
 declare i32 @llvm.cttz.i32(i32, i1)
 
 define signext i32 @cttz_i32(i32 signext %a) nounwind {
 ; RV64I-LABEL: cttz_i32:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    addi sp, sp, -16
 ; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
-; RV64I-NEXT:    beqz a0, .LBB6_2
+; RV64I-NEXT:    sext.w a1, a0
+; RV64I-NEXT:    beqz a1, .LBB6_2
 ; RV64I-NEXT:  # %bb.1: # %cond.false
 ; RV64I-NEXT:    addiw a1, a0, -1
 ; RV64I-NEXT:    not a0, a0
 ; RV64I-NEXT:    and a0, a0, a1
 ; RV64I-NEXT:    srli a1, a0, 1
 ; RV64I-NEXT:    lui a2, 349525
 ; RV64I-NEXT:    addiw a2, a2, 1365
 ; RV64I-NEXT:    and a1, a1, a2

llvm/test/Transforms/CodeGenPrepare/X86/cttz-ctlz.ll

 
 define i64 @cttz(i64 %A) {
 ; SLOW-LABEL: @cttz(
 ; SLOW-NEXT:  entry:
 ; SLOW-NEXT:    [[A_FR:%.*]] = freeze i64 [[A:%.*]]
 ; SLOW-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[A_FR]], 0
 ; SLOW-NEXT:    br i1 [[CMPZ]], label [[COND_END:%.*]], label [[COND_FALSE:%.*]]
 ; SLOW:       cond.false:
-; SLOW-NEXT:    [[Z:%.*]] = call i64 @llvm.cttz.i64(i64 [[A]], i1 true)
+; SLOW-NEXT:    [[Z:%.*]] = call i64 @llvm.cttz.i64(i64 [[A_FR]], i1 true)
 ; SLOW-NEXT:    br label [[COND_END]]
 ; SLOW:       cond.end:
 ; SLOW-NEXT:    [[CTZ:%.*]] = phi i64 [ 64, [[ENTRY:%.*]] ], [ [[Z]], [[COND_FALSE]] ]
 ; SLOW-NEXT:    ret i64 [[CTZ]]
 ;
 ; FAST_TZ-LABEL: @cttz(
 ; FAST_TZ-NEXT:  entry:
 ; FAST_TZ-NEXT:    [[Z:%.*]] = call i64 @llvm.cttz.i64(i64 [[A:%.*]], i1 false)
 ; FAST_TZ-NEXT:    ret i64 [[Z]]
 ;
 ; FAST_LZ-LABEL: @cttz(
 ; FAST_LZ-NEXT:  entry:
 ; FAST_LZ-NEXT:    [[A_FR:%.*]] = freeze i64 [[A:%.*]]
 ; FAST_LZ-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[A_FR]], 0
 ; FAST_LZ-NEXT:    br i1 [[CMPZ]], label [[COND_END:%.*]], label [[COND_FALSE:%.*]]
 ; FAST_LZ:       cond.false:
-; FAST_LZ-NEXT:    [[Z:%.*]] = call i64 @llvm.cttz.i64(i64 [[A]], i1 true)
+; FAST_LZ-NEXT:    [[Z:%.*]] = call i64 @llvm.cttz.i64(i64 [[A_FR]], i1 true)
 ; FAST_LZ-NEXT:    br label [[COND_END]]
 ; FAST_LZ:       cond.end:
 ; FAST_LZ-NEXT:    [[CTZ:%.*]] = phi i64 [ 64, [[ENTRY:%.*]] ], [ [[Z]], [[COND_FALSE]] ]
 ; FAST_LZ-NEXT:    ret i64 [[CTZ]]
 ;
 entry:
   %z = call i64 @llvm.cttz.i64(i64 %A, i1 false)
   ret i64 %z
 }
 
 define i64 @ctlz(i64 %A) {
 ; SLOW-LABEL: @ctlz(
 ; SLOW-NEXT:  entry:
 ; SLOW-NEXT:    [[A_FR:%.*]] = freeze i64 [[A:%.*]]
 ; SLOW-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[A_FR]], 0
 ; SLOW-NEXT:    br i1 [[CMPZ]], label [[COND_END:%.*]], label [[COND_FALSE:%.*]]
 ; SLOW:       cond.false:
-; SLOW-NEXT:    [[Z:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A]], i1 true)
+; SLOW-NEXT:    [[Z:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A_FR]], i1 true)
 ; SLOW-NEXT:    br label [[COND_END]]
 ; SLOW:       cond.end:
 ; SLOW-NEXT:    [[CTZ:%.*]] = phi i64 [ 64, [[ENTRY:%.*]] ], [ [[Z]], [[COND_FALSE]] ]
 ; SLOW-NEXT:    ret i64 [[CTZ]]
 ;
 ; FAST_TZ-LABEL: @ctlz(
 ; FAST_TZ-NEXT:  entry:
 ; FAST_TZ-NEXT:    [[A_FR:%.*]] = freeze i64 [[A:%.*]]
 ; FAST_TZ-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[A_FR]], 0
 ; FAST_TZ-NEXT:    br i1 [[CMPZ]], label [[COND_END:%.*]], label [[COND_FALSE:%.*]]
 ; FAST_TZ:       cond.false:
-; FAST_TZ-NEXT:    [[Z:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A]], i1 true)
+; FAST_TZ-NEXT:    [[Z:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A_FR]], i1 true)
 ; FAST_TZ-NEXT:    br label [[COND_END]]
 ; FAST_TZ:       cond.end:
 ; FAST_TZ-NEXT:    [[CTZ:%.*]] = phi i64 [ 64, [[ENTRY:%.*]] ], [ [[Z]], [[COND_FALSE]] ]
 ; FAST_TZ-NEXT:    ret i64 [[CTZ]]
 ;
 ; FAST_LZ-LABEL: @ctlz(
 ; FAST_LZ-NEXT:  entry:
 ; FAST_LZ-NEXT:    [[Z:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A:%.*]], i1 false)