[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

16271

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

15

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

19

Is this function call relevant?

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16217

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

16271

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

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16271

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

16269

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

16217–16223

isAllInactivePredicate is unused now?

16225

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

16271–16274

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

16217–16223

No, still in use in performVSelectCombine.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16219–16221

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

9

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

2

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 SDValue findRootNonReinterpretNode(SDValue Op) {
+  if (Op.getNode()->getOpcode() != AArch64ISD::REINTERPRET_CAST)
+    return Op;
+  return findRootNonReinterpretNode(Op->getOperand(0));
+}
+
 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();
 
   // Zero/any extend of an unsigned unpack
   if (Opc == AArch64ISD::UUNPKHI || Opc == AArch64ISD::UUNPKLO) {
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
+  SDValue Op0 = findRootNonReinterpretNode(Src);
+  SDValue Op1 = findRootNonReinterpretNode(N->getOperand(1));
+  if ((ISD::isConstantSplatVectorAllOnes(Op0.getNode()) ||
+       Op0.getNode()->getOpcode() == AArch64ISD::PTRUE) &&
+      (ISD::isConstantSplatVectorAllOnes(Op1.getNode()) ||
+       Op1.getNode()->getOpcode() == AArch64ISD::PTRUE)) {
+    EVT OpVT0 = Op0.getValueType();
+    EVT OpVT1 = Op1.getValueType();
+    if (OpVT0.getVectorElementType() == MVT::i1 &&
+        OpVT1.getVectorElementType() == MVT::i1) {
+      SDLoc DL(N);
+      return (OpVT0.getVectorElementCount().getKnownMinValue() >
+              OpVT1.getVectorElementCount().getKnownMinValue())
+                 ? N->getOperand(1)
+                 : Src;
+    }
+  }
+
   if (!EnableCombineMGatherIntrinsics)
     return SDValue();
 
   SDValue Mask = N->getOperand(1);
 
   if (!Src.hasOneUse())
     return SDValue();
 

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
+
+; Function Attrs: nounwind uwtable vscale_range(1,16)
+define void @bar() #0 {
+; CHECK-LABEL: bar:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    b foo
+entry:
+  %0 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 31)
+  %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %0)
+  %2 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
+  %and = and <vscale x 16 x i1> %2, %1
+  tail call void @foo(<vscale x 16 x i1> %and) #3
+  ret void
+}
+
+declare void @foo(<vscale x 16 x i1>)
+
+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)
+
+
+attributes #0 = { "target-features"="+sve" }