[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