[CGP] Freeze condition when despeculating ctlz/cttz

[nikic]

Summary

Freeze the condition of the newly introduced conditional branch, to avoid immediate undefined behavior if the input to ctlz/cttz was originally poison.

[spatel]

LGTM - is it worth adding a CodeGen equivalent test to show that doesn't change (assuming there is no change)?

[fhahn]

LGTM, thanks!

llvm/lib/CodeGen/CodeGenPrepare.cpp

2042

nit: while you are here, llvm:: could be dropped.

2064

nit: might be worth to name the freeze adding a fr. prefix or something to the name of the frozen value.

[nikic]

In D125887#3524715, @spatel wrote:

LGTM - is it worth adding a CodeGen equivalent test to show that doesn't change (assuming there is no change)?

It looks like we already have a couple of CodeGen tests for this, e.g. llvm/test/CodeGen/X86/clz.ll.

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.
-  if (llvm::isKnownNonZero(CountZeros->getOperand(0), *DL))
+  Value *Op = CountZeros->getOperand(0);
+  if (llvm::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);
-  Value *Cmp = Builder.CreateICmpEQ(CountZeros->getOperand(0), Zero, "cmpz");
+  if (!isGuaranteedNotToBeUndefOrPoison(Op))
+    Op = Builder.CreateFreeze(Op); // Avoid introducing branch on poison.
+  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.
   Builder.SetInsertPoint(&EndBlock->front());
   PHINode *PN = Builder.CreatePHI(Ty, 2, "ctz");
   CountZeros->replaceAllUsesWith(PN);

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -codegenprepare < %s | FileCheck %s --check-prefix=SLOW
 ; RUN: opt -S -codegenprepare -mattr=+bmi < %s | FileCheck %s --check-prefix=FAST_TZ
 ; RUN: opt -S -codegenprepare -mattr=+lzcnt < %s | FileCheck %s --check-prefix=FAST_LZ
 
 target triple = "x86_64-unknown-unknown"
 target datalayout = "e-n32:64"
 
 ; If the intrinsic is cheap, nothing should change.
 ; If the intrinsic is expensive, check if the input is zero to avoid the call.
 ; This is undoing speculation that may have been created by SimplifyCFG + InstCombine.
 
 define i64 @cttz(i64 %A) {
 ; SLOW-LABEL: @cttz(
 ; SLOW-NEXT:  entry:
-; SLOW-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[A:%.*]], 0
+; SLOW-NEXT:    [[TMP0:%.*]] = freeze i64 [[A:%.*]]
+; SLOW-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[TMP0]], 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:    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:    [[CMPZ:%.*]] = icmp eq i64 [[A:%.*]], 0
+; FAST_LZ-NEXT:    [[TMP0:%.*]] = freeze i64 [[A:%.*]]
+; FAST_LZ-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[TMP0]], 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:    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:    [[CMPZ:%.*]] = icmp eq i64 [[A:%.*]], 0
+; SLOW-NEXT:    [[TMP0:%.*]] = freeze i64 [[A:%.*]]
+; SLOW-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[TMP0]], 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:    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:    [[CMPZ:%.*]] = icmp eq i64 [[A:%.*]], 0
+; FAST_TZ-NEXT:    [[TMP0:%.*]] = freeze i64 [[A:%.*]]
+; FAST_TZ-NEXT:    [[CMPZ:%.*]] = icmp eq i64 [[TMP0]], 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:    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]]
 ;