[llvm] [SDAG] Fix CSE for ADDRSPACECAST nodes (PR #122912)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/122912

Correct CSE in SelectionDAG can make DAG combining more effective and reduces the size of the DAG and thus should improve compile time.

This change is a prerequisite for https://github.com/llvm/llvm-project/pull/121710 as it enables the removal or redundant loads of byval arguments in some cases.

>From 7176ce0db58dca7531c37cb0ce2fe326b21134f4 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Tue, 14 Jan 2025 15:03:35 +0000
Subject: [PATCH] [SDAG] Fix CSE for ADDRSPACECAST nodes

---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  6 ++++++
 llvm/test/CodeGen/NVPTX/addrspacecast-cse.ll  | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 llvm/test/CodeGen/NVPTX/addrspacecast-cse.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 0dfd0302ae5438..743ae4895a1b1c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -954,6 +954,12 @@ static void AddNodeIDCustom(FoldingSetNodeID &ID, const SDNode *N) {
       ID.AddInteger(M);
     break;
   }
+  case ISD::ADDRSPACECAST: {
+    const AddrSpaceCastSDNode *ASC = cast<AddrSpaceCastSDNode>(N);
+    ID.AddInteger(ASC->getSrcAddressSpace());
+    ID.AddInteger(ASC->getDestAddressSpace());
+    break;
+  }
   case ISD::TargetBlockAddress:
   case ISD::BlockAddress: {
     const BlockAddressSDNode *BA = cast<BlockAddressSDNode>(N);
diff --git a/llvm/test/CodeGen/NVPTX/addrspacecast-cse.ll b/llvm/test/CodeGen/NVPTX/addrspacecast-cse.ll
new file mode 100644
index 00000000000000..a1bf8042817b43
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/addrspacecast-cse.ll
@@ -0,0 +1,19 @@
+; RUN: llc < %s -mcpu=sm_80 -mattr=+ptx73 -debug-only=isel -o /dev/null 2>&1 | FileCheck %s
+
+; REQUIRES: asserts
+
+target triple = "nvptx64-nvidia-cuda"
+
+;; Selection DAG CSE is hard to test since we run CSE/GVN on the IR before and
+;; after selection DAG ISel so most cases will be handled by one of these.
+define void @foo(ptr %p) {
+; CHECK-LABEL: Optimized legalized selection DAG: %bb.0 'foo:'
+; CHECK:       addrspacecast[0 -> 5]
+; CHECK-NOT:   addrspacecast[0 -> 5]
+; CHECK-LABEL: ===== Instruction selection begins
+;
+  %a1 = addrspacecast ptr %p to ptr addrspace(5)
+  call void @llvm.stackrestore(ptr %p)
+  store ptr %p, ptr addrspace(5) %a1
+  ret void
+}