[AArch64] Try to re-use extended operand for SETCC with vector ops.

[fhahn]

Summary

Try to re-use an already extended operand for SetCC with vector operands
feeding an extended select. Doing so avoids requiring another full
extension of the SET_CC result when lowering the select.

This improves lowering for certain extend/cmp/select patterns operating.
For example with v16i8, this replaces 6 instructions for the extra extension
with 4 separate selects.

This improves the generated code for loops like the one below in
combination with D96522.

int foo(uint8_t *p, int N) {
  unsigned long long sum = 0;
  for (int i = 0; i < N ; i++, p++) {
    unsigned int v = *p;
    sum += (v < 127) ? v : 256 - v;
  }
  return sum;
}

https://clang.godbolt.org/z/Wco866MjY

On the AArch64 cores I have access to, the patch improves performance of
the vector loop by ~10%.

This could be generalized per follow-ups, but the initial version
targets one of the more important cases in combination with D96522.

Alive2 modeling:

• sext EQ https://alive2.llvm.org/ce/z/5upBvb
• sext NE https://alive2.llvm.org/ce/z/zbEcJp
• zext EQ https://alive2.llvm.org/ce/z/_xMwof
• zext NE https://alive2.llvm.org/ce/z/5FwKfc
• zext unsigned predicate: https://alive2.llvm.org/ce/z/iEwLU3
• sext signed predicate: https://alive2.llvm.org/ce/z/aMBega

[dmgreen]

Sorry for the delay. It took a while to find the time to figure out what was wrong, but it turned out that was in the other patch.

The change sounds OK to me, although it might be useful to generalize it to handle more types. This seems like it is useful in a few more cases: https://godbolt.org/z/fo8MMrqrE. I'm not sure if this generalises to not just select uses.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17062

Do you have a test for multiple uses?

18741

Whitespace

[fhahn]

Rebased on top of new test case, removed whitespace change.

In D120481#3353908, @dmgreen wrote:

Sorry for the delay. It took a while to find the time to figure out what was wrong, but it turned out that was in the other patch.

The change sounds OK to me, although it might be useful to generalize it to handle more types. This seems like it is useful in a few more cases: https://godbolt.org/z/fo8MMrqrE. I'm not sure if this generalises to not just select uses.

Thanks for taking a look! Agreed that it will be helpful for a couple of other type combinations. But maybe the generalization should be done as follow-ups?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17062

I added one in 3836003e87ee, thanks!

18741

Removed, thanks!

[dmgreen]

Thanks for taking a look! Agreed that it will be helpful for a couple of other type combinations. But maybe the generalization should be done as follow-ups?

Any reason to do that in a followup, as opposed to doing it here?

[fhahn]

Update to lift restrictions on specific vector types.

In D120481#3366872, @dmgreen wrote:

Thanks for taking a look! Agreed that it will be helpful for a couple of other type combinations. But maybe the generalization should be done as follow-ups?

Any reason to do that in a followup, as opposed to doing it here?

Mostly time constraints :)

I finally had some time to generalize this to support vector types in general

[dmgreen]

This sounds good as far as I can tell. If we are dealing with a dag combine, do we need to worry about odd types at all? And I think we can include eq and ne too, if alive is to be believed.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17069

I think this can skip the getScalarType call: UseMVT.getScalarSizeInBits() <= Op0MVT.getScalarSizeInBits()

17091

Could these include eq and ne conditions too?

llvm/test/CodeGen/AArch64/vselect-ext.ll

103–104

Can you remove all of these with the lower cases.

[fhahn]

In D120481#3632283, @dmgreen wrote:

This sounds good as far as I can tell. If we are dealing with a dag combine, do we need to worry about odd types at all?

Do you mean worry in terms of whether it is optimal? There are a few test cases with odd types, and it seems like they are handled OK on that front.

And I think we can include eq and ne too, if alive is to be believed.

Thanks, I added support for EQ/NE. The type of extend shouldn't matter, as long as we chose the existing one.

• sext EQ https://alive2.llvm.org/ce/z/5upBvb
• sext NE https://alive2.llvm.org/ce/z/zbEcJp
• zext EQ https://alive2.llvm.org/ce/z/_xMwof
• zext NE https://alive2.llvm.org/ce/z/5FwKfc

