From a264065fbbe8aae8b084bbac31aadc05e1b314c0 Mon Sep 17 00:00:00 2001 From: Philip Reames Date: Wed, 27 Mar 2024 14:27:06 -0700 Subject: [PATCH 1/2] [RISCV] Support 3-argument associative add for transformAddShlImm This transform is looking for the shNadd idiom for zba, but that can be obscured if there's another value being added to the result. The choice to restrict to one level of association is tactical - we could of course do more, but there's the usual compile time tradeoff, and this covers the motivating example. This is a solution to a reduced test case originally flagged in the description of https://github.com/llvm/llvm-project/pull/85734. --- llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 54 ++++++++++++++++----- llvm/test/CodeGen/RISCV/rv64zba.ll | 10 ++-- 2 files changed, 47 insertions(+), 17 deletions(-) diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp index 564fda674317f..14ace1d30a411 100644 --- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp +++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp @@ -12733,20 +12733,13 @@ static SDValue combineBinOpToReduce(SDNode *N, SelectionDAG &DAG, // Optimize (add (shl x, c0), (shl y, c1)) -> // (SLLI (SH*ADD x, y), c0), if c1-c0 equals to [1|2|3]. -static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG, +static SDValue transformAddShlImm(SDValue N0, SDValue N1, SDLoc DL, + SelectionDAG &DAG, const RISCVSubtarget &Subtarget) { - // Perform this optimization only in the zba extension. - if (!Subtarget.hasStdExtZba()) - return SDValue(); - // Skip for vector types and larger types. - EVT VT = N->getValueType(0); - if (VT.isVector() || VT.getSizeInBits() > Subtarget.getXLen()) - return SDValue(); + EVT VT = N0.getValueType(); // The two operand nodes must be SHL and have no other use. - SDValue N0 = N->getOperand(0); - SDValue N1 = N->getOperand(1); if (N0->getOpcode() != ISD::SHL || N1->getOpcode() != ISD::SHL || !N0->hasOneUse() || !N1->hasOneUse()) return SDValue(); @@ -12768,7 +12761,6 @@ static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG, return SDValue(); // Build nodes. - SDLoc DL(N); SDValue NS = (C0 < C1) ? N0->getOperand(0) : N1->getOperand(0); SDValue NL = (C0 > C1) ? N0->getOperand(0) : N1->getOperand(0); SDValue NA0 = @@ -12777,6 +12769,46 @@ static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG, return DAG.getNode(ISD::SHL, DL, VT, NA1, DAG.getConstant(Bits, DL, VT)); } +// Generalized form of above which looks through one level of add +// reassociation for oppurtunities. +static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG, + const RISCVSubtarget &Subtarget) { + // Perform this optimization only in the zba extension. + if (!Subtarget.hasStdExtZba()) + return SDValue(); + + // Skip for vector types and larger types. + EVT VT = N->getValueType(0); + if (VT.isVector() || VT.getSizeInBits() > Subtarget.getXLen()) + return SDValue(); + + // We're look for two SHL nodes in the add tree with all nodes + // involved having no other use. + SDValue N0 = N->getOperand(0); + SDValue N1 = N->getOperand(1); + if (N0->getOpcode() != ISD::SHL) + std::swap(N0, N1); + + if (SDValue Res = transformAddShlImm(N0, N1, SDLoc(N), DAG, Subtarget)) + return Res; + + if (N0->getOpcode() != ISD::SHL || N1->getOpcode() != ISD::ADD || + !N1->hasOneUse()) + return SDValue(); + + // Allow reassociation for a 3-argument add + SDLoc DL(N1); + SDValue A = N1->getOperand(0); + SDValue B = N1->getOperand(1); + if (SDValue Res = transformAddShlImm(N0, A, SDLoc(N), DAG, Subtarget)) + return DAG.getNode(ISD::ADD, DL, VT, Res, B); + + if (SDValue Res = transformAddShlImm(N0, B, SDLoc(N), DAG, Subtarget)) + return DAG.getNode(ISD::ADD, DL, VT, Res, A); + + return SDValue(); +} + // Combine a constant select operand into its use: // // (and (select cond, -1, c), x) diff --git a/llvm/test/CodeGen/RISCV/rv64zba.ll b/llvm/test/CodeGen/RISCV/rv64zba.ll index d9d83633a8537..c09cf1b6f4844 100644 --- a/llvm/test/CodeGen/RISCV/rv64zba.ll +++ b/llvm/test/CodeGen/RISCV/rv64zba.ll @@ -1315,9 +1315,8 @@ define i64 @sh6_sh3_add2(i64 noundef %x, i64 noundef %y, i64 noundef %z) { ; ; RV64ZBA-LABEL: sh6_sh3_add2: ; RV64ZBA: # %bb.0: # %entry -; RV64ZBA-NEXT: slli a1, a1, 6 -; RV64ZBA-NEXT: add a0, a1, a0 -; RV64ZBA-NEXT: sh3add a0, a2, a0 +; RV64ZBA-NEXT: sh3add a1, a1, a2 +; RV64ZBA-NEXT: sh3add a0, a1, a0 ; RV64ZBA-NEXT: ret entry: %shl = shl i64 %z, 3 @@ -1360,9 +1359,8 @@ define i64 @sh6_sh3_add4(i64 noundef %x, i64 noundef %y, i64 noundef %z) { ; ; RV64ZBA-LABEL: sh6_sh3_add4: ; RV64ZBA: # %bb.0: # %entry -; RV64ZBA-NEXT: slli a1, a1, 6 -; RV64ZBA-NEXT: sh3add a0, a2, a0 -; RV64ZBA-NEXT: add a0, a0, a1 +; RV64ZBA-NEXT: sh3add a1, a1, a2 +; RV64ZBA-NEXT: sh3add a0, a1, a0 ; RV64ZBA-NEXT: ret entry: %shl = shl i64 %z, 3 From 244d52c80614668654e3104590744d1172c998aa Mon Sep 17 00:00:00 2001 From: Philip Reames Date: Wed, 27 Mar 2024 15:34:03 -0700 Subject: [PATCH 2/2] Address review comments --- llvm/lib/Target/RISCV/RISCVISelLowering.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp index 14ace1d30a411..bab09bf642179 100644 --- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp +++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp @@ -12733,7 +12733,7 @@ static SDValue combineBinOpToReduce(SDNode *N, SelectionDAG &DAG, // Optimize (add (shl x, c0), (shl y, c1)) -> // (SLLI (SH*ADD x, y), c0), if c1-c0 equals to [1|2|3]. -static SDValue transformAddShlImm(SDValue N0, SDValue N1, SDLoc DL, +static SDValue transformAddShlImm(SDValue N0, SDValue N1, const SDLoc &DL, SelectionDAG &DAG, const RISCVSubtarget &Subtarget) { @@ -12770,7 +12770,7 @@ static SDValue transformAddShlImm(SDValue N0, SDValue N1, SDLoc DL, } // Generalized form of above which looks through one level of add -// reassociation for oppurtunities. +// reassociation for opportunities. static SDValue transformAddShlImm(SDNode *N, SelectionDAG &DAG, const RISCVSubtarget &Subtarget) { // Perform this optimization only in the zba extension.