[ValueTracking] Recognize more integer abs idioms

[frasercrmck]

Summary

This patch supports additional integer abs idioms where the comparison
is performed in a wider type than the abs result, i.e. where select's
comparison operands are sign-extended versions of the select's
true/false operands:

  %neg = sub i16 0, %a
  %a.prom = sext i16 %a to i32
  %abs.cond = icmp sge i32 %a.prom, 0
  %abs = select i1 %abs.cond, i16 %a, i16 %neg

[nikic]

This pattern already gets canonicalized by InstCombine (and there are no multi-use problems): https://llvm.godbolt.org/z/od9oq4 As such, it shouldn't be reaching the backend in this form.

[frasercrmck]

In D89955#2348423, @nikic wrote:

This pattern already gets canonicalized by InstCombine (and there are no multi-use problems): https://llvm.godbolt.org/z/od9oq4 As such, it shouldn't be reaching the backend in this form.

That's a good point, thanks. Sadly I don't have access to my motivating example as I've lost access to the code. I suppose it's possible that there was a pass between InstCombine and the SelectionDAG lowering which undid this canonicalization, but it'll be hard to find.

Out of curiosity, what do you mean by multi-use problems? Profitability or correctness?

[shchenz]

I suppose it's possible that there was a pass between InstCombine and the SelectionDAG lowering which undid this canonicalization, but it'll be hard to find.

The issue here is there is no instruction combine/simplify pass in llc pipeline which eliminates the %a.prom = sext i16 %a to i32 before we canonicalize abs.

In opt pipeline, InstCombine will eliminate the sext first and then can recognize it as a abs.

If we compile with clang which contain inst combine pass, we will get expected abs instruction.

[lebedev.ri]

Please can you explain the bigger problem here?
I'm not sure why the solution isn't: "just run the middle-end optimization passes, or don't expect good back-end codegen".

[frasercrmck]

In D89955#2355653, @lebedev.ri wrote:

Please can you explain the bigger problem here?
I'm not sure why the solution isn't: "just run the middle-end optimization passes, or don't expect good back-end codegen".

The original motivation was that when we moved from LLVM7 to LLVM9 we observed a degradation in the recognition of abs idioms (using clang, so including InstCombine @shchenz). I tracked it to the time at which abs idioms were moved out of DAGCombine and into the initial lowering from IR, where I found a potential for more patterns. It's possible something happened to our code in between InstCombine and ISel as our target had support for i8/i16 but only when extended/truncated to/from i32: exts & truncs were common and it was generally awkward in LLVM. Sadly this is all hypothetical since I've lost access to the code so I can't get at the full reproduction details.

I put this patch up anyway, since I thought that the ValueTracking pattern analysis had a function outside of the specific position(s) in the pipeline it's currently run. I reasoned that you might want to match this select pattern before InstCombine, for instance. If that's not ValueTracking's agreed-upon contract (which I'd understand since then its scope becomes much larger: it has to settle on some form of canonicalization) then your solution is the one we should take. It sounds to me like a question of determining which pass is responsible for what.

[lebedev.ri]

Since as @nikic said this pattern is canonicalized to the intrinsic already, i don't think this should be needed.

llvm/lib/Analysis/ValueTracking.cpp

static SelectPatternResult matchSelectPattern(CmpInst::Predicate Pred,
 
   if (isKnownNegation(TrueVal, FalseVal)) {
     // Sign-extending LHS does not change its sign, so TrueVal/FalseVal can
     // match against either LHS or sext(LHS).
     auto MaybeSExtCmpLHS =
         m_CombineOr(m_Specific(CmpLHS), m_SExt(m_Specific(CmpLHS)));
     auto ZeroOrAllOnes = m_CombineOr(m_ZeroInt(), m_AllOnes());
     auto ZeroOrOne = m_CombineOr(m_ZeroInt(), m_One());
-    if (match(TrueVal, MaybeSExtCmpLHS)) {
+    if (match(TrueVal, MaybeSExtCmpLHS) ||
+        match(CmpLHS, m_SExt(m_Specific(TrueVal)))) {
       // Set the return values. If the compare uses the negated value (-X >s 0),
       // swap the return values because the negated value is always 'RHS'.
       LHS = TrueVal;
       RHS = FalseVal;
       if (match(CmpLHS, m_Neg(m_Specific(FalseVal))))
         std::swap(LHS, RHS);
 
       // (X >s 0) ? X : -X or (X >s -1) ? X : -X --> ABS(X)
       // (-X >s 0) ? -X : X or (-X >s -1) ? -X : X --> ABS(X)
       if (Pred == ICmpInst::ICMP_SGT && match(CmpRHS, ZeroOrAllOnes))
         return {SPF_ABS, SPNB_NA, false};
 
       // (X >=s 0) ? X : -X or (X >=s 1) ? X : -X --> ABS(X)
       if (Pred == ICmpInst::ICMP_SGE && match(CmpRHS, ZeroOrOne))
         return {SPF_ABS, SPNB_NA, false};
 
       // (X <s 0) ? X : -X or (X <s 1) ? X : -X --> NABS(X)
       // (-X <s 0) ? -X : X or (-X <s 1) ? -X : X --> NABS(X)
       if (Pred == ICmpInst::ICMP_SLT && match(CmpRHS, ZeroOrOne))
         return {SPF_NABS, SPNB_NA, false};
-    }
-    else if (match(FalseVal, MaybeSExtCmpLHS)) {
+    } else if (match(FalseVal, MaybeSExtCmpLHS) ||
+               match(CmpLHS, m_SExt(m_Specific(FalseVal)))) {
       // Set the return values. If the compare uses the negated value (-X >s 0),
       // swap the return values because the negated value is always 'RHS'.
       LHS = FalseVal;
       RHS = TrueVal;
       if (match(CmpLHS, m_Neg(m_Specific(TrueVal))))
         std::swap(LHS, RHS);
 
       // (X >s 0) ? -X : X or (X >s -1) ? -X : X --> NABS(X)

llvm/test/CodeGen/NVPTX/idioms.ll

 ; Check that various LLVM idioms get lowered to NVPTX as expected.
 
 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s
 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s
 
 ; CHECK-LABEL: abs_i16(
 define i16 @abs_i16(i16 %a) {
 ; CHECK: abs.s16
   %neg = sub i16 0, %a
   %abs.cond = icmp sge i16 %a, 0
   %abs = select i1 %abs.cond, i16 %a, i16 %neg
   ret i16 %abs
 }
 
+; CHECK-LABEL: abs_extcmp_i16(
+define i16 @abs_extcmp_i16(i16 %a) {
+; CHECK: abs.s16
+  %neg = sub i16 0, %a
+  %a.prom = sext i16 %a to i32
+  %abs.cond = icmp sge i32 %a.prom, 0
+  %abs = select i1 %abs.cond, i16 %a, i16 %neg
+  ret i16 %abs
+}
+
 ; CHECK-LABEL: abs_i32(
 define i32 @abs_i32(i32 %a) {
 ; CHECK: abs.s32
   %neg = sub i32 0, %a
   %abs.cond = icmp sge i32 %a, 0
   %abs = select i1 %abs.cond, i32 %a, i32 %neg
   ret i32 %abs
 }