[AArch64][SVE] Avoid AND operation if both side are splat of i1 or PTRUE

[dtemirbulatov]

Summary

If both sides of AND operations are i1 splat_vectors or PTRUE node then we can produce just i1 splat_vector as the result.

[dtemirbulatov]

Changed proposed test in order to remove -O3 out of llc parameters.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16380

I think you can use the existing function isAllActivePredicate to check if either of the operands is an all active-predicate. This function already looks through reinterpret nodes and already checks both the splat/ptrue case.

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

16

It would also be good to have a test that uses a shufflevector to create the splat.

20

Is this function call relevant?

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16275

Rather than use a recursive function call you could just do:

static SDValue findRootNonReinterpretNode(SDValue Op) {
  while (Op.getOpcode() == AArch64ISD::REINTERPRET_CAST)
    Op = Op->getOperand(0);
  return Op;
}

Perhaps that's a bit friendlier to the compiler too?

[dtemirbulatov]

Addressed comments.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16380

but I need to check and return ptrue/splat node and isAllActivePredicate() just checks.

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16380

Hi @dtemirbulatov, I'm re-reading this code and isAllActivePredicate. I actually think @sdesmalen is right and we should use isAllActivePredicate here, because your current code doesn't take into account various cases such as:

1. When reinterpreting from a type with fewer elements the "new" elements are not active, so I think we should bail out of this optimisation in this case.
2. Your code below doesn't consider what type of pattern is being passed to ptrue, but we should be. isAllActivePredicate only returns true for AArch64SVEPredPattern::all.

For any cases where we bail out I think it's worth adding a negative test as well.

[dtemirbulatov]

Addressed coments.

[david-arm]

Hi @dtemirbulatov, this is looking better now thanks! I just have a couple more comments ...

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16378

I think this comment suggests we only support the case where both operands are i1 splat_vectors, but your code below supports only one operand being an all-active predicate. Can you update the comment to reflect that? For example, something like

// If either side of the AND operation is a i1 splat_vector then return the other
// operand.

For example, and %splat_i1, %b -> %b and and %a, %splat_i1 -> %a. Can you add twmo test cases to cover both of these cases too?

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

36

This looks like a negative test. Can you add a comment explaining why the DAG combine doesn't happen?

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16275–16281

isAllInactivePredicate is unused now?

16283

Can you commit the moving of this function as a non-functional change (as a precursory patch)?

16380–16383

nit: can you replace Src with N->getOperand(0) here? It's equivalent, but it makes it a bit more intuitive to understand.

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

41

I would have expected all ptrues to have folded away and to only have:

ptrue p0.s

Why are these AND instructions not removed?

46

nit: you can remove tail for these function calls.

[sdesmalen]

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

41

In second consideration, I see that only one of these AND's can be removed. The semantics of llvm.aarch64.sve.convert.to.svbool.nxv2i1 are such that it explicitly zeroes the inactive lanes, i.e. they don't map 1-1 to REINTERPRET_CAST.

[dtemirbulatov]

Added combine for AArch64ISD::REINTERPRET_CAST nodes, addressed remarks.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16275–16281

No, still in use in performVSelectCombine.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16324–16326

This code would simplify:

t0: nxv2i1 ...
  t1: nxv16i1 REINTERPRET_CAST t0
    t2: nxv4i1 REINTERPRET_CAST t1
      t3: nxv8i1 REINTERPRET_CAST t2
        t4: nxv2i1 REINTERPRET_CAST t3

into: t0

But not something like this

t0: nxv2i1 ...
  t1: nxv16i1 REINTERPRET_CAST t0
    t2: nxv4i1 REINTERPRET_CAST t1
      t3: nxv8i1 REINTERPRET_CAST t2
        t4: nxv16i1 REINTERPRET_CAST t3

into:

t0: nxv2i1 ...
  t1: nxv16i1 REINTERPRET_CAST t0

I'm not sure if this is common, but it's probably more robust to do something like this:

while(Op.getOpcode() == AArch64ISD::REINTERPRET_CAST &&
      LeafOp.getValueType() != Op.getValueType())
  Op = Op->getOperand(0);

llvm/test/CodeGen/AArch64/sve-intrinsics-reinterpret.ll

128

Nice :)

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

8

Can you give these tests more meaningful names, e.g. @fold_away_ptrue_and_ptrue for @foo and @fold_away_ptrue_and_splat_predicate for bar.

[david-arm]

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

1

I think this is still missing tests for the more general case where only one AND operand is an all-active predicate. Your logic on line 16380 permits this more general behaviour so we ought to test it.

[dtemirbulatov]

Addressed comments.

[sdesmalen]

LGTM with comment on the test addressed!

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

37

Can you give @foo and @bar some more meaningful names as well?

[dtemirbulatov]

Renamed two functions in sve-splat-one-and-ptrue.ll

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static bool isConstantSplatVectorMaskForType(SDNode *N, EVT MemVT) {
 
   if (N->getOpcode() == AArch64ISD::DUP || N->getOpcode() == ISD::SPLAT_VECTOR)
     if (auto *Op0 = dyn_cast<ConstantSDNode>(N->getOperand(0)))
       return Op0->getAPIntValue().getLimitedValue() == MaskForTy;
 
   return false;
 }
 
 static bool isAllInactivePredicate(SDValue N) {
   // Look through cast.
   while (N.getOpcode() == AArch64ISD::REINTERPRET_CAST)
     N = N.getOperand(0);
 
   return ISD::isConstantSplatVectorAllZeros(N.getNode());
 }
 
 static bool isAllActivePredicate(SelectionDAG &DAG, SDValue N) {
   unsigned NumElts = N.getValueType().getVectorMinNumElements();
 
   // Look through cast.
   while (N.getOpcode() == AArch64ISD::REINTERPRET_CAST) {
     N = N.getOperand(0);
     // When reinterpreting from a type with fewer elements the "new" elements
     // are not active, so bail if they're likely to be used.
     if (N.getValueType().getVectorMinNumElements() < NumElts)
if (MaxSVESize && MinSVESize == MaxSVESize) {
           getNumElementsFromSVEPredPattern(N.getConstantOperandVal(0));
       return PatNumElts == (NumElts * VScale);
     }
   }
 
   return false;
 }
 
