[AArch64][SVE] Combine add and index_vector

[junparser]

Summary

This patch tries to combine pattern add(index_vector(zero, step), dup(X)) into index_vector(X, step)

TestPlan: check-llvm

[junparser]

fix typo.

[junparser]

This comment was removed by junparser.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

Hi @junparser, there are a few things to take into account here:

• Folding the add into the index_vector may be more expensive in practice, if index_vector(zero, step) is created once, then subsequent adds may be cheaper than having multiple index_vector instructions which each need to calculate the expansion of the (very similar) series. It's probably best to check that the node (add) has only a single result.
• The form where it uses an add may come in useful when selecting an addressing mode for the gather/scatter operations, in which case: add(nxv4i32 dup(X), nxv4i32 index_vector(zero, step)) can be implemented with a scalar+vector addressing mode such as [x0, z0.s]. You may want to add a check that the result is not used as the address in a load/store operation.

This can probably be implemented quite easily in TableGen with some patterns that try to do the fold, but where the PatFrag itself has a condition that it only has a single use that this use is not the address in a MemSDNode. See e.g. AArch64mul_p_oneuse where it also checks for a single use before selecting an MLA/MLS.

Because of how we have currently implemented the legalization of gathers/scatters, I don't think you can currently test the condition that the node is not used in a MemSDNode, but I'm not sure if that matters much.

[junparser]

What I want know something more is what is the boundary between tablegen pattern match and dag combine? I never figure this out. Just use this case as example, we can implement the feature in both place, but I don't know how to handle commutivity in tablegen (maybe SDNPCommutative? I don't know.) , and their effects on the next step: combine load/store with index_vector.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

Hi @sdesmalen It looks reasonable to me to check add has only one use. as for addressing mode for gather/scatter, have we already implemented it now? I haven't see it in trunk. Since my next plan is to combine load/store with index_vector.

13437

Hi @junparser, there are a few things to take into account here:

• Folding the add into the index_vector may be more expensive in practice, if index_vector(zero, step) is created once, then subsequent adds may be cheaper than having multiple index_vector instructions which each need to calculate the expansion of the (very similar) series. It's probably best to check that the node (add) has only a single result.
• The form where it uses an add may come in useful when selecting an addressing mode for the gather/scatter operations, in which case: add(nxv4i32 dup(X), nxv4i32 index_vector(zero, step)) can be implemented with a scalar+vector addressing mode such as [x0, z0.s]. You may want to add a check that the result is not used as the address in a load/store operation.

This can probably be implemented quite easily in TableGen with some patterns that try to do the fold, but where the PatFrag itself has a condition that it only has a single use that this use is not the address in a MemSDNode. See e.g. AArch64mul_p_oneuse where it also checks for a single use before selecting an MLA/MLS.

Because of how we have currently implemented the legalization of gathers/scatters, I don't think you can currently test the condition that the node is not used in a MemSDNode, but I'm not sure if that matters much.

[sdesmalen]

In D100107#2682168, @junparser wrote:

What I want know something more is what is the boundary between tablegen pattern match and dag combine? I never figure this out. Just use this case as example, we can implement the feature in both place, but I don't know how to handle commutivity in tablegen (maybe SDNPCommutative? I don't know.) , and their effects on the next step: combine load/store with index_vector.

There isn't a firm rule to follow re tablegen pattern vs manual dag combine, but I think the steer is that when the combine is needed after legalization and it's possible to specify as a pattern, a pattern is preferred. In this case, probably a pattern would be a clean way to implement it. I haven't used SDNPCommutative before, but that looks like it should do the trick.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

The gather/scatter nodes are not implemented as patterns at the moment, only as custom dagcombine code. So adding a check that it isn't used in an addressing mode would currently have no effect, but it probably would if we change our gather/scatter implementation to not use the intrinsics, but use common nodes instead. This is something we want to do in the not-too-distant future. You can alternatively also leave a FIXME.

Since my next plan is to combine load/store with index_vector

Can you elaborate what you mean by that?

[junparser]

In D100107#2685779, @sdesmalen wrote:

In D100107#2682168, @junparser wrote:

What I want know something more is what is the boundary between tablegen pattern match and dag combine? I never figure this out. Just use this case as example, we can implement the feature in both place, but I don't know how to handle commutivity in tablegen (maybe SDNPCommutative? I don't know.) , and their effects on the next step: combine load/store with index_vector.

There isn't a firm rule to follow re tablegen pattern vs manual dag combine, but I think the steer is that when the combine is needed after legalization and it's possible to specify as a pattern, a pattern is preferred. In this case, probably a pattern would be a clean way to implement it. I haven't used SDNPCommutative before, but that looks like it should do the trick.

Thanks for explain this.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

We already have node like GLD1_MERGE_ZERO and combine function like performGLD1Combine, so rather than implementing it as tablegen pattern, i prefer to do it in here, and then do some more work in performGLD1Combine with index_vector, that the plan.

Also what is the difference between scalar+vector and. vector+ index addressing mode? since we can also use vector+ index addressing mode after this combination.

[junparser]

Address comment.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

Hi @junparser, please can you provide an example for what you're ultimately trying to achieve.

As Sander points out the decision as to where things should live is not not clear cut and I've nothing against doing things within DAGCombine if it helps reduce isel pattern explosion. That said I am very keen to cut down the need for duplicated combines due to the lack of a canonicalised DAG. With regards to the addressing modes we should already have good support by looking for ADD based arithmetic used to compute addresses. Unless there's a good reason I'd rather not have to duplicate this logic to handle INDEX_VECTOR.

Personally I see INDEX_VECTOR as a last stage optimisation whereby everything else has had a chance to remove the ADD (or MUL if we're talking about a stride) and thus the last option is to merge stray ADDs into the INDEX_VECTOR.

For what it's worth the target specific INDEX_VECTOR existed before the common STEP_VECTOR nodes so I think there's a good chance we'll put effort into ensuring any early uses of INDEX_VECTOR are first converted to STEP_VECTOR to maximum the effectiveness of future common STEP_VECTOR combines.

[junparser]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13437

Hi @junparser, please can you provide an example for what you're ultimately trying to achieve.

combine add+index_vector here, and do some address mode choice in selectGatherScatterAddrMode between scalar+vector and. vector+ index , and convert it back to index_vector + gather.

As Sander points out the decision as to where things should live is not not clear cut and I've nothing against doing things within DAGCombine if it helps reduce isel pattern explosion. That said I am very keen to cut down the need for duplicated combines due to the lack of a canonicalised DAG. With regards to the addressing modes we should already have good support by looking for ADD based arithmetic used to compute addresses. Unless there's a good reason I'd rather not have to duplicate this logic to handle INDEX_VECTOR.

Personally I see INDEX_VECTOR as a last stage optimisation whereby everything else has had a chance to remove the ADD (or MUL if we're talking about a stride) and thus the last option is to merge stray ADDs into the INDEX_VECTOR.

sounds reasonable to me.

For what it's worth the target specific INDEX_VECTOR existed before the common STEP_VECTOR nodes so I think there's a good chance we'll put effort into ensuring any early uses of INDEX_VECTOR are first converted to STEP_VECTOR to maximum the effectiveness of future common STEP_VECTOR combines.

I see we lower step_vector to index_vector currently, do you mean that we only generate index_vector in instruction selection, and only use step_vector before that? Hmm, maybe it's a good idea.

[junparser]

Address comments. using tablegen pattern match.

[sdesmalen]

Thanks, that looks quite neat. Can you also add a few tests when there is >1 use of the stepvector (e.g. using the stepvector in two adds), so we can test the fold indeed doesn't happen?
One test for each of the instructions should be sufficient.

llvm/test/CodeGen/AArch64/sve-stepvector.ll

145

nit: add_stepvector_nxv8i8_2_commutative ?

[junparser]

Address comments.

[junparser]

Update testcase.

[junparser]

@sdesmalen kindly ping~

[sdesmalen]

Other than my comment on two missing tests, the patch looks good to me.

llvm/test/CodeGen/AArch64/sve-stepvector.ll

132

nit: I think you're still missing a test for the scalar+scalar case and one for imm+scalar (for start/stride respectively)

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (!LHS.getNode())
       return SDValue();
 
     LHS = DAG.getNode(ExtType, SDLoc(N), VT, LHS);
   }
 
   return DAG.getNode(N->getOpcode(), SDLoc(N), VT, LHS, RHS);
 }
 
+// add(index_vector(zero, step), dup(X)) -> index_vector(X, step)
+static SDValue performAddIndexVectorCombine(SDNode *N, SelectionDAG &DAG) {
+  if (N->getOpcode() != ISD::ADD)
+    return SDValue();
+
+  SDValue IV = N->getOperand(0);
+  SDValue A = N->getOperand(1);
+  // Handle commutivity
+  if (IV.getOpcode() != AArch64ISD::INDEX_VECTOR)
+    std::swap(IV, A);
+  if (IV.getOpcode() != AArch64ISD::INDEX_VECTOR)
+    return SDValue();
+
+  EVT VT = N->getValueType(0);
+  assert(VT.isScalableVector() &&
+         "Only expect scalable vectors for STEP_VECTOR");
+
+  auto *BaseVal = dyn_cast<ConstantSDNode>(IV->getOperand(0));
+  if (BaseVal && BaseVal->isNullValue() &&
+      (A.getOpcode() == AArch64ISD::DUP || A.getOpcode() == ISD::SPLAT_VECTOR))
+    return DAG.getNode(AArch64ISD::INDEX_VECTOR, SDLoc(N), VT, A->getOperand(0),
+                       IV->getOperand(1));
+
+  return SDValue();
+}
+
 static SDValue performAddSubCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI,
                                     SelectionDAG &DAG) {
   // Try to change sum of two reductions.
   if (SDValue Val = performUADDVCombine(N, DAG))
     return Val;
   if (SDValue Val = performAddDotCombine(N, DAG))
     return Val;
 
+  // Try to combine add with index_vector
+  if (SDValue Val = performAddIndexVectorCombine(N, DAG))
+    return Val;
+
   return performAddSubLongCombine(N, DCI, DAG);
 }
 
 // Massage DAGs which we can use the high-half "long" operations on into
 // something isel will recognize better. E.g.
 //
 // (aarch64_neon_umull (extract_high vec) (dupv64 scalar)) -->
 //   (aarch64_neon_umull (extract_high (v2i64 vec)))

llvm/test/CodeGen/AArch64/sve-stepvector.ll

 }
 
 ; ILLEGAL INTEGER TYPES
 
 define <vscale x 4 x i64> @stepvector_nxv4i64() {
 ; CHECK-LABEL: stepvector_nxv4i64:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    mov z1.d, x8
+; CHECK-NEXT:    index z1.d, x8, #1
 ; CHECK-NEXT:    index z0.d, #0, #1
-; CHECK-NEXT:    add z1.d, z0.d, z1.d
 ; CHECK-NEXT:    ret
 entry:
   %0 = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
   ret <vscale x 4 x i64> %0
 }
 
 define <vscale x 16 x i32> @stepvector_nxv16i32() {
 ; CHECK-LABEL: stepvector_nxv16i32:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    cntw x9
 ; CHECK-NEXT:    cnth x8
+; CHECK-NEXT:    cntw x9
+; CHECK-NEXT:    mov z0.s, w8
+; CHECK-NEXT:    index z1.s, w9, #1
+; CHECK-NEXT:    index z2.s, w8, #1
+; CHECK-NEXT:    add z3.s, z1.s, z0.s
 ; CHECK-NEXT:    index z0.s, #0, #1
-; CHECK-NEXT:    mov z1.s, w9
-; CHECK-NEXT:    mov z3.s, w8
-; CHECK-NEXT:    add z1.s, z0.s, z1.s
-; CHECK-NEXT:    add z2.s, z0.s, z3.s
-; CHECK-NEXT:    add z3.s, z1.s, z3.s
 ; CHECK-NEXT:    ret
 entry:
   %0 = call <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
   ret <vscale x 16 x i32> %0
 }
 
 define <vscale x 2 x i32> @stepvector_nxv2i32() {
 ; CHECK-LABEL: stepvector_nxv2i32:

 entry:
   %0 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
   %1 = add <vscale x 8 x i8> %p, %0
   %2 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
   %3 = add <vscale x 8 x i8> %1, %2
   ret <vscale x 8 x i8> %3
 }
 
 define <vscale x 8 x i8> @mul_stepvector_nxv8i8() {
 ; CHECK-LABEL: mul_stepvector_nxv8i8:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    index z0.h, #0, #2
 ; CHECK-NEXT:    ret
 entry:
   %0 = insertelement <vscale x 8 x i8> poison, i8 2, i32 0
   %1 = shufflevector <vscale x 8 x i8> %0, <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer
   %2 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
   %3 = mul <vscale x 8 x i8> %2, %1
   ret <vscale x 8 x i8> %3
 }
 
 define <vscale x 8 x i8> @shl_stepvector_nxv8i8() {
 ; CHECK-LABEL: shl_stepvector_nxv8i8:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    index z0.h, #0, #4
 ; CHECK-NEXT:    ret
 entry:
   %0 = insertelement <vscale x 8 x i8> poison, i8 2, i32 0
   %1 = shufflevector <vscale x 8 x i8> %0, <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer
   %2 = call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()