[ValueTracking] Match select abs pattern when there's an sext involved

[john.brawn]

Summary

When checking a select to see if it matches an abs, allow the true/false values to be a sign-extension of the comparison value instead of requiring that they're directly the comparison value, as all the comparison cares about is the sign of the value.

This fixes a regression due to r333702, where we were no longer generating ctlz due to isKnownNonNegative failing to match such a pattern.

[craig.topper]

lib/Analysis/ValueTracking.cpp

4618 ↗

(On Diff #149447)

I think it might be better to make LHS assigned to the sign extended value here.

[craig.topper]

lib/Analysis/ValueTracking.cpp

4612 ↗

(On Diff #149447)

Can we just check m_Neg(TrueVal)?

[john.brawn]

Update according to review comments.

[craig.topper]

LGTM

[mstorsjo]

It seems like this started triggering assertions in builds, I was about to file a bug about it, but it seems someone else already beat me to it, in https://bugs.llvm.org/show_bug.cgi?id=37686.

[john.brawn]

In D47631#1122027, @mstorsjo wrote:

It seems like this started triggering assertions in builds, I was about to file a bug about it, but it seems someone else already beat me to it, in https://bugs.llvm.org/show_bug.cgi?id=37686.

Fixed in https://reviews.llvm.org/rL334019.

llvm/trunk/lib/Analysis/ValueTracking.cpp

if (TrueVal == CmpLHS && FalseVal == CmpRHS) {
     case FCmpInst::FCMP_ULE:
     case FCmpInst::FCMP_OLT:
     case FCmpInst::FCMP_OLE: return {SPF_FMINNUM, NaNBehavior, Ordered};
     }
   }
 
   const APInt *C1;
   if (match(CmpRHS, m_APInt(C1))) {
-    if ((CmpLHS == TrueVal && match(FalseVal, m_Neg(m_Specific(CmpLHS)))) ||
-        (CmpLHS == FalseVal && match(TrueVal, m_Neg(m_Specific(CmpLHS))))) {
-      // Set RHS to the negate operand. LHS was assigned to CmpLHS earlier.
-      RHS = (CmpLHS == TrueVal) ? FalseVal : TrueVal;
+    // Sign-extending LHS does not change its sign, so TrueVal/FalseVal can
+    // match against either LHS or sext(LHS).
+    auto MaybeSExtLHS = m_CombineOr(m_Specific(CmpLHS),
+                                    m_SExt(m_Specific(CmpLHS)));
+    if ((match(TrueVal, MaybeSExtLHS) &&
+         match(FalseVal, m_Neg(m_Specific(TrueVal)))) ||
+        (match(FalseVal, MaybeSExtLHS) &&
+         match(TrueVal, m_Neg(m_Specific(FalseVal))))) {
+      // Set LHS and RHS so that RHS is the negated operand of the select
+      if (match(TrueVal, MaybeSExtLHS)) {
+        LHS = TrueVal;
+        RHS = FalseVal;
+      } else {
+        LHS = FalseVal;
+        RHS = TrueVal;
+      }
 
       // ABS(X) ==> (X >s 0) ? X : -X and (X >s -1) ? X : -X
       // NABS(X) ==> (X >s 0) ? -X : X and (X >s -1) ? -X : X
       if (Pred == ICmpInst::ICMP_SGT &&
           (C1->isNullValue() || C1->isAllOnesValue())) {
-        return {(CmpLHS == TrueVal) ? SPF_ABS : SPF_NABS, SPNB_NA, false};
+        return {(LHS == TrueVal) ? SPF_ABS : SPF_NABS, SPNB_NA, false};
       }
 
       // ABS(X) ==> (X <s 0) ? -X : X and (X <s 1) ? -X : X
       // NABS(X) ==> (X <s 0) ? X : -X and (X <s 1) ? X : -X
       if (Pred == ICmpInst::ICMP_SLT &&
           (C1->isNullValue() || C1->isOneValue())) {
-        return {(CmpLHS == FalseVal) ? SPF_ABS : SPF_NABS, SPNB_NA, false};
+        return {(LHS == FalseVal) ? SPF_ABS : SPF_NABS, SPNB_NA, false};
       }
     }
   }
 
   if (CmpInst::isIntPredicate(Pred))
     return matchMinMax(Pred, CmpLHS, CmpRHS, TrueVal, FalseVal, LHS, RHS, Depth);
 
   // According to (IEEE 754-2008 5.3.1), minNum(0.0, -0.0) and similar

llvm/trunk/test/Transforms/LoopIdiom/ARM/ctlz.ll

while.cond:                                       ; preds = %while.cond, %entry
   %shr = ashr i32 %n.addr.0, 1
   %tobool = icmp eq i32 %shr, 0
   %inc = add nsw i32 %i.0, 1
   br i1 %tobool, label %while.end, label %while.cond
 
 while.end:                                        ; preds = %while.cond
   ret i32 %i.0
 }
+
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_sext(short in)
+; {
+;   int n = in;
+;   if (in < 0)
+;     n = -n;
+;   int i = 0;
+;   while(n >>= 1) {
+;     i++;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = ashr i32 %abs_n, 1
+; ALL-NEXT:  %1 = call i32 @llvm.ctlz.i32(i32 %0, i1 false)
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL-NEXT:  %3 = add i32 %2, 1
+; ALL:  %i.0.lcssa = phi i32 [ %2, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_sext(i16 %in) {
+entry:
+  %n = sext i16 %in to i32
+  %c = icmp sgt i16 %in, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %abs_n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %while.cond ]
+  %shr = ashr i32 %n.addr.0, 1
+  %tobool = icmp eq i32 %shr, 0
+  %inc = add nsw i32 %i.0, 1
+  br i1 %tobool, label %while.end, label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 %i.0
+}

llvm/trunk/test/Transforms/LoopIdiom/X86/ctlz.ll

while.cond:                                       ; preds = %while.cond, %entry
   %tobool = icmp eq i32 %shr, 0
   %inc = add nsw i32 %i.0, 1
   br i1 %tobool, label %while.end, label %while.cond
 
 while.end:                                        ; preds = %while.cond
   ret i32 %i.0
 }
 
+; Recognize CTLZ builtin pattern.
+; Here it will replace the loop -
+; assume builtin is always profitable.
+;
+; int ctlz_sext(short in)
+; {
+;   int n = in;
+;   if (in < 0)
+;     n = -n;
+;   int i = 0;
+;   while(n >>= 1) {
+;     i++;
+;   }
+;   return i;
+; }
+;
+; ALL:  entry
+; ALL:  %0 = ashr i32 %abs_n, 1
+; ALL-NEXT:  %1 = call i32 @llvm.ctlz.i32(i32 %0, i1 false)
+; ALL-NEXT:  %2 = sub i32 32, %1
+; ALL-NEXT:  %3 = add i32 %2, 1
+; ALL:  %i.0.lcssa = phi i32 [ %2, %while.cond ]
+; ALL:  ret i32 %i.0.lcssa
+
+; Function Attrs: norecurse nounwind readnone uwtable
+define i32 @ctlz_sext(i16 %in) {
+entry:
+  %n = sext i16 %in to i32
+  %c = icmp sgt i16 %in, 0
+  %negn = sub nsw i32 0, %n
+  %abs_n = select i1 %c, i32 %n, i32 %negn
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.cond, %entry
+  %n.addr.0 = phi i32 [ %abs_n, %entry ], [ %shr, %while.cond ]
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %while.cond ]
+  %shr = ashr i32 %n.addr.0, 1
+  %tobool = icmp eq i32 %shr, 0
+  %inc = add nsw i32 %i.0, 1
+  br i1 %tobool, label %while.end, label %while.cond
+
+while.end:                                        ; preds = %while.cond
+  ret i32 %i.0
+}
+
 ; This loop contains a volatile store. If x is initially negative,
 ; the code will be an infinite loop because the ashr will eventually produce
 ; all ones and continue doing so. This prevents the loop from terminating. If
 ; we convert this to a countable loop using ctlz that loop will only run 32
 ; times. This is different than the infinite number of times of the original.
 define i32 @foo(i32 %x) {
 ; LZCNT-LABEL: @foo(
 ; LZCNT-NEXT:  entry: