From f86da4cb7d6433dab10a91e33b4c24e87ec799d8 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Mon, 12 Aug 2024 15:13:15 +0400
Subject: [PATCH] StructurizeCFG: Add SkipUniformRegions pass parameter to new
 PM version (#102812)

Keep respecting the old cl::opt for now.
---
 .../llvm/Transforms/Scalar/StructurizeCFG.h   |  9 ++++++
 llvm/lib/Passes/PassBuilder.cpp               |  5 ++++
 llvm/lib/Passes/PassRegistry.def              |  7 ++++-
 llvm/lib/Transforms/Scalar/StructurizeCFG.cpp | 28 ++++++++++++++++++-
 .../StructurizeCFG/AMDGPU/uniform-regions.ll  |  2 ++
 5 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Scalar/StructurizeCFG.h b/llvm/include/llvm/Transforms/Scalar/StructurizeCFG.h
index 50d41acd529e769..f68067d93545830 100644
--- a/llvm/include/llvm/Transforms/Scalar/StructurizeCFG.h
+++ b/llvm/include/llvm/Transforms/Scalar/StructurizeCFG.h
@@ -13,6 +13,15 @@
 
 namespace llvm {
 struct StructurizeCFGPass : PassInfoMixin<StructurizeCFGPass> {
+private:
+  bool SkipUniformRegions;
+
+public:
+  StructurizeCFGPass(bool SkipUniformRegions = false);
+
+  void printPipeline(raw_ostream &OS,
+                     function_ref<StringRef(StringRef)> MapClassName2PassName);
+
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 } // namespace llvm
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 7bc1c870ce5191a..46f43f3de4705c3 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1085,6 +1085,11 @@ Expected<bool> parseSeparateConstOffsetFromGEPPassOptions(StringRef Params) {
                                             "SeparateConstOffsetFromGEP");
 }
 
+Expected<bool> parseStructurizeCFGPassOptions(StringRef Params) {
+  return PassBuilder::parseSinglePassOption(Params, "skip-uniform-regions",
+                                            "StructurizeCFG");
+}
+
 Expected<OptimizationLevel>
 parseFunctionSimplificationPipelineOptions(StringRef Params) {
   std::optional<OptimizationLevel> L = parseOptLevel(Params);
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 95842d15a35bf60..0cec9fbd7cd05e4 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -458,7 +458,6 @@ FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("slsr", StraightLineStrengthReducePass())
 FUNCTION_PASS("stack-protector", StackProtectorPass(TM))
 FUNCTION_PASS("strip-gc-relocates", StripGCRelocates())
-FUNCTION_PASS("structurizecfg", StructurizeCFGPass())
 FUNCTION_PASS("tailcallelim", TailCallElimPass())
 FUNCTION_PASS("tlshoist", TLSVariableHoistPass())
 FUNCTION_PASS("transform-warning", WarnMissedTransformationsPass())
@@ -586,6 +585,12 @@ FUNCTION_PASS_WITH_PARAMS(
     "sroa", "SROAPass",
     [](SROAOptions PreserveCFG) { return SROAPass(PreserveCFG); },
     parseSROAOptions, "preserve-cfg;modify-cfg")
+FUNCTION_PASS_WITH_PARAMS(
+  "structurizecfg", "StructurizeCFG",
+  [](bool SkipUniformRegions) {
+    return StructurizeCFGPass(SkipUniformRegions);
+  },
+  parseStructurizeCFGPassOptions, "skip-uniform-regions")
 FUNCTION_PASS_WITH_PARAMS(
     "win-eh-prepare", "WinEHPreparePass",
     [](bool DemoteCatchSwitchPHIOnly) {
diff --git a/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp b/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
index 9c711ec183821ff..a6ed58ac9d47f29 100644
--- a/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
+++ b/llvm/lib/Transforms/Scalar/StructurizeCFG.cpp
@@ -1212,20 +1212,46 @@ static void addRegionIntoQueue(Region &R, std::vector<Region *> &Regions) {
     addRegionIntoQueue(*E, Regions);
 }
 
+StructurizeCFGPass::StructurizeCFGPass(bool SkipUniformRegions_)
+    : SkipUniformRegions(SkipUniformRegions_) {
+  if (ForceSkipUniformRegions.getNumOccurrences())
+    SkipUniformRegions = ForceSkipUniformRegions.getValue();
+}
+
+void StructurizeCFGPass::printPipeline(
+    raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
+  static_cast<PassInfoMixin<StructurizeCFGPass> *>(this)->printPipeline(
+      OS, MapClassName2PassName);
+  if (SkipUniformRegions)
+    OS << "<skip-uniform-regions>";
+}
+
 PreservedAnalyses StructurizeCFGPass::run(Function &F,
                                           FunctionAnalysisManager &AM) {
 
   bool Changed = false;
   DominatorTree *DT = &AM.getResult<DominatorTreeAnalysis>(F);
   auto &RI = AM.getResult<RegionInfoAnalysis>(F);
+
+  UniformityInfo *UI = nullptr;
+  if (SkipUniformRegions)
+    UI = &AM.getResult<UniformityInfoAnalysis>(F);
+
   std::vector<Region *> Regions;
   addRegionIntoQueue(*RI.getTopLevelRegion(), Regions);
   while (!Regions.empty()) {
     Region *R = Regions.back();
+    Regions.pop_back();
+
     StructurizeCFG SCFG;
     SCFG.init(R);
+
+    if (SkipUniformRegions && SCFG.makeUniformRegion(R, *UI)) {
+      Changed = true; // May have added metadata.
+      continue;
+    }
+
     Changed |= SCFG.run(R, DT);
-    Regions.pop_back();
   }
   if (!Changed)
     return PreservedAnalyses::all();
diff --git a/llvm/test/Transforms/StructurizeCFG/AMDGPU/uniform-regions.ll b/llvm/test/Transforms/StructurizeCFG/AMDGPU/uniform-regions.ll
index ae73eedd4f502b5..34c73ab8fd74f35 100644
--- a/llvm/test/Transforms/StructurizeCFG/AMDGPU/uniform-regions.ll
+++ b/llvm/test/Transforms/StructurizeCFG/AMDGPU/uniform-regions.ll
@@ -1,5 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -mtriple=amdgcn-- -S -o - -structurizecfg -structurizecfg-skip-uniform-regions -structurizecfg-relaxed-uniform-regions < %s | FileCheck %s
+; RUN: opt -mtriple=amdgcn-- -S -o - -passes='structurizecfg<skip-uniform-regions>' -structurizecfg-relaxed-uniform-regions < %s | FileCheck %s
+; RUN: opt -mtriple=amdgcn-- -S -o - -passes=structurizecfg -structurizecfg-skip-uniform-regions -structurizecfg-relaxed-uniform-regions < %s | FileCheck %s
 
 define amdgpu_cs void @uniform(i32 inreg %v) {
 ; CHECK-LABEL: @uniform(