[llvm] [AMDGPU] Add register pressure guard on LLVM-IR level to prevent harmful optimizations (PR #171267)

[Matt Arsenault via llvm-commits]

================
@@ -0,0 +1,322 @@
+//===- AMDGPURegPressureGuard.cpp - Register Pressure Guarded Pass Wrapper ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements a guard mechanism for IR transformations that measures
+/// VGPR register pressure before and after applying a pass, reverting the
+/// transformation if pressure increases beyond a configurable threshold.
+///
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "AMDGPURegPressureGuard.h"
+#include "AMDGPURegPressureEstimator.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/UniformityAnalysis.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/Function.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include <memory>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-reg-pressure-guard"
+
+static cl::opt<bool> EnableRegPressureGuard(
+    "amdgpu-enable-reg-pressure-guard",
+    cl::desc("Enable AMDGPU register pressure guard to revert transformations "
+             "that increase VGPR pressure beyond threshold"),
+    cl::init(true));
+
+static cl::opt<unsigned> MaxPercentIncreaseOpt(
+    "amdgpu-reg-pressure-max-increase",
+    cl::desc("Maximum allowed percentage increase in VGPR pressure"),
+    cl::init(20));
+
+static cl::opt<unsigned> MinBaselineVGPRsOpt(
+    "amdgpu-reg-pressure-min-baseline",
+    cl::desc("Minimum baseline VGPRs required to enable guard"),
+    cl::init(96));
+
+STATISTIC(NumFunctionsGuarded, "Number of functions checked by guard");
+STATISTIC(NumTransformationsReverted, "Number of transformations reverted");
+STATISTIC(NumTransformationsKept, "Number of transformations kept");
+
+namespace llvm {
+namespace AMDGPURegPressureGuardHelper {
+
+bool isEnabled() { return EnableRegPressureGuard; }
+
+bool shouldGuardFunction(const AMDGPURegPressureGuardConfig &Config,
+                         Function &F, unsigned BaselineVGPRs) {
+  unsigned MinBaseline = Config.MinBaselineVGPRs > 0
+                             ? Config.MinBaselineVGPRs
+                             : MinBaselineVGPRsOpt;
+
+  if (BaselineVGPRs < MinBaseline) {
+    LLVM_DEBUG(dbgs() << "AMDGPURegPressureGuard: Skipping " << F.getName()
+                      << " - baseline VGPRs (" << BaselineVGPRs
+                      << ") below threshold (" << MinBaseline
+                      << ")\n");
+    return false;
+  }
+  return true;
+}
+
+bool shouldRevert(const AMDGPURegPressureGuardConfig &Config,
+                  unsigned BaselineVGPRs, unsigned NewVGPRs) {
+  if (NewVGPRs <= BaselineVGPRs) {
+    LLVM_DEBUG(dbgs() << "AMDGPURegPressureGuard: Keeping transformation - "
+                      << "VGPRs decreased from " << BaselineVGPRs << " to "
+                      << NewVGPRs << "\n");
+    return false;
+  }
+
+  if (BaselineVGPRs == 0)
+    return false;
+
+  unsigned PercentIncrease = ((NewVGPRs - BaselineVGPRs) * 100) / BaselineVGPRs;
+  unsigned MaxIncrease = Config.MaxPercentIncrease > 0
+                             ? Config.MaxPercentIncrease
+                             : MaxPercentIncreaseOpt;
+
+  if (PercentIncrease > MaxIncrease) {
+    LLVM_DEBUG(dbgs() << "AMDGPURegPressureGuard: Reverting transformation - "
+                      << "VGPR increase " << PercentIncrease
+                      << "% exceeds limit " << MaxIncrease
+                      << "%\n");
+    return true;
+  }
+
+  LLVM_DEBUG(dbgs() << "AMDGPURegPressureGuard: Keeping transformation - "
+                    << "VGPR increase " << PercentIncrease
+                    << "% within limit\n");
+  return false;
+}
+
+void restoreFunction(Function &F, Function &BackupFunc) {
----------------
arsenm wrote:

A function pass is not allowed to do any of this 

https://github.com/llvm/llvm-project/pull/171267