+static SDValue performReinterpretCastCombine(SDNode *N) {
+  SDValue LeafOp = SDValue(N, 0);
+  SDValue Op = N->getOperand(0);
+  while (Op.getOpcode() == AArch64ISD::REINTERPRET_CAST)
+    Op = Op->getOperand(0);
+  if (LeafOp.getValueType() == Op.getValueType())
+    return Op;
+  return SDValue();
+}
+
 static SDValue performSVEAndCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   SelectionDAG &DAG = DCI.DAG;
   SDValue Src = N->getOperand(0);
   unsigned Opc = Src->getOpcode();
Dup = DAG.getNode(ISD::SPLAT_VECTOR, DL, UnpkOp->getValueType(0),
                       DAG.getConstant(Mask.zextOrTrunc(32), DL, MVT::i32));
 
     SDValue And = DAG.getNode(ISD::AND, DL,
                               UnpkOp->getValueType(0), UnpkOp, Dup);
 
     return DAG.getNode(Opc, DL, N->getValueType(0), And);
   }
 
+  // If both sides of AND operations are i1 splat_vectors then
+  // we can produce just i1 splat_vector as the result.
+  if (isAllActivePredicate(DAG, N->getOperand(0)))
+    return N->getOperand(1);
+  if (isAllActivePredicate(DAG, N->getOperand(1)))
+    return N->getOperand(0);
+
   if (!EnableCombineMGatherIntrinsics)
     return SDValue();
 
   SDValue Mask = N->getOperand(1);
 
   if (!Src.hasOneUse())
     return SDValue();
 
case AArch64ISD::SPLICE:
     return performSpliceCombine(N, DAG);
   case AArch64ISD::UUNPKLO:
   case AArch64ISD::UUNPKHI:
     return performUnpackCombine(N, DAG, Subtarget);
   case AArch64ISD::UZP1:
     return performUzpCombine(N, DAG);
   case AArch64ISD::SETCC_MERGE_ZERO:
     return performSetccMergeZeroCombine(N, DCI);
+  case AArch64ISD::REINTERPRET_CAST:
+    return performReinterpretCastCombine(N);
   case AArch64ISD::GLD1_MERGE_ZERO:
   case AArch64ISD::GLD1_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_IMM_MERGE_ZERO:
   case AArch64ISD::GLD1S_MERGE_ZERO:

llvm/test/CodeGen/AArch64/sve-intrinsics-reinterpret.ll

; CHECK-NEXT:    ret
   %2 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %1)
   ret <vscale x 16 x i1> %2
 }
 
 ; The first reinterpret should prevent the second one from being simplified as a nop
 define <vscale x 16 x i1> @chained_reinterpret() {
 ; CHECK-LABEL: chained_reinterpret:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptrue p0.b
-; CHECK-NEXT:    ptrue p1.d
-; CHECK-NEXT:    and p0.b, p0/z, p0.b, p1.b
+; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    ret
   %in = tail call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
   %cast2 = call <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1> %in)
   %out = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1> %cast2)
   ret <vscale x 16 x i1> %out
 }
 
 declare <vscale x 8 x i1> @llvm.aarch64.sve.ptrue.nxv8i1(i32 immarg)

llvm/test/CodeGen/AArch64/sve-splat-one-and-ptrue.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+; Ensure that a no-op 'and' get removed with vector splat of 1 or ptrue with proper constant
+
+define <vscale x 16 x i1> @foo() #0 {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 31)
+  %1 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %0)
+  %2 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %and = and <vscale x 16 x i1> %2, %1
+  ret <vscale x 16 x i1> %and
+}
+
+define <vscale x 16 x i1> @bar() #0 {
+; CHECK-LABEL: bar:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ret
+entry:
+  %ins = insertelement <vscale x 4 x i1> undef, i1 1, i32 0
+  %splat = shufflevector <vscale x 4 x i1> %ins, <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer
+  %0 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %splat)
+  %1 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %and = and <vscale x 16 x i1> %0, %1
+  ret <vscale x 16 x i1> %and
+}
+
+; Ensure that one AND operation remain for inactive lanes zeroing with 2 x i1 type (llvm.aarch64.sve.convert.to.svbool.nxv2i1).
+define <vscale x 16 x i1> @foo1() #0 {
+; CHECK-LABEL: foo1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    and p0.b, p1/z, p1.b, p0.b
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 31)
+  %1 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %0)
+  %2 = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %3 = call <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1> %2)
+  %4 = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1> %3)
+  %and = and <vscale x 16 x i1> %4, %1
+  ret <vscale x 16 x i1> %and
+}
+
+declare <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 immarg)
+
+declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1>)
+
+declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 immarg)
+
+declare <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1>)
+
+declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1>)
+
+
+attributes #0 = { "target-features"="+sve" }