[llvm] [NVPTX] Emit `NVPTXISD::DYNAMIC_STACKALLOC`'s chain (PR #101714)

[Justin Fargnoli via llvm-commits]

https://github.com/justinfargnoli created https://github.com/llvm/llvm-project/pull/101714

`LowerDYNAMIC_STACKALLOC()` emits the `dynamic_stackalloc` chain operand instead of the chain produced by the `NVPTXISD::DYNAMIC_STACKALLOC`. Fix this behavior and don't produce an unnecessary `ISD::MERGE_VALUES`.

>From 5a07d8168820cde69c592dc008a97aa8251cf3d4 Mon Sep 17 00:00:00 2001
From: Justin Fargnoli <jfargnoli at nvidia.com>
Date: Wed, 31 Jul 2024 13:28:33 -0700
Subject: [PATCH 1/3] [NVPTX] Emit `NVPTXISD::DYNAMIC_STACKALLOC`'s chain

---
 llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp   | 15 ++++------
 .../NVPTX/dynamic-stackalloca-regression.ll   | 29 +++++++++++++++++++
 2 files changed, 34 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 6975412ce5d35..8a0f1691dc8a6 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -2232,18 +2232,13 @@ SDValue NVPTXTargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
   SDLoc DL(Op.getNode());
 
   // The size for ptx alloca instruction is 64-bit for m64 and 32-bit for m32.
-  if (nvTM->is64Bit())
-    Size = DAG.getZExtOrTrunc(Size, DL, MVT::i64);
-  else
-    Size = DAG.getZExtOrTrunc(Size, DL, MVT::i32);
+  MVT ValueSizeTy = nvTM->is64Bit() ? MVT::i64 : MVT::i32;
 
-  SDValue AllocOps[] = {Chain, Size,
+  SDValue AllocOps[] = {Chain, DAG.getZExtOrTrunc(Size, DL, ValueSizeTy),
                         DAG.getTargetConstant(Align, DL, MVT::i32)};
-  SDValue Alloca = DAG.getNode(NVPTXISD::DYNAMIC_STACKALLOC, DL,
-                               nvTM->is64Bit() ? MVT::i64 : MVT::i32, AllocOps);
-
-  SDValue MergeOps[] = {Alloca, Chain};
-  return DAG.getMergeValues(MergeOps, DL);
+  EVT RetTypes[] = {ValueSizeTy, MVT::Other};  
+  return DAG.getNode(NVPTXISD::DYNAMIC_STACKALLOC, DL,
+                               RetTypes, AllocOps);
 }
 
 // By default CONCAT_VECTORS is lowered by ExpandVectorBuildThroughStack()
diff --git a/llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll b/llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll
new file mode 100644
index 0000000000000..16255893b1c3d
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll
@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -march=nvptx64 -mattr=+ptx73 -mcpu=sm_52 | FileCheck %s
+
+target triple = "nvptx64-nvidia-cuda"
+
+define void @foo(i64 %a, ptr %p0, ptr %p1) {
+; CHECK-LABEL: foo(
+; CHECK:       {
+; CHECK-NEXT:    .reg .b64 %rd<8>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    ld.param.u64 %rd1, [foo_param_0];
+; CHECK-NEXT:    add.s64 %rd2, %rd1, 7;
+; CHECK-NEXT:    and.b64 %rd3, %rd2, -8;
+; CHECK-NEXT:    alloca.u64 %rd4, %rd3, 16;
+; CHECK-NEXT:    cvta.local.u64 %rd4, %rd4;
+; CHECK-NEXT:    ld.param.u64 %rd5, [foo_param_1];
+; CHECK-NEXT:    alloca.u64 %rd6, %rd3, 16;
+; CHECK-NEXT:    cvta.local.u64 %rd6, %rd6;
+; CHECK-NEXT:    ld.param.u64 %rd7, [foo_param_2];
+; CHECK-NEXT:    st.u64 [%rd5], %rd4;
+; CHECK-NEXT:    st.u64 [%rd7], %rd6;
+; CHECK-NEXT:    ret;
+  %b = alloca i8, i64 %a, align 16
+  %c = alloca i8, i64 %a, align 16
+  store ptr %b, ptr %p0, align 8
+  store ptr %c, ptr %p1, align 8
+  ret void
+}

>From 7ff6fd8ff3433c63dc5cbb351bcee2b8a2cd837d Mon Sep 17 00:00:00 2001
From: Justin Fargnoli <jfargnoli at nvidia.com>
Date: Wed, 31 Jul 2024 13:28:54 -0700
Subject: [PATCH 2/3] clang-format

---
 llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 8a0f1691dc8a6..516fc7339a4bf 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -2236,9 +2236,8 @@ SDValue NVPTXTargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
 
   SDValue AllocOps[] = {Chain, DAG.getZExtOrTrunc(Size, DL, ValueSizeTy),
                         DAG.getTargetConstant(Align, DL, MVT::i32)};
-  EVT RetTypes[] = {ValueSizeTy, MVT::Other};  
-  return DAG.getNode(NVPTXISD::DYNAMIC_STACKALLOC, DL,
-                               RetTypes, AllocOps);
+  EVT RetTypes[] = {ValueSizeTy, MVT::Other};
+  return DAG.getNode(NVPTXISD::DYNAMIC_STACKALLOC, DL, RetTypes, AllocOps);
 }
 
 // By default CONCAT_VECTORS is lowered by ExpandVectorBuildThroughStack()

>From bef15d79551663cb0a535bbdc11de2c842921a0e Mon Sep 17 00:00:00 2001
From: Justin Fargnoli <jfargnoli at nvidia.com>
Date: Wed, 31 Jul 2024 17:12:46 -0700
Subject: [PATCH 3/3] Rename file

---
 ...stackalloca-regression.ll => dynamic-stackalloc-regression.ll} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename llvm/test/CodeGen/NVPTX/{dynamic-stackalloca-regression.ll => dynamic-stackalloc-regression.ll} (100%)

diff --git a/llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll b/llvm/test/CodeGen/NVPTX/dynamic-stackalloc-regression.ll
similarity index 100%
rename from llvm/test/CodeGen/NVPTX/dynamic-stackalloca-regression.ll
rename to llvm/test/CodeGen/NVPTX/dynamic-stackalloc-regression.ll