Add decompose-structs option to SROA

Author: Icohedron

llvm/include/llvm/Transforms/Scalar.h

@@ -44,7 +44,8 @@ LLVM_ABI FunctionPass *createDeadStoreEliminationPass();
 //
 // SROA - Replace aggregates or pieces of aggregates with scalar SSA values.
 //
-LLVM_ABI FunctionPass *createSROAPass(bool PreserveCFG = true);
+LLVM_ABI FunctionPass *createSROAPass(bool PreserveCFG = true,
+                                      bool DecomposeStructs = false);
 
 //===----------------------------------------------------------------------===//
 //

llvm/include/llvm/Transforms/Scalar/SROA.h

@@ -21,15 +21,31 @@ namespace llvm {
 
 class Function;
 
-enum class SROAOptions : bool { ModifyCFG, PreserveCFG };
+struct SROAOptions {
+  enum PreserveCFGOption : bool { ModifyCFG, PreserveCFG };
+  enum DecomposeStructsOption : bool { NoDecomposeStructs, DecomposeStructs };
+  PreserveCFGOption PCFGOption;
+  DecomposeStructsOption DSOption;
+  SROAOptions(PreserveCFGOption PCFGOption)
+      : PCFGOption(PCFGOption), DSOption(NoDecomposeStructs) {}
+  SROAOptions(PreserveCFGOption PCFGOption, DecomposeStructsOption DSOption)
+      : PCFGOption(PCFGOption), DSOption(DSOption) {}
+};
 
 class SROAPass : public PassInfoMixin<SROAPass> {
-  const SROAOptions PreserveCFG;
+  const SROAOptions Options;
 
 public:
   /// If \p PreserveCFG is set, then the pass is not allowed to modify CFG
   /// in any way, even if it would update CFG analyses.
-  SROAPass(SROAOptions PreserveCFG);
+  SROAPass(SROAOptions::PreserveCFGOption PreserveCFG);
+
+  /// If \p Options.PreserveCFG is set, then the pass is not allowed to modify
+  /// CFG in any way, even if it would update CFG analyses.
+  /// If \p Options.DecomposeStructs is set, then the pass will decompose
+  /// structs allocas into its constituent components regardless of whether or
+  /// not pointer offsets into them are known at compile time.
+  SROAPass(const SROAOptions &Options);
 
   /// Run the pass over the function.
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

llvm/lib/Passes/PassBuilder.cpp

@@ -1353,16 +1353,29 @@ Expected<ScalarizerPassOptions> parseScalarizerOptions(StringRef Params) {
 }
 
 Expected<SROAOptions> parseSROAOptions(StringRef Params) {
-  if (Params.empty() || Params == "modify-cfg")
-    return SROAOptions::ModifyCFG;
-  if (Params == "preserve-cfg")
-    return SROAOptions::PreserveCFG;
-  return make_error<StringError>(
-      formatv("invalid SROA pass parameter '{}' (either preserve-cfg or "
-              "modify-cfg can be specified)",
-              Params)
-          .str(),
-      inconvertibleErrorCode());
+  SROAOptions::PreserveCFGOption PreserveCFG = SROAOptions::ModifyCFG;
+  SROAOptions::DecomposeStructsOption DecomposeStructs =
+      SROAOptions::NoDecomposeStructs;
+
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
+
+    if (ParamName.consume_front("preserve-cfg"))
+      PreserveCFG = SROAOptions::PreserveCFG;
+    else if (ParamName.consume_front("modify-cfg"))
+      PreserveCFG = SROAOptions::ModifyCFG;
+    else if (ParamName.consume_front("no-decompose-structs"))
+      DecomposeStructs = SROAOptions::NoDecomposeStructs;
+    else if (ParamName.consume_front("decompose-structs"))
+      DecomposeStructs = SROAOptions::DecomposeStructs;
+    else
+      return make_error<StringError>(
+          formatv("invalid SROA pass option '{}'", ParamName).str(),
+          inconvertibleErrorCode());
+  }
+
+  return SROAOptions(PreserveCFG, DecomposeStructs);
 }
 
 Expected<StackLifetime::LivenessType>

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -174,6 +174,7 @@ class SROA {
   DomTreeUpdater *const DTU;
   AssumptionCache *const AC;
   const bool PreserveCFG;
+  const bool DecomposeStructs;
 
   /// Worklist of alloca instructions to simplify.
   ///
@@ -236,9 +237,10 @@ class SROA {
 
 public:
   SROA(LLVMContext *C, DomTreeUpdater *DTU, AssumptionCache *AC,
-       SROAOptions PreserveCFG_)
+       const SROAOptions &Options)
       : C(C), DTU(DTU), AC(AC),
-        PreserveCFG(PreserveCFG_ == SROAOptions::PreserveCFG) {}
+        PreserveCFG(Options.PCFGOption == SROAOptions::PreserveCFG),
+        DecomposeStructs(Options.DSOption == SROAOptions::DecomposeStructs) {}
 
   /// Main run method used by both the SROAPass and by the legacy pass.
   std::pair<bool /*Changed*/, bool /*CFGChanged*/> runSROA(Function &F);
@@ -6040,7 +6042,7 @@ PreservedAnalyses SROAPass::run(Function &F, FunctionAnalysisManager &AM) {
   AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   auto [Changed, CFGChanged] =
-      SROA(&F.getContext(), &DTU, &AC, PreserveCFG).runSROA(F);
+      SROA(&F.getContext(), &DTU, &AC, Options).runSROA(F);
   if (!Changed)
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
@@ -6054,23 +6056,29 @@ void SROAPass::printPipeline(
     raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
   static_cast<PassInfoMixin<SROAPass> *>(this)->printPipeline(
       OS, MapClassName2PassName);
-  OS << (PreserveCFG == SROAOptions::PreserveCFG ? "<preserve-cfg>"
-                                                 : "<modify-cfg>");
+  OS << (Options.PCFGOption == SROAOptions::PreserveCFG
+             ? "<preserve-cfg>"
+             : "<modify-cfg>");
 }
 
-SROAPass::SROAPass(SROAOptions PreserveCFG) : PreserveCFG(PreserveCFG) {}
+SROAPass::SROAPass(SROAOptions::PreserveCFGOption PreserveCFG)
+    : Options({PreserveCFG, SROAOptions::NoDecomposeStructs}) {}
+SROAPass::SROAPass(const SROAOptions &Options) : Options(Options) {}
 
 namespace {
 
 /// A legacy pass for the legacy pass manager that wraps the \c SROA pass.
 class SROALegacyPass : public FunctionPass {
-  SROAOptions PreserveCFG;
+  SROAOptions Options;
 
 public:
   static char ID;
 
-  SROALegacyPass(SROAOptions PreserveCFG = SROAOptions::PreserveCFG)
-      : FunctionPass(ID), PreserveCFG(PreserveCFG) {
+  SROALegacyPass(
+      const SROAOptions &Options =
+          {SROAOptions::PreserveCFGOption::PreserveCFG,
+           SROAOptions::DecomposeStructsOption::NoDecomposeStructs})
+      : FunctionPass(ID), Options(Options) {
     initializeSROALegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
@@ -6083,7 +6091,7 @@ class SROALegacyPass : public FunctionPass {
         getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
     DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
     auto [Changed, _] =
-        SROA(&F.getContext(), &DTU, &AC, PreserveCFG).runSROA(F);
+        SROA(&F.getContext(), &DTU, &AC, Options).runSROA(F);
     return Changed;
   }
 
@@ -6101,9 +6109,11 @@ class SROALegacyPass : public FunctionPass {
 
 char SROALegacyPass::ID = 0;
 
-FunctionPass *llvm::createSROAPass(bool PreserveCFG) {
-  return new SROALegacyPass(PreserveCFG ? SROAOptions::PreserveCFG
-                                        : SROAOptions::ModifyCFG);
+FunctionPass *llvm::createSROAPass(bool PreserveCFG, bool DecomposeStructs) {
+  return new SROALegacyPass(
+      {PreserveCFG ? SROAOptions::PreserveCFG : SROAOptions::ModifyCFG,
+       DecomposeStructs ? SROAOptions::DecomposeStructs
+                        : SROAOptions::NoDecomposeStructs});
 }
 
 INITIALIZE_PASS_BEGIN(SROALegacyPass, "sroa",