[SVE] Add a DAG combiner fold to visitADD for vscale with truncate

[Allen]

Summary

Fold a+truncate(vscale(c1))+truncate(vscale(c2)) to a+truncate(vscale(c1+c2))
As the vscale Constant is legalized with type i64 DAG, so it inserts a truncate
in function LowerVSCALE, which cause the mismatch in D82792.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2622

You can't assert that VSCALE returns an i64. Probably not even on Arm, but definitely not in target-independent code.

2631

Do you really need to explicitly use the number "64" to make this work?

[reames]

I don't claim to fully understand this, so my comment here might be off base.

I suspect your fold can be generalized as: fold a+truncate(vscale(c1))+truncate(vscale(c2)) to a+truncate(vscale(c1)+vscale(c2))

The vscale(c1)+vscale(c2) to vscale(C1 + C2) is handled separately above already.

If this is true, that your transform reduces to proving that it's legal to common the truncate. However, as a far as I known trunc(x) + trunc(y) is always equal to trunc(x+y). So why do we need this transform at all? Shouldn't this be covered by generic trunc folds and the existing rule?

Anyways, I'm clearly missing something here. Any idea what?

[Allen]

In D126532#3544144, @reames wrote:

I don't claim to fully understand this, so my comment here might be off base.

I suspect your fold can be generalized as: fold a+truncate(vscale(c1))+truncate(vscale(c2)) to a+truncate(vscale(c1)+vscale(c2))

The vscale(c1)+vscale(c2) to vscale(C1 + C2) is handled separately above already.

If this is true, that your transform reduces to proving that it's legal to common the truncate. However, as a far as I known trunc(x) + trunc(y) is always equal to trunc(x+y). So why do we need this transform at all? Shouldn't this be covered by generic trunc folds and the existing rule?

Anyways, I'm clearly missing something here. Any idea what?

Thanks for your attention. With my debug, I find vscale is a little especial. As the vscale(C1) is not a const type node before LowerVSCALE, so we missing the transform of trunc(x) + trunc(y) --> trunc(x+y) with instcombine.
And in the dagcombine, it match the DAG IR from bottom to top, so we can capture the former a+truncate(vscale(c1))+truncate(vscale(c2)), but not truncate(vscale(c1))+truncate(vscale(c2)).

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2622

Thanks, apply your commit

2631

Yes, it can be deleted

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2635

Do you need to check that both vscales have the same types?

[Allen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2635

Thanks for your attention.
As only the same type node can be on the both side of ISD::ADD, so I think it can be ignored. Am I missing something?

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2635

There's no guarantee the input types of the two truncates is the same. You could have something like (i16 trunc (i32 vscale)) and (i16 trunc (i64 vscale)).

[Allen]

Add new condtion to check that both vscales have the same types

[Allen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2635

Thanks for detail example. Apply your comment

[Allen]

ping ?

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2629–2630

Would the following combine be safe?

ty1 truncate(ty2 vscale(c1)) -> ty1 vscale(c1)

I ask because then the combine just above the new one would just work?

I guess the problem is that operation legalisation might be the thing introducing the truncate but then we can just limit the combine to before then. I wouldn't expect the combine to be all that useful after legalisation anyway, although am happy to be proven wrong if you've a test case.

2631

Is the VT.isScalarInteger() check necessary? I figure the later ISD::VSCALE requirement will guaranteed such, plus I don't see anything in the if block that actually cares.

[Allen]

Add the check Level > AfterLegalizeVectorOps

[Allen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2618

With Further debuging , I found this pattern doen't match during llvm::BeforeLegalizeTypes because when the node t26 is updated, it only add its user node t9 to the Worklist, without its recursive user node t10
https://github.com/llvm/llvm-project/blob/main/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L1607

Surely, It's unreasonable to recursively add nodes t10 to the worklist, as it may lead to a significant increase in compile time.

SelectionDAG has 16 nodes:
  t0: ch = EntryToken
        t2: i32,ch = CopyFromReg t0, Register:i32 %0
        t26: i32 = vscale Constant:i32<8>  --- update last
      t9: i32 = add nuw t2, t26
      t19: i32 = vscale Constant:i32<2>
    t10: i32 = add nuw t9, t19
  t12: ch,glue = CopyToReg t0, Register:i32 $w0, t10
    t4: i32 = vscale Constant:i32<1>
  t24: i32 = shl t4, Constant:i64<3>
  t13: ch = AArch64ISD::RET_FLAG t12, Register:i32 $w0, t12:1

2629–2630

thanks @paulwalker-arm for your idea, I thinks ty1 truncate(ty2 vscale(c1)) -> ty1 vscale(c1) is safe as C1 is a const.
But If we combine this, the ISel may crash as we don't defined such pattern.

Yes, the truncate is introduced in AArch64TargetLowering::LowerVSCALE, where addressing the legalisation for the new added case combine_add_vscale_C_i32.

BTW: Added the check Level > AfterLegalizeVectorOps according the comment

2631

Thanks, apply your comment, and deleted the unnecessary check.

[Allen]

adopt as I don't work on it now

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitADD(SDNode *N) {
 
   // Fold (add (vscale * C0), (vscale * C1)) to (vscale * (C0 + C1)).
   if (N0.getOpcode() == ISD::VSCALE && N1.getOpcode() == ISD::VSCALE) {
     const APInt &C0 = N0->getConstantOperandAPInt(0);
     const APInt &C1 = N1->getConstantOperandAPInt(0);
     return DAG.getVScale(DL, VT, C0 + C1);
   }
 
   // fold a+vscale(c1)+vscale(c2) -> a+vscale(c1+c2)
   if ((N0.getOpcode() == ISD::ADD) &&
       (N0.getOperand(1).getOpcode() == ISD::VSCALE) &&
       (N1.getOpcode() == ISD::VSCALE)) {
+    assert(VT == MVT::i64 && "Unexpected element type!");
     const APInt &VS0 = N0.getOperand(1)->getConstantOperandAPInt(0);
     const APInt &VS1 = N1->getConstantOperandAPInt(0);
     SDValue VS = DAG.getVScale(DL, VT, VS0 + VS1);
     return DAG.getNode(ISD::ADD, DL, VT, N0.getOperand(0), VS);
   }
 
+  // fold a+truncate(vscale(c1))+truncate(vscale(c2))
+  // to a+truncate(vscale(c1+c2))
+  if (VT.isScalarInteger() && VT.getSizeInBits() < 64 &&
+      (N0.getOpcode() == ISD::ADD) &&
+      (N0.getOperand(1).getOpcode() == ISD::TRUNCATE) &&
+      (N0.getOperand(1).getOperand(0).getOpcode() == ISD::VSCALE) &&
+      (N1.getOpcode() == ISD::TRUNCATE) &&
+      (N1.getOperand(0).getOpcode() == ISD::VSCALE)) {
+    const APInt &VS0 =
+        N0.getOperand(1).getOperand(0)->getConstantOperandAPInt(0);
+    const APInt &VS1 = N1.getOperand(0)->getConstantOperandAPInt(0);
+    SDValue VS =
+        DAG.getZExtOrTrunc(DAG.getVScale(DL, MVT::i64, VS0 + VS1), DL, VT);
+    return DAG.getNode(ISD::ADD, DL, VT, N0.getOperand(0), VS);
+  }
+
   // Fold (add step_vector(c1), step_vector(c2)  to step_vector(c1+c2))
   if (N0.getOpcode() == ISD::STEP_VECTOR &&
       N1.getOpcode() == ISD::STEP_VECTOR) {
     const APInt &C0 = N0->getConstantOperandAPInt(0);
     const APInt &C1 = N1->getConstantOperandAPInt(0);
     APInt NewStep = C0 + C1;
     return DAG.getStepVector(DL, VT, NewStep);
   }

llvm/test/CodeGen/AArch64/sve-vscale-combine.ll

 ; CHECK-LABEL: combine_shl_vscale_i32:
 ; CHECK-NOT:   shl
 ; CHECK-NEXT:  rdvl  x0, #1
 ; CHECK-NEXT:  ret
  %vscale = call i32 @llvm.vscale.i32()
  %shl = shl i32 %vscale, 4
  ret i32 %shl
 }
+
+; Fold a+truncate(vscale(c1))+truncate(vscale(c2)) to a+truncate(vscale(c1+c2))
+define i32 @combine_add_vscale_C_i32(i32 %index) nounwind {
+; CHECK-LABEL: combine_add_vscale_C_i32:
+; CHECK-NEXT:    cntd x8, all, mul #5
+; CHECK-NEXT:    add w0, w0, w8
+; CHECK-NEXT:    ret
+  %vscale = call i32 @llvm.vscale.i32()
+  %mul8 = mul i32 %vscale, 8
+  %mul2 = mul i32 %vscale, 2
+  %index.next = add nuw i32 %index, %mul8
+  %add = add nuw i32 %index.next, %mul2
+  ret i32 %add
+}