From aea2c115453cce45a66ad79ccadde0aad181a068 Mon Sep 17 00:00:00 2001
From: XiaShark <xiajingze1@huawei.com>
Date: Sat, 11 Oct 2025 16:34:27 +0800
Subject: [PATCH] [CodeGen] Add DebugCounter support for machine-cse

---
 llvm/lib/CodeGen/MachineCSE.cpp               |  9 +++++++++
 .../AArch64/debugcounter-machine-cse.mir      | 20 +++++++++++++++++++
 2 files changed, 29 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/debugcounter-machine-cse.mir

diff --git a/llvm/lib/CodeGen/MachineCSE.cpp b/llvm/lib/CodeGen/MachineCSE.cpp
index 780ed92bc8581..9e67d830baf27 100644
--- a/llvm/lib/CodeGen/MachineCSE.cpp
+++ b/llvm/lib/CodeGen/MachineCSE.cpp
@@ -40,6 +40,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugCounter.h"
 #include "llvm/Support/RecyclingAllocator.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
@@ -60,6 +61,9 @@ STATISTIC(NumCrossBBCSEs,
           "Number of cross-MBB physreg referencing CS eliminated");
 STATISTIC(NumCommutes,  "Number of copies coalesced after commuting");
 
+DEBUG_COUNTER(MachineCSECounter, "machine-cse",
+              "Controls which instructions are removed");
+
 // Threshold to avoid excessive cost to compute isProfitableToCSE.
 static cl::opt<int>
     CSUsesThreshold("csuses-threshold", cl::Hidden, cl::init(1024),
@@ -668,6 +672,11 @@ bool MachineCSEImpl::ProcessBlockCSE(MachineBasicBlock *MBB) {
       --NumDefs;
     }
 
+    if (DoCSE && !DebugCounter::shouldExecute(MachineCSECounter)) {
+      LLVM_DEBUG(dbgs() << "Skip CSE due to debug counter\n");
+      DoCSE = false;
+    }
+
     // Actually perform the elimination.
     if (DoCSE) {
       for (const std::pair<Register, Register> &CSEPair : CSEPairs) {
diff --git a/llvm/test/CodeGen/AArch64/debugcounter-machine-cse.mir b/llvm/test/CodeGen/AArch64/debugcounter-machine-cse.mir
new file mode 100644
index 0000000000000..89f2fa15be158
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/debugcounter-machine-cse.mir
@@ -0,0 +1,20 @@
+# RUN: llc -mtriple=aarch64 -run-pass=machine-cse -debug-counter=machine-cse=0-1 -o - %s | FileCheck %s
+
+---
+name:             debug_counter_machine_cse
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: debug_counter_machine_cse
+    ; CHECK: %0:gpr32 = MOVi32imm 1
+    ; CHECK-NEXT: %1:gpr32 = ADDWrr %0, %0
+    ; CHECK-NEXT: %4:gpr32 = ADDWrr %0, %0
+    ; CHECK-NEXT: $w0 = COPY %4
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+    %0:gpr32 = MOVi32imm 1
+    %1:gpr32 = ADDWrr %0, %0
+    %2:gpr32 = ADDWrr %0, %0
+    %3:gpr32 = ADDWrr %0, %0
+    %4:gpr32 = ADDWrr %0, %0
+    $w0 = COPY %4
+    RET_ReallyLR implicit $w0
+...