[llvm] [MergeFuncs] Account for elementwise loads in FunctionComparator (PR #211478)

[Harrison Hao via llvm-commits]

https://github.com/harrisonGPU updated https://github.com/llvm/llvm-project/pull/211478

>From 2b561d5fdec0ed56254ae759ce80faf3181d55e2 Mon Sep 17 00:00:00 2001
From: Harrison Hao <tsworld1314 at gmail.com>
Date: Thu, 23 Jul 2026 15:13:29 +0800
Subject: [PATCH 1/2] [MergeFuncs] Account for elementwise loads in
 FunctionComparator

---
 .../Transforms/Utils/FunctionComparator.cpp   |  3 +++
 .../MergeFunc/atomic-load-elementwise.ll      | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll

diff --git a/llvm/lib/Transforms/Utils/FunctionComparator.cpp b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
index 80b02d579f5da..f3b595e05d881 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -699,6 +699,9 @@ int FunctionComparator::cmpOperations(const Instruction *L,
   if (const LoadInst *LI = dyn_cast<LoadInst>(L)) {
     if (int Res = cmpNumbers(LI->isVolatile(), cast<LoadInst>(R)->isVolatile()))
       return Res;
+    if (int Res =
+            cmpNumbers(LI->isElementwise(), cast<LoadInst>(R)->isElementwise()))
+      return Res;
     if (int Res = cmpAligns(LI->getAlign(), cast<LoadInst>(R)->getAlign()))
       return Res;
     if (int Res =
diff --git a/llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll b/llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll
new file mode 100644
index 0000000000000..7bb76f61a8a49
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll
@@ -0,0 +1,22 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=mergefunc %s | FileCheck %s
+
+define internal <2 x i32> @whole_vector_load(ptr %p) {
+; CHECK-LABEL: define internal <2 x i32> @whole_vector_load(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:    [[VALUE:%.*]] = load atomic <2 x i32>, ptr [[P]] monotonic, align 4
+; CHECK-NEXT:    ret <2 x i32> [[VALUE]]
+;
+  %value = load atomic <2 x i32>, ptr %p monotonic, align 4
+  ret <2 x i32> %value
+}
+
+define internal <2 x i32> @elementwise_load(ptr %p) {
+; CHECK-LABEL: define internal <2 x i32> @elementwise_load(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:    [[VALUE:%.*]] = load atomic elementwise <2 x i32>, ptr [[P]] monotonic, align 4
+; CHECK-NEXT:    ret <2 x i32> [[VALUE]]
+;
+  %value = load atomic elementwise <2 x i32>, ptr %p monotonic, align 4
+  ret <2 x i32> %value
+}

>From b476a593bf0e5ba7f336eef8708ccde273f40c33 Mon Sep 17 00:00:00 2001
From: Harrison Hao <tsworld1314 at gmail.com>
Date: Thu, 23 Jul 2026 17:26:35 +0800
Subject: [PATCH 2/2] Support atomicrmw

---
 .../Transforms/Utils/FunctionComparator.cpp   |  3 +++
 ...d-elementwise.ll => atomic-elementwise.ll} | 20 +++++++++++++++++++
 2 files changed, 23 insertions(+)
 rename llvm/test/Transforms/MergeFunc/{atomic-load-elementwise.ll => atomic-elementwise.ll} (50%)

diff --git a/llvm/lib/Transforms/Utils/FunctionComparator.cpp b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
index f3b595e05d881..c9cc62adc2604 100644
--- a/llvm/lib/Transforms/Utils/FunctionComparator.cpp
+++ b/llvm/lib/Transforms/Utils/FunctionComparator.cpp
@@ -800,6 +800,9 @@ int FunctionComparator::cmpOperations(const Instruction *L,
     if (int Res = cmpNumbers(RMWI->isVolatile(),
                              cast<AtomicRMWInst>(R)->isVolatile()))
       return Res;
+    if (int Res = cmpNumbers(RMWI->isElementwise(),
+                             cast<AtomicRMWInst>(R)->isElementwise()))
+      return Res;
     if (int Res = cmpOrderings(RMWI->getOrdering(),
                                cast<AtomicRMWInst>(R)->getOrdering()))
       return Res;
diff --git a/llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll b/llvm/test/Transforms/MergeFunc/atomic-elementwise.ll
similarity index 50%
rename from llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll
rename to llvm/test/Transforms/MergeFunc/atomic-elementwise.ll
index 7bb76f61a8a49..9f756c3aaa25b 100644
--- a/llvm/test/Transforms/MergeFunc/atomic-load-elementwise.ll
+++ b/llvm/test/Transforms/MergeFunc/atomic-elementwise.ll
@@ -20,3 +20,23 @@ define internal <2 x i32> @elementwise_load(ptr %p) {
   %value = load atomic elementwise <2 x i32>, ptr %p monotonic, align 4
   ret <2 x i32> %value
 }
+
+define internal <2 x i32> @whole_vector_atomicrmw(ptr %p, <2 x i32> %value) {
+; CHECK-LABEL: define internal <2 x i32> @whole_vector_atomicrmw(
+; CHECK-SAME: ptr [[P:%.*]], <2 x i32> [[VALUE:%.*]]) {
+; CHECK-NEXT:    [[OLD:%.*]] = atomicrmw add ptr [[P]], <2 x i32> [[VALUE]] monotonic, align 8
+; CHECK-NEXT:    ret <2 x i32> [[OLD]]
+;
+  %old = atomicrmw add ptr %p, <2 x i32> %value monotonic, align 8
+  ret <2 x i32> %old
+}
+
+define internal <2 x i32> @elementwise_atomicrmw(ptr %p, <2 x i32> %value) {
+; CHECK-LABEL: define internal <2 x i32> @elementwise_atomicrmw(
+; CHECK-SAME: ptr [[P:%.*]], <2 x i32> [[VALUE:%.*]]) {
+; CHECK-NEXT:    [[OLD:%.*]] = atomicrmw elementwise add ptr [[P]], <2 x i32> [[VALUE]] monotonic, align 8
+; CHECK-NEXT:    ret <2 x i32> [[OLD]]
+;
+  %old = atomicrmw elementwise add ptr %p, <2 x i32> %value monotonic, align 8
+  ret <2 x i32> %old
+}