[AArch64][SelectionDAG] Add support for 8to64 partial reduction cases #138269

NickGuy-Arm · 2025-05-02T13:08:11Z

No description provided.

llvmbot · 2025-05-02T13:08:51Z

@llvm/pr-subscribers-backend-aarch64

Author: Nicholas Guy (NickGuy-Arm)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/138269.diff

3 Files Affected:

(modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+28)
(modified) llvm/lib/Target/AArch64/AArch64ISelLowering.h (+1)
(modified) llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll (+12-74)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 0126b97c9fb9a..ca63473ab085b 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1867,6 +1867,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     // Other pairs will default to 'Expand'.
     setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv8i16, Legal);
     setPartialReduceMLAAction(MVT::nxv4i32, MVT::nxv16i8, Legal);
+
+    setPartialReduceMLAAction(MVT::nxv2i64, MVT::nxv16i8, Custom);
   }
 
   // Handle operations that are only available in non-streaming SVE mode.
@@ -7767,6 +7769,9 @@ SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
     return LowerFLDEXP(Op, DAG);
   case ISD::EXPERIMENTAL_VECTOR_HISTOGRAM:
     return LowerVECTOR_HISTOGRAM(Op, DAG);
+  case ISD::PARTIAL_REDUCE_SMLA:
+  case ISD::PARTIAL_REDUCE_UMLA:
+    return LowerPARTIAL_REDUCE_MLA(Op, DAG);
   }
 }
 
@@ -29509,6 +29514,29 @@ SDValue AArch64TargetLowering::LowerVECTOR_HISTOGRAM(SDValue Op,
   return Scatter;
 }
 
+SDValue
+AArch64TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
+                                               SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+
+  auto Acc = Op.getOperand(0);
+  auto LHS = Op.getOperand(1);
+  auto RHS = Op.getOperand(2);
+
+  auto ResultVT = Op.getValueType();
+
+  assert(ResultVT == MVT::nxv2i64 && LHS.getValueType() == MVT::nxv16i8);
+
+  auto NewAcc = DAG.getConstant(0, DL, MVT::nxv4i32);
+  auto DotNode =
+      DAG.getNode(Op.getOpcode(), DL, MVT::nxv4i32, NewAcc, LHS, RHS);
+
+  auto Lo = DAG.getNode(AArch64ISD::UUNPKLO, DL, ResultVT, DotNode);
+  auto Hi = DAG.getNode(AArch64ISD::UUNPKHI, DL, ResultVT, DotNode);
+  auto Extended = DAG.getNode(ISD::ADD, DL, ResultVT, Lo, Hi);
+  return DAG.getNode(ISD::ADD, DL, ResultVT, Acc, Extended);
+}
+
 SDValue
 AArch64TargetLowering::LowerFixedLengthFPToIntToSVE(SDValue Op,
                                                     SelectionDAG &DAG) const {
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index d9b535b910b80..9d8d1c22258be 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -1181,6 +1181,7 @@ class AArch64TargetLowering : public TargetLowering {
   SDValue LowerVECTOR_DEINTERLEAVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECTOR_INTERLEAVE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECTOR_HISTOGRAM(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerPARTIAL_REDUCE_MLA(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDIV(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerMUL(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorSRA_SRL_SHL(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll b/llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll
index 039cac01008b8..4d5996ee955ce 100644
--- a/llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll
+++ b/llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll
@@ -198,43 +198,12 @@ define <vscale x 4 x i64> @udot_8to64(<vscale x 4 x i64> %acc, <vscale x 16 x i8
 ;
 ; CHECK-NEWLOWERING-LABEL: udot_8to64:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z4.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    uunpklo z2.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z5.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    uunpklo z3.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    ptrue p0.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z6.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z7.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z24.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z25.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z4.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z5.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    uunpklo z26.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z27.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z28.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z29.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z6.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z7.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z24.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z25.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z26.d, z4.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z28.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z29.d
-; CHECK-NEWLOWERING-NEXT:    uunpklo z27.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    uunpklo z29.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z4.d, z4.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z5.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    uunpkhi z3.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z6.d, z24.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z7.d, z25.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z29.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z4.d, z5.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z2.d, z3.d
+; CHECK-NEWLOWERING-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-NEXT:    udot z4.s, z2.b, z3.b
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z3.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    add z2.d, z3.d, z2.d
+; CHECK-NEWLOWERING-NEXT:    add z0.d, z0.d, z2.d
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
   %a.wide = zext <vscale x 16 x i8> %a to <vscale x 16 x i64>
@@ -258,43 +227,12 @@ define <vscale x 4 x i64> @sdot_8to64(<vscale x 4 x i64> %acc, <vscale x 16 x i8
 ;
 ; CHECK-NEWLOWERING-LABEL: sdot_8to64:
 ; CHECK-NEWLOWERING:       // %bb.0: // %entry
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z4.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    sunpklo z2.h, z2.b
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z5.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    sunpklo z3.h, z3.b
-; CHECK-NEWLOWERING-NEXT:    ptrue p0.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z6.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z7.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z24.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z25.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z4.s, z4.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z2.s, z2.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z5.s, z5.h
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.s, z3.h
-; CHECK-NEWLOWERING-NEXT:    sunpklo z26.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z27.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z28.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z29.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z6.d, z6.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z7.d, z7.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z24.d, z24.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z25.d, z25.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z26.d, z4.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z28.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z29.d
-; CHECK-NEWLOWERING-NEXT:    sunpklo z27.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    sunpklo z29.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z4.d, z4.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z2.d, z2.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z5.d, z5.s
-; CHECK-NEWLOWERING-NEXT:    sunpkhi z3.d, z3.s
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z6.d, z24.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z7.d, z25.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z26.d, z28.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z27.d, z29.d
-; CHECK-NEWLOWERING-NEXT:    mla z1.d, p0/m, z4.d, z5.d
-; CHECK-NEWLOWERING-NEXT:    mla z0.d, p0/m, z2.d, z3.d
+; CHECK-NEWLOWERING-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-NEXT:    sdot z4.s, z2.b, z3.b
+; CHECK-NEWLOWERING-NEXT:    uunpkhi z2.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uunpklo z3.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    add z2.d, z3.d, z2.d
+; CHECK-NEWLOWERING-NEXT:    add z0.d, z0.d, z2.d
 ; CHECK-NEWLOWERING-NEXT:    ret
 entry:
   %a.wide = sext <vscale x 16 x i8> %a to <vscale x 16 x i64>

MacDue · 2025-05-02T13:39:26Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+  auto Lo = DAG.getNode(AArch64ISD::UUNPKLO, DL, ResultVT, DotNode);
+  auto Hi = DAG.getNode(AArch64ISD::UUNPKHI, DL, ResultVT, DotNode);
+  auto Extended = DAG.getNode(ISD::ADD, DL, ResultVT, Lo, Hi);


Maybe an extension... but I think this could use the top/bottom wide adds rather than two unpacks then an add (for SVE2/SME).

Also, is UUNPKLO/HI correct for sdot? IIRC it only zero extends, I think that case needs SUNPKLO/HI.

The resulting assembly is certainly cleaner, done.
(As well as respecting whether it's signed/unsigned)

huntergr-arm · 2025-05-02T15:23:11Z

llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll

+; CHECK-NEWLOWERING-NEXT:    movi v4.2d, #0000000000000000
+; CHECK-NEWLOWERING-NEXT:    udot z4.s, z2.b, z3.b
+; CHECK-NEWLOWERING-NEXT:    uaddwb z0.d, z0.d, z4.s
+; CHECK-NEWLOWERING-NEXT:    uaddwt z0.d, z0.d, z4.s


I'm wondering slightly about only updating half the accumulator (<vscale x 4 x i64> will be distributed across two registers), but I believe that's fine under the definition of a partial reduction. Certainly much nicer looking than the previous codegen.

MacDue · 2025-05-02T15:32:50Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+  unsigned LoOpcode = IsUnsigned ? AArch64ISD::UADDWB : AArch64ISD::SADDWB;
+  unsigned HiOpcode = IsUnsigned ? AArch64ISD::UADDWT : AArch64ISD::SADDWT;


You need to check Subtarget->hasSVE2() || Subtarget->isStreamingSVEAvailable() before using the wide adds IIRC. They're gated under if !IsFeatureImplemented(FEAT_SVE2) && !IsFeatureImplemented(FEAT_SME).

I had that at first, but removed it as the PartialReduceMLAAction was only set if we had SVE, not realising that these instructions were specifically SVE2. Added, as well as a fallback to the old unpack/add approach for when we don't have SVE2.

sdesmalen-arm

That's a nice improvement!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

sdesmalen-arm

Just left two nits. Otherwise it looks good to me!

sdesmalen-arm · 2025-05-02T19:50:27Z

llvm/test/CodeGen/AArch64/sve-partial-reduce-dot-product.ll

 ; RUN: llc -mtriple=aarch64 -mattr=+sve2 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-NOI8MM
-; RUN: llc -mtriple=aarch64 -mattr=+sve2,+i8mm -aarch64-enable-partial-reduce-nodes %s -o - | FileCheck %s --check-prefixes=CHECK-NEWLOWERING
+; RUN: llc -mtriple=aarch64 -mattr=+sve,+i8mm -aarch64-enable-partial-reduce-nodes %s -o - | FileCheck %s --check-prefixes=CHECK-NEWLOWERING,CHECK-NEWLOWERING-SVE
+; RUN: llc -mtriple=aarch64 -mattr=+sve2,+i8mm -aarch64-enable-partial-reduce-nodes %s -o - | FileCheck %s --check-prefixes=CHECK-NEWLOWERING,CHECK-NEWLOWERING-SVE2


Could you also add a run line for +sme -force-streaming [...] --check-prefixes=CHECK-NEWLOWERING,CHECK-NEWLOWERING-SME ?

sdesmalen-arm · 2025-05-02T19:55:05Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+/// If a PARTIAL_REDUCE_MLA node comes in with an accumulator type that is too
+/// wide to be used for (u|s)dot, we can still make use of the dot product


nit: rather than speaking of 'types that are too wide', I'd suggest being explicit about the types, e.g. we can't accumulate from nxv16i8 -> nxv2i64 in one step, but we can accumulate in two steps: nxv16i8 -> nxv4i32 -> nxv2i64.

…llvm#138269) --------- Co-authored-by: James Chesterman <[email protected]>

[AArch64][SelectionDAG] Add support for 8to64 partial reduction cases

3187e07

NickGuy-Arm requested review from JamesChesterman, MacDue, SamTebbs33, huntergr-arm and sdesmalen-arm May 2, 2025 13:08

llvmbot added the backend:AArch64 label May 2, 2025

MacDue reviewed May 2, 2025

View reviewed changes

Switch to using wide adds instead of unpacking

00371df

huntergr-arm reviewed May 2, 2025

View reviewed changes

MacDue reviewed May 2, 2025

View reviewed changes

sdesmalen-arm reviewed May 2, 2025

View reviewed changes

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp Outdated Show resolved Hide resolved

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp Show resolved Hide resolved

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp Outdated Show resolved Hide resolved

Address comments and add SVE2/SME guard & fallback

1d0994b

sdesmalen-arm approved these changes May 2, 2025

View reviewed changes

Update doc-comment and add sme test run line

466eb9f

NickGuy-Arm merged commit e86a92f into llvm:main May 6, 2025
9 of 10 checks passed

GeorgeARM pushed a commit to GeorgeARM/llvm-project that referenced this pull request May 7, 2025

[AArch64][SelectionDAG] Add support for 8to64 partial reduction cases (…

bae827f

…llvm#138269) --------- Co-authored-by: James Chesterman <[email protected]>

		unsigned LoOpcode = IsUnsigned ? AArch64ISD::UADDWB : AArch64ISD::SADDWB;
		unsigned HiOpcode = IsUnsigned ? AArch64ISD::UADDWT : AArch64ISD::SADDWT;

		/// If a PARTIAL_REDUCE_MLA node comes in with an accumulator type that is too
		/// wide to be used for (u\|s)dot, we can still make use of the dot product

[AArch64][SelectionDAG] Add support for 8to64 partial reduction cases #138269

[AArch64][SelectionDAG] Add support for 8to64 partial reduction cases #138269

Uh oh!

Conversation

NickGuy-Arm commented May 2, 2025

Uh oh!

llvmbot commented May 2, 2025

Uh oh!

MacDue May 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MacDue May 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NickGuy-Arm May 2, 2025

Choose a reason for hiding this comment

Uh oh!

huntergr-arm May 2, 2025

Choose a reason for hiding this comment

Uh oh!

MacDue May 2, 2025

Choose a reason for hiding this comment

Uh oh!

NickGuy-Arm May 2, 2025

Choose a reason for hiding this comment

Uh oh!

sdesmalen-arm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sdesmalen-arm left a comment

Choose a reason for hiding this comment

Uh oh!

sdesmalen-arm May 2, 2025

Choose a reason for hiding this comment

Uh oh!

sdesmalen-arm May 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

MacDue May 2, 2025 •

edited

Loading

MacDue May 2, 2025 •

edited

Loading