[dmgreen]

In D120481#3634566, @fhahn wrote:

In D120481#3632283, @dmgreen wrote:

This sounds good as far as I can tell. If we are dealing with a dag combine, do we need to worry about odd types at all?

Do you mean worry in terms of whether it is optimal? There are a few test cases with odd types, and it seems like they are handled OK on that front.

I meant for awkward types like i23 and i128. They don't have to be optimal, so long as they work OK. Just so long as we have a couple of tests, to make sure it doesnt run into problems.

The change sounds good with a small simplification suggestion. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17090

I think this might be able to combine the if blocks together for the ne/eq conditions:

if (Op0SExt && (isSignedIntSetCC(CC) || isIntEqualitySetCC(CC))) {
  Op0ExtV = SDValue(Op0SExt, 0);
  Op1ExtV = DAG.getNode(ISD::SIGN_EXTEND, DL, UseMVT, Op->getOperand(1));
} else if (Op0ZExt && (isUnsignedIntSetCC(CC) || isIntEqualitySetCC(CC))) {
  Op0ExtV = SDValue(Op0ZExt, 0);
  Op1ExtV = DAG.getNode(ISD::ZERO_EXTEND, DL, UseMVT, Op->getOperand(1));
} else
  return SDValue();

[fhahn]

Rebase on top of additional tests.

[fhahn]

Simplify check as suggested.

[fhahn]

In D120481#3634726, @dmgreen wrote:

In D120481#3634566, @fhahn wrote:

In D120481#3632283, @dmgreen wrote:

This sounds good as far as I can tell. If we are dealing with a dag combine, do we need to worry about odd types at all?

Do you mean worry in terms of whether it is optimal? There are a few test cases with odd types, and it seems like they are handled OK on that front.

I meant for awkward types like i23 and i128. They don't have to be optimal, so long as they work OK. Just so long as we have a couple of tests, to make sure it doesnt run into problems.

Ah got it! I added a few tests with vectors of i13 and i15. Looks like this handled well, because they are first converted to legal vector types.

The change sounds good with a small simplification suggestion. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17069

Thanks, simplified!

17090

Simplified as suggested, thanks!

17091

Yes, updated!

llvm/test/CodeGen/AArch64/vselect-ext.ll

103–104

Yeah I removed those. They were added by accident :(

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue performCSELCombine(SDNode *N,
                                   SelectionDAG &DAG) {
   // CSEL x, x, cc -> x
   if (N->getOperand(0) == N->getOperand(1))
     return N->getOperand(0);
 
   return performCONDCombine(N, DCI, DAG, 2, 3);
 }
 
+// Try to re-use an already extended operand of a v16i8 SetCC feeding a
+// extended select. Doing so avoids requiring another full extension of the
+// SET_CC result when lowering the select.
+static SDValue tryToWidenSetCCOperands(SDNode *Op, SelectionDAG &DAG) {
+  if (Op->getOperand(0).getValueType() != MVT::v16i8)
+    return SDValue();
+
+  if (any_of(Op->uses(), [](const SDNode *N) {
+        return N->getOpcode() != ISD::VSELECT ||
+               N->getValueType(0) != MVT::v16i32;
+      }))
+    return SDValue();
+
+  APInt V;
+  if (!ISD::isConstantSplatVector(Op->getOperand(1).getNode(), V))
+    return SDValue();
+
+  SDLoc DL(Op);
+  SDValue Op0ExtV;
+  SDValue Op1ExtV;
+  ISD::CondCode CC = cast<CondCodeSDNode>(Op->getOperand(2))->get();
+  // Check if the first operand of the SET_CC is already extended. If it is,
+  // split the SET_CC and re-use the extended version of the operand.
+  SDNode *Op0SExt = DAG.getNodeIfExists(
+      ISD::SIGN_EXTEND, DAG.getVTList(MVT::v16i32), Op->getOperand(0));
+  SDNode *Op0ZExt = DAG.getNodeIfExists(
+      ISD::ZERO_EXTEND, DAG.getVTList(MVT::v16i32), Op->getOperand(0));
+  if (Op0SExt && isSignedIntSetCC(CC)) {
+    Op0ExtV = SDValue(Op0SExt, 0);
+    Op1ExtV = DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::v16i32, Op->getOperand(1));
+  } else if (Op0ZExt && isUnsignedIntSetCC(CC)) {
+    Op0ExtV = SDValue(Op0ZExt, 0);
+    Op1ExtV = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::v16i32, Op->getOperand(1));
+  } else
+    return SDValue();
+
+  return DAG.getNode(ISD::SETCC, DL, MVT::v16i1, Op0ExtV, Op1ExtV,
+                     Op->getOperand(2));
+}
+
 static SDValue performSETCCCombine(SDNode *N, SelectionDAG &DAG) {
   assert(N->getOpcode() == ISD::SETCC && "Unexpected opcode!");
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   ISD::CondCode Cond = cast<CondCodeSDNode>(N->getOperand(2))->get();
 
+  if (SDValue V = tryToWidenSetCCOperands(N, DAG))
+    return V;
+
   // setcc (csel 0, 1, cond, X), 1, ne ==> csel 0, 1, !cond, X
   if (Cond == ISD::SETNE && isOneConstant(RHS) &&
       LHS->getOpcode() == AArch64ISD::CSEL &&
       isNullConstant(LHS->getOperand(0)) && isOneConstant(LHS->getOperand(1)) &&
       LHS->hasOneUse()) {
     SDLoc DL(N);
 
     // Invert CSEL's condition.
case ISD::BITCAST:
     return;
   case ISD::VECREDUCE_ADD:
   case ISD::VECREDUCE_SMAX:
   case ISD::VECREDUCE_SMIN:
   case ISD::VECREDUCE_UMAX:
   case ISD::VECREDUCE_UMIN:
     Results.push_back(LowerVECREDUCE(SDValue(N, 0), DAG));
     return;
-
   case ISD::CTPOP:
     if (SDValue Result = LowerCTPOP(SDValue(N, 0), DAG))
       Results.push_back(Result);
     return;
   case AArch64ISD::SADDV:
     ReplaceReductionResults(N, Results, DAG, ISD::ADD, AArch64ISD::SADDV);
     return;
   case AArch64ISD::UADDV:

llvm/test/CodeGen/AArch64/vselect-ext.ll

entry:
   %cmp = icmp sgt <16 x i8> %a, <i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10>
   %sel = select <16 x i1> %cmp, <16 x i32> %ext, <16 x i32> zeroinitializer
   ret <16 x i32> %sel
 }
 
 define <16 x i32> @same_zext_used_in_cmp_unsigned_pred_and_select(<16 x i8> %a) {
 ; CHECK-LABEL: same_zext_used_in_cmp_unsigned_pred_and_select:
 ; CHECK:       ; %bb.0: ; %entry
-; CHECK-NEXT:    movi.16b v1, #10
-; CHECK-NEXT:    ushll.8h v2, v0, #0
-; CHECK-NEXT:    ushll2.8h v3, v0, #0
-; CHECK-NEXT:    ushll.4s v4, v2, #0
-; CHECK-NEXT:    cmhi.16b v0, v0, v1
-; CHECK-NEXT:    ushll.4s v5, v3, #0
-; CHECK-NEXT:    ushll2.4s v1, v3, #0
-; CHECK-NEXT:    sshll.8h v3, v0, #0
-; CHECK-NEXT:    sshll2.8h v0, v0, #0
-; CHECK-NEXT:    ushll2.4s v2, v2, #0
-; CHECK-NEXT:    sshll.4s v6, v3, #0
-; CHECK-NEXT:    sshll.4s v7, v0, #0
-; CHECK-NEXT:    sshll2.4s v0, v0, #0
-; CHECK-NEXT:    sshll2.4s v16, v3, #0
-; CHECK-NEXT:    and.16b v3, v1, v0
-; CHECK-NEXT:    and.16b v1, v2, v16
-; CHECK-NEXT:    and.16b v2, v5, v7
-; CHECK-NEXT:    and.16b v0, v4, v6
+; CHECK-NEXT:    movi.4s v1, #10
+; CHECK-NEXT:    ushll2.8h v2, v0, #0
+; CHECK-NEXT:    ushll.8h v0, v0, #0
+; CHECK-NEXT:    ushll2.4s v3, v2, #0
+; CHECK-NEXT:    ushll2.4s v4, v0, #0
+; CHECK-NEXT:    ushll.4s v0, v0, #0
+; CHECK-NEXT:    ushll.4s v2, v2, #0
+; CHECK-NEXT:    cmhi.4s v5, v0, v1
+; CHECK-NEXT:    cmhi.4s v6, v2, v1
+; CHECK-NEXT:    cmhi.4s v7, v3, v1
+; CHECK-NEXT:    cmhi.4s v1, v4, v1
+; CHECK-NEXT:    and.16b v3, v3, v7
+; CHECK-NEXT:    and.16b v1, v4, v1
+; CHECK-NEXT:    and.16b v2, v2, v6
+; CHECK-NEXT:    and.16b v0, v0, v5
 ; CHECK-NEXT:    ret
 entry:
   %ext = zext <16 x i8> %a to <16 x i32>
   %cmp = icmp ugt <16 x i8> %a, <i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10>
   %sel = select <16 x i1> %cmp, <16 x i32> %ext, <16 x i32> zeroinitializer
   ret <16 x i32> %sel
 }
 
 define <16 x i32> @same_sext_used_in_cmp_signed_pred_and_select(<16 x i8> %a) {
 ; CHECK-LABEL: same_sext_used_in_cmp_signed_pred_and_select:
 ; CHECK:       ; %bb.0: ; %entry
-; CHECK-NEXT:    movi.16b v1, #10
-; CHECK-NEXT:    sshll.8h v3, v0, #0
+; CHECK-NEXT:    movi.4s v1, #10
 ; CHECK-NEXT:    sshll2.8h v2, v0, #0
-; CHECK-NEXT:    cmgt.16b v0, v0, v1
-; CHECK-NEXT:    ext.16b v1, v3, v3, #8
-; CHECK-NEXT:    sshll.8h v5, v0, #0
-; CHECK-NEXT:    sshll2.8h v0, v0, #0
-; CHECK-NEXT:    ext.16b v4, v2, v2, #8
-; CHECK-NEXT:    ext.16b v6, v5, v5, #8
-; CHECK-NEXT:    ext.16b v7, v0, v0, #8
-; CHECK-NEXT:    and.8b v0, v2, v0
-; CHECK-NEXT:    sshll.4s v2, v0, #0
-; CHECK-NEXT:    and.8b v0, v3, v5
-; CHECK-NEXT:    and.8b v1, v1, v6
-; CHECK-NEXT:    and.8b v3, v4, v7
+; CHECK-NEXT:    sshll.8h v0, v0, #0
+; CHECK-NEXT:    sshll2.4s v3, v2, #0
+; CHECK-NEXT:    sshll2.4s v4, v0, #0
 ; CHECK-NEXT:    sshll.4s v0, v0, #0
-; CHECK-NEXT:    sshll.4s v1, v1, #0
-; CHECK-NEXT:    sshll.4s v3, v3, #0
+; CHECK-NEXT:    sshll.4s v2, v2, #0
+; CHECK-NEXT:    cmgt.4s v5, v0, v1
+; CHECK-NEXT:    cmgt.4s v6, v2, v1
+; CHECK-NEXT:    cmgt.4s v7, v3, v1
+; CHECK-NEXT:    cmgt.4s v1, v4, v1
+; CHECK-NEXT:    and.16b v3, v3, v7
+; CHECK-NEXT:    and.16b v1, v4, v1
+; CHECK-NEXT:    and.16b v2, v2, v6
+; CHECK-NEXT:    and.16b v0, v0, v5
 ; CHECK-NEXT:    ret
 entry:
   %ext = sext <16 x i8> %a to <16 x i32>
   %cmp = icmp sgt <16 x i8> %a, <i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10, i8 10>
   %sel = select <16 x i1> %cmp, <16 x i32> %ext, <16 x i32> zeroinitializer
   ret <16 x i32> %sel
 }