[InstCombine] Relax cttz/ctlz with select on zero

[nikic]

Summary

The cttz/ctlz intrinsics have a parameter specifying whether the result is undefined for zero. cttz(x, false) can be relaxed to cttz(x, true) if x is known non-zero, and in fact such an optimization is already performed. However, this currently doesn't work if x is non-zero as a result of a select rather than an explicit branch. This patch adds handling for this case, thus allowing x != 0 ? cttz(x, false) : y to simplify to x != 0 ? cttz(x, true) : y.

[nikic]

Ping

[spatel]

LGTM - please modify or add a test to use vector types, so we have some coverage for that possibility.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (Pred == ICmpInst::ICMP_NE)
     std::swap(Count, ValueOnZero);
 
   // Skip zero extend/truncate.
   Value *V = nullptr;
   if (match(Count, m_ZExt(m_Value(V))) ||
       match(Count, m_Trunc(m_Value(V))))
     Count = V;
 
-  // Check if the value propagated on zero is a constant number equal to the
-  // sizeof in bits of 'Count'.
-  unsigned SizeOfInBits = Count->getType()->getScalarSizeInBits();
-  if (!match(ValueOnZero, m_SpecificInt(SizeOfInBits)))
-    return nullptr;
-
   // Check that 'Count' is a call to intrinsic cttz/ctlz. Also check that the
   // input to the cttz/ctlz is used as LHS for the compare instruction.
-  if (match(Count, m_Intrinsic<Intrinsic::cttz>(m_Specific(CmpLHS))) ||
-      match(Count, m_Intrinsic<Intrinsic::ctlz>(m_Specific(CmpLHS)))) {
-    IntrinsicInst *II = cast<IntrinsicInst>(Count);
+  if (!match(Count, m_Intrinsic<Intrinsic::cttz>(m_Specific(CmpLHS))) &&
+      !match(Count, m_Intrinsic<Intrinsic::ctlz>(m_Specific(CmpLHS))))
+    return nullptr;
+
+
+  // Check if the value propagated on zero is a constant number equal to the
+  // sizeof in bits of 'Count'.
+  unsigned SizeOfInBits = Count->getType()->getScalarSizeInBits();
+  if (match(ValueOnZero, m_SpecificInt(SizeOfInBits))) {
     // Explicitly clear the 'undef_on_zero' flag.
     IntrinsicInst *NewI = cast<IntrinsicInst>(II->clone());
     NewI->setArgOperand(1, ConstantInt::getFalse(NewI->getContext()));
     Builder.Insert(NewI);
     return Builder.CreateZExtOrTrunc(NewI, ValueOnZero->getType());
   }
 
+  // If the ValueOnZero is not the bitwidth, we can at least make use of the
+  // fact that the cttz/ctlz result will not be used if the input is zero, so
+  // it's okay to relax it to undef for that case.
+  if (II->hasOneUse() && !match(II->getArgOperand(1), m_One()))
+    II->setArgOperand(1, ConstantInt::getTrue(II->getContext()));
+
   return nullptr;
 }
 
 /// Return true if we find and adjust an icmp+select pattern where the compare
 /// is with a constant that can be incremented or decremented to match the
 /// minimum or maximum idiom.
 static bool adjustMinMax(SelectInst &Sel, ICmpInst &Cmp) {
   ICmpInst::Predicate Pred = Cmp.getPredicate();

test/Transforms/InstCombine/select-cmp-cttz-ctlz.ll

;
   %ct = tail call i128 @llvm.cttz.i128(i128 %x, i1 true)
   %tobool = icmp ne i128 %x, 0
   %cond = select i1 %tobool, i128 %ct, i128 128
   ret i128 %cond
 }
 
 define i32 @test_ctlz_not_bw(i32 %x) {
 ; CHECK-LABEL: @test_ctlz_not_bw(
-; CHECK-NEXT:    [[CT:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range !1
+; CHECK-NEXT:    [[CT:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 true), !range !1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 0
 ; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 123, i32 [[CT]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %ct = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %cmp = icmp ne i32 %x, 0
   %res = select i1 %cmp, i32 %ct, i32 123
   ret i32 %res
;
   %cmp = icmp ne i32 %x, 0
   %sel = select i1 %cmp, i32 %ct, i32 123
   %res = or i32 %sel, %ct
   ret i32 %res
 }
 
 define i32 @test_cttz_not_bw(i32 %x) {
 ; CHECK-LABEL: @test_cttz_not_bw(
-; CHECK-NEXT:    [[CT:%.*]] = tail call i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 false), !range !1
+; CHECK-NEXT:    [[CT:%.*]] = tail call i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 true), !range !1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X]], 0
 ; CHECK-NEXT:    [[RES:%.*]] = select i1 [[CMP]], i32 123, i32 [[CT]]
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %ct = tail call i32 @llvm.cttz.i32(i32 %x, i1 false)
   %cmp = icmp ne i32 %x, 0
   %res = select i1 %cmp, i32 %ct, i32 123
   ret i32 %res