[llvm] [NVPTX] Improve folding to mad with immediate 1 (PR #93628)
Alex MacLean via llvm-commits
llvm-commits at lists.llvm.org
Tue May 28 18:01:32 PDT 2024
https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/93628
Extend NVPTX DAG combining logic to distribute a mul instruction across an add of 1 into a mad where possible. In addition, add support for transposing a mul through a select with an option of 1, if that would allow further mul folding.
>From 5b9b98a3d9f75ea225cb40eb8d0092d21089a00f Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Wed, 29 May 2024 00:42:25 +0000
Subject: [PATCH] [NVPTX] Improve folding to mad with immediate 1
---
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 93 ++++++++++++++++--
llvm/test/CodeGen/NVPTX/combine-mad.ll | 101 ++++++++++++++++++++
2 files changed, 188 insertions(+), 6 deletions(-)
create mode 100644 llvm/test/CodeGen/NVPTX/combine-mad.ll
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 1e7477cf9d60e..304d1984edd54 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -5614,17 +5614,98 @@ static SDValue TryMULWIDECombine(SDNode *N,
return DCI.DAG.getNode(Opc, DL, MulType, TruncLHS, TruncRHS);
}
+static SDValue matchMADConstOnePattern(SDValue X, SDValue Add) {
+ if (Add->getOpcode() != ISD::ADD)
+ return SDValue();
+
+ SDValue Y = Add->getOperand(0);
+ ConstantSDNode *Const = dyn_cast<ConstantSDNode>(Add->getOperand(1));
+ if (!Const || Const->getZExtValue() != 1)
+ return SDValue();
+
+ return Y;
+}
+
+static SDValue combineMADConstOne(SDValue X, SDValue Add, EVT VT, SDLoc DL,
+ TargetLowering::DAGCombinerInfo &DCI) {
+
+ if (SDValue Y = matchMADConstOnePattern(X, Add))
+ return DCI.DAG.getNode(NVPTXISD::IMAD, DL, VT, X, Y, X);
+
+ return SDValue();
+}
+
+static SDValue combineMulSelectConstOne(SDValue X, SDValue Select, EVT VT,
+ SDLoc DL,
+ TargetLowering::DAGCombinerInfo &DCI) {
+ if (Select->getOpcode() != ISD::SELECT)
+ return SDValue();
+
+ SDValue Cond = Select->getOperand(0);
+
+ unsigned ConstOpNo = 1;
+ auto *Const = dyn_cast<ConstantSDNode>(Select->getOperand(ConstOpNo));
+ if (!Const || Const->getZExtValue() != 1) {
+ ConstOpNo = 2;
+ Const = dyn_cast<ConstantSDNode>(Select->getOperand(ConstOpNo));
+ if (!Const || Const->getZExtValue() != 1)
+ return SDValue();
+ }
+
+ SDValue Y = Select->getOperand((ConstOpNo == 1) ? 2 : 1);
+
+ // Do not combine if the resulting sequence is not obviously profitable.
+ if (!matchMADConstOnePattern(X, Y))
+ return SDValue();
+
+ SDValue NewMul = DCI.DAG.getNode(ISD::MUL, DL, VT, X, Y);
+
+ return DCI.DAG.getNode(ISD::SELECT, DL, VT, Cond,
+ (ConstOpNo == 1) ? X : NewMul,
+ (ConstOpNo == 1) ? NewMul : X);
+}
+
+static SDValue
+PerformMULCombineWithOperands(SDNode *N, SDValue N0, SDValue N1,
+ TargetLowering::DAGCombinerInfo &DCI) {
+
+ EVT VT = N0.getValueType();
+ if (VT.isVector())
+ return SDValue();
+
+ if (VT != MVT::i16 && VT != MVT::i32 && VT != MVT::i64)
+ return SDValue();
+
+ SDLoc DL(N);
+
+ // (mul x, (add y, 1)) -> (mad x, y, x)
+ if (SDValue Res = combineMADConstOne(N0, N1, VT, DL, DCI))
+ return Res;
+ if (SDValue Res = combineMADConstOne(N1, N0, VT, DL, DCI))
+ return Res;
+
+ // (mul x, (select y, 1)) -> (select (mul x, y), x)
+ if (SDValue Res = combineMulSelectConstOne(N0, N1, VT, DL, DCI))
+ return Res;
+ if (SDValue Res = combineMulSelectConstOne(N1, N0, VT, DL, DCI))
+ return Res;
+
+ return SDValue();
+}
+
/// PerformMULCombine - Runs PTX-specific DAG combine patterns on MUL nodes.
static SDValue PerformMULCombine(SDNode *N,
TargetLowering::DAGCombinerInfo &DCI,
CodeGenOptLevel OptLevel) {
- if (OptLevel > CodeGenOptLevel::None) {
- // Try mul.wide combining at OptLevel > 0
- if (SDValue Ret = TryMULWIDECombine(N, DCI))
- return Ret;
- }
+ if (OptLevel == CodeGenOptLevel::None)
+ return SDValue();
- return SDValue();
+ if (SDValue Ret = TryMULWIDECombine(N, DCI))
+ return Ret;
+
+ SDValue N0 = N->getOperand(0);
+ SDValue N1 = N->getOperand(1);
+ return PerformMULCombineWithOperands(N, N0, N1, DCI);
}
/// PerformSHLCombine - Runs PTX-specific DAG combine patterns on SHL nodes.
diff --git a/llvm/test/CodeGen/NVPTX/combine-mad.ll b/llvm/test/CodeGen/NVPTX/combine-mad.ll
new file mode 100644
index 0000000000000..382856dfe76c3
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/combine-mad.ll
@@ -0,0 +1,101 @@
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -O1 | FileCheck %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -O1 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -O1 | %ptxas-verify %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 -O1 | %ptxas-verify %}
+
+define i32 @test1(i32 %n, i32 %m) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test1_param_0];
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test1_param_1];
+; CHECK: mad.lo.s32 %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: st.param.b32 [func_retval0+0], %[[MAD]];
+;
+ %add = add i32 %n, 1
+ %mul = mul i32 %add, %m
+ ret i32 %mul
+}
+
+define i32 @test1_rev(i32 %n, i32 %m) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test1_rev_param_0];
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test1_rev_param_1];
+; CHECK: mad.lo.s32 %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: st.param.b32 [func_retval0+0], %[[MAD]];
+;
+ %add = add i32 %n, 1
+ %mul = mul i32 %m, %add
+ ret i32 %mul
+}
+
+; Transpose (mul (select)) if it can then be folded to mad
+define i32 @test2(i32 %n, i32 %m, i32 %s) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test2_param_0];
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test2_param_1];
+; CHECK: ld.param.u32 %[[S:r[0-9]+]], [test2_param_2];
+; CHECK: setp.lt.s32 %[[COND:p[0-9]+]], %[[S]], 1;
+; CHECK: mad.lo.s32 %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: selp.b32 %[[SEL:r[0-9]+]], %[[M]], %[[MAD]], %[[COND]];
+; CHECK: st.param.b32 [func_retval0+0], %[[SEL]];
+;
+ %add = add i32 %n, 1
+ %cond = icmp slt i32 %s, 1
+ %sel = select i1 %cond, i32 1, i32 %add
+ %mul = mul i32 %sel, %m
+ ret i32 %mul
+}
+
+;; Transpose (mul (select)) if it can then be folded to mad
+define i32 @test2_rev1(i32 %n, i32 %m, i32 %s) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test2_rev1_param_0];
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test2_rev1_param_1];
+; CHECK: ld.param.u32 %[[S:r[0-9]+]], [test2_rev1_param_2];
+; CHECK: setp.lt.s32 %[[COND:p[0-9]+]], %[[S]], 1;
+; CHECK: mad.lo.s32 %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: selp.b32 %[[SEL:r[0-9]+]], %[[MAD]], %[[M]], %[[COND]];
+; CHECK: st.param.b32 [func_retval0+0], %[[SEL]];
+;
+ %add = add i32 %n, 1
+ %cond = icmp slt i32 %s, 1
+ %sel = select i1 %cond, i32 %add, i32 1
+ %mul = mul i32 %sel, %m
+ ret i32 %mul
+}
+
+;; Transpose (mul (select)) if it can then be folded to mad
+define i32 @test2_rev2(i32 %n, i32 %m, i32 %s) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test2_rev2_param_0];
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test2_rev2_param_1];
+; CHECK: ld.param.u32 %[[S:r[0-9]+]], [test2_rev2_param_2];
+; CHECK: setp.lt.s32 %[[COND:p[0-9]+]], %[[S]], 1;
+; CHECK: mad.lo.s32 %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: selp.b32 %[[SEL:r[0-9]+]], %[[MAD]], %[[M]], %[[COND]];
+; CHECK: st.param.b32 [func_retval0+0], %[[SEL]];
+;
+ %add = add i32 %n, 1
+ %cond = icmp slt i32 %s, 1
+ %sel = select i1 %cond, i32 %add, i32 1
+ %mul = mul i32 %m, %sel
+ ret i32 %mul
+}
+
+;; Leave (mul (select)) intact if it transposing is not profitable
+define i32 @test3(i32 %n, i32 %m, i32 %s) {
+;
+; CHECK: ld.param.u32 %[[N:r[0-9]+]], [test3_param_0];
+; CHECK: add.s32 %[[ADD:r[0-9]+]], %[[N]], 3;
+; CHECK: ld.param.u32 %[[M:r[0-9]+]], [test3_param_1];
+; CHECK: ld.param.u32 %[[S:r[0-9]+]], [test3_param_2];
+; CHECK: setp.lt.s32 %[[COND:p[0-9]+]], %[[S]], 1;
+; CHECK: selp.b32 %[[SEL:r[0-9]+]], 1, %[[ADD]], %[[COND]];
+; CHECK: mul.lo.s32 %[[MUL:r[0-9]+]], %[[SEL]], %[[M]];
+; CHECK: st.param.b32 [func_retval0+0], %[[MUL]];
+;
+ %add = add i32 %n, 3
+ %cond = icmp slt i32 %s, 1
+ %sel = select i1 %cond, i32 1, i32 %add
+ %mul = mul i32 %sel, %m
+ ret i32 %mul
+}
More information about the llvm-commits
mailing list