[SelectionDAG] Enable WidenVecOp_VECREDUCE_SEQ for scalable vector

[Jimerlife]

Summary

Enable WidenVecOp_VECREDUCE_SEQ for scalable vector by inserting some splatted neutral vectors.
e.g.
when widen <vscale x 10 x half> to <vscale x 16 x half> in index (10, 12, 14) inserted nxv2f16:splat<neutral element>

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6153

Why was this blank line removed?

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6152

I may be missing something here, but I'm not sure I understand the need for a GCD or for a loop that generates WideVT.getVectorMinNumElements() INSERT_SUBVECTORs. Should this code not just insert VecOp into a splat(NeutralElem) of type WideVT using INSERT_SUBVECTOR?

i.e.

vecreduce_seq(<vscale x 10 x half> op)
    <=>
vecreduce_seq(
    insert_subvector(op, <vscale x 16 x half> splat(%neutral.elem), /*idx=*/0))

where <vscale x 10 x half> is OrigVT and <vscale x 16 x half> is WideVT.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6152

I think you meant insert_subvector(<vscale x 16 x half> splat(%neutral.elem), op, /*idx=*/0). op would be the subvector so should be the second argument.

That would need to go through WidenVecOp_INSERT_SUBVECTOR, but it will fail because InVec isn't undef.

[Jimerlife]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6152

If we do like this vecreduce_seq(insert_subvector(<vscale x 16 x half> splat, vecop, /*idx=*/0)), that would go through WidenVecOp_INSERT_SUBVECTOR because of vecop need to widen.
But now, WidenVecOp_INSERT_SUBVECTOR only support insert subvector into undefined vector.

[Jimerlife]

address comment

[Jimerlife]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6153

Mistakenly deleted blank line. Restored back.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6152

Thanks, that makes sense.

6153

For SVE there is still the problem that the type action of SplatVT may be TypeWiden as well, because we haven't implemented all support for nxv1 types yet. The approach taken here seems valid though, it just means that we can't have the same tests for SVE yet.

llvm/test/CodeGen/RISCV/rvv/vreductions-fp-sdnode.ll

1081

Could you also add @vreduce_ord_fadd_nxv6f16 and @vreduce_ord_fadd_nxv10f16 to llvm/test/CodeGen/AArch64/vecreduce-fadd-legalization-strict.ll as well ? (and then just re-generate the CHECK lines using the update_lcc_test_checks.py script)

[sdesmalen]

llvm/test/CodeGen/RISCV/rvv/vreductions-fp-sdnode.ll

1081

And @vreduce_ord_fadd_nxv12f16 as well.

[Jimerlife]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6153

Should I need to add getTypeAction(SplatVT) != TargetLowering::TypeWidenVector to avoid this problem now?

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

6153

No I don't think that's necessary, but thanks for the suggestion! When we fix the nxv1 legalisation of insert_subvector, we'll get the support added in this patch for free. And the compiler will currently crash with an error either way :) It was more just an observation.

[Jimerlife]

address comment and add test for AArch64 SVE

[Jimerlife]

llvm/test/CodeGen/RISCV/rvv/vreductions-fp-sdnode.ll

1081

This vecreduce-fadd-legalization-strict.ll file is only for NEON feature. So, I add tests in sve-fp-reduce.ll. Is it OK?

[sdesmalen]

LGTM!

llvm/test/CodeGen/RISCV/rvv/vreductions-fp-sdnode.ll

1081

Thanks, I hadn't seen that. That's fine!

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

SDValue DAGTypeLegalizer::WidenVecOp_VECREDUCE_SEQ(SDNode *N) {
   EVT ElemVT = OrigVT.getVectorElementType();
   SDNodeFlags Flags = N->getFlags();
 
   unsigned Opc = N->getOpcode();
   unsigned BaseOpc = ISD::getVecReduceBaseOpcode(Opc);
   SDValue NeutralElem = DAG.getNeutralElement(BaseOpc, dl, ElemVT, Flags);
 
   // Pad the vector with the neutral element.
-  unsigned OrigElts = OrigVT.getVectorNumElements();
-  unsigned WideElts = WideVT.getVectorNumElements();
+  unsigned OrigElts = OrigVT.getVectorMinNumElements();
+  unsigned WideElts = WideVT.getVectorMinNumElements();
+
+  if (WideVT.isScalableVector()) {
+    unsigned GCD = greatestCommonDivisor(OrigElts, WideElts);
+    EVT SplatVT = EVT::getVectorVT(*DAG.getContext(), ElemVT,
+                                   ElementCount::getScalable(GCD));
+    SDValue SplatNeutral = DAG.getSplatVector(SplatVT, dl, NeutralElem);
+    for (unsigned Idx = OrigElts; Idx < WideElts; Idx = Idx + GCD)
+      Op = DAG.getNode(ISD::INSERT_SUBVECTOR, dl, WideVT, Op, SplatNeutral,
+                       DAG.getVectorIdxConstant(Idx, dl));
+    return DAG.getNode(Opc, dl, N->getValueType(0), AccOp, Op, Flags);
+  }
+
   for (unsigned Idx = OrigElts; Idx < WideElts; Idx++)
     Op = DAG.getNode(ISD::INSERT_VECTOR_ELT, dl, WideVT, Op, NeutralElem,
                      DAG.getVectorIdxConstant(Idx, dl));
 
   return DAG.getNode(Opc, dl, N->getValueType(0), AccOp, Op, Flags);
 }
 
 SDValue DAGTypeLegalizer::WidenVecOp_VP_REDUCE(SDNode *N) {
   assert(N->isVPOpcode() && "Expected VP opcode");
 
   SDLoc dl(N);
   SDValue Op = GetWidenedVector(N->getOperand(1));

llvm/test/CodeGen/RISCV/rvv/vreductions-fp-sdnode.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=riscv32 -mattr=+d,+zfh,+experimental-zvfh,+v -target-abi=ilp32d \
+; RUN: llc -mtriple=riscv32 -mattr=+d,+zfh,+experimental-zvfh,+v,+m -target-abi=ilp32d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s
-; RUN: llc -mtriple=riscv64 -mattr=+d,+zfh,+experimental-zvfh,+v -target-abi=lp64d \
+; RUN: llc -mtriple=riscv64 -mattr=+d,+zfh,+experimental-zvfh,+v,+m -target-abi=lp64d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s
 
 declare half @llvm.vector.reduce.fadd.nxv1f16(half, <vscale x 1 x half>)
 
 define half @vreduce_fadd_nxv1f16(<vscale x 1 x half> %v, half %s) {
 ; CHECK-LABEL: vreduce_fadd_nxv1f16:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 1, e16, m1, ta, mu

 ; CHECK-NEXT:    vfmv.s.f v9, fa0
 ; CHECK-NEXT:    vsetvli a0, zero, e32, mf2, ta, mu
 ; CHECK-NEXT:    vfredusum.vs v8, v8, v9
 ; CHECK-NEXT:    vfmv.f.s fa0, v8
 ; CHECK-NEXT:    ret
   %red = call reassoc nsz float @llvm.vector.reduce.fadd.nxv1f32(float %s, <vscale x 1 x float> %v)
   ret float %red
 }
+
+; Test Widen VECREDUCE_SEQ_FADD
+declare half @llvm.vector.reduce.fadd.nxv3f16(half, <vscale x 3 x half>)
+
+define half @vreduce_ord_fadd_nxv3f16(<vscale x 3 x half> %v, half %s) {
+; CHECK-LABEL: vreduce_ord_fadd_nxv3f16:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli a0, a0, 3
+; CHECK-NEXT:    slli a1, a0, 1
+; CHECK-NEXT:    add a1, a1, a0
+; CHECK-NEXT:    add a0, a1, a0
+; CHECK-NEXT:    fmv.h.x ft0, zero
+; CHECK-NEXT:    fneg.h ft0, ft0
+; CHECK-NEXT:    vsetvli a2, zero, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.v.f v9, ft0
+; CHECK-NEXT:    vsetvli zero, a0, e16, m1, tu, mu
+; CHECK-NEXT:    vslideup.vx v8, v9, a1
+; CHECK-NEXT:    vsetivli zero, 1, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.s.f v9, fa0
+; CHECK-NEXT:    vsetvli a0, zero, e16, m1, ta, mu
+; CHECK-NEXT:    vfredosum.vs v8, v8, v9
+; CHECK-NEXT:    vfmv.f.s fa0, v8
+; CHECK-NEXT:    ret
+  %red = call half @llvm.vector.reduce.fadd.nxv3f16(half %s, <vscale x 3 x half> %v)
+  ret half %red
+}
+
+declare half @llvm.vector.reduce.fadd.nxv6f16(half, <vscale x 6 x half>)
+
+define half @vreduce_ord_fadd_nxv6f16(<vscale x 6 x half> %v, half %s) {
+; CHECK-LABEL: vreduce_ord_fadd_nxv6f16:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli a0, a0, 2
+; CHECK-NEXT:    add a1, a0, a0
+; CHECK-NEXT:    fmv.h.x ft0, zero
+; CHECK-NEXT:    fneg.h ft0, ft0
+; CHECK-NEXT:    vsetvli a2, zero, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.v.f v10, ft0
+; CHECK-NEXT:    vsetvli zero, a1, e16, m1, tu, mu
+; CHECK-NEXT:    vslideup.vx v9, v10, a0
+; CHECK-NEXT:    vsetivli zero, 1, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.s.f v10, fa0
+; CHECK-NEXT:    vsetvli a0, zero, e16, m2, ta, mu
+; CHECK-NEXT:    vfredosum.vs v8, v8, v10
+; CHECK-NEXT:    vfmv.f.s fa0, v8
+; CHECK-NEXT:    ret
+  %red = call half @llvm.vector.reduce.fadd.nxv6f16(half %s, <vscale x 6 x half> %v)
+  ret half %red
+}
+
+declare half @llvm.vector.reduce.fadd.nxv10f16(half, <vscale x 10 x half>)
+
+define half @vreduce_ord_fadd_nxv10f16(<vscale x 10 x half> %v, half %s) {
+; CHECK-LABEL: vreduce_ord_fadd_nxv10f16:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    srli a0, a0, 2
+; CHECK-NEXT:    add a1, a0, a0
+; CHECK-NEXT:    fmv.h.x ft0, zero
+; CHECK-NEXT:    fneg.h ft0, ft0
+; CHECK-NEXT:    vsetvli a2, zero, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.v.f v12, ft0
+; CHECK-NEXT:    vsetvli zero, a1, e16, m1, tu, mu
+; CHECK-NEXT:    vslideup.vx v10, v12, a0
+; CHECK-NEXT:    vsetvli zero, a0, e16, m1, tu, mu
+; CHECK-NEXT:    vslideup.vi v11, v12, 0
+; CHECK-NEXT:    vsetvli zero, a1, e16, m1, tu, mu
+; CHECK-NEXT:    vslideup.vx v11, v12, a0
+; CHECK-NEXT:    vsetivli zero, 1, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.s.f v12, fa0
+; CHECK-NEXT:    vsetvli a0, zero, e16, m4, ta, mu
+; CHECK-NEXT:    vfredosum.vs v8, v8, v12
+; CHECK-NEXT:    vfmv.f.s fa0, v8
+; CHECK-NEXT:    ret
+  %red = call half @llvm.vector.reduce.fadd.nxv10f16(half %s, <vscale x 10 x half> %v)
+  ret half %red
+}
+
+declare half @llvm.vector.reduce.fadd.nxv12f16(half, <vscale x 12 x half>)
+
+define half @vreduce_ord_fadd_nxv12f16(<vscale x 12 x half> %v, half %s) {
+; CHECK-LABEL: vreduce_ord_fadd_nxv12f16:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 1, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.s.f v12, fa0
+; CHECK-NEXT:    fmv.h.x ft0, zero
+; CHECK-NEXT:    fneg.h ft0, ft0
+; CHECK-NEXT:    vsetvli a0, zero, e16, m1, ta, mu
+; CHECK-NEXT:    vfmv.v.f v11, ft0
+; CHECK-NEXT:    vsetvli a0, zero, e16, m4, ta, mu
+; CHECK-NEXT:    vfredosum.vs v8, v8, v12
+; CHECK-NEXT:    vfmv.f.s fa0, v8
+; CHECK-NEXT:    ret
+  %red = call half @llvm.vector.reduce.fadd.nxv12f16(half %s, <vscale x 12 x half> %v)
+  ret half %red
+}