[SROA] Canonicalize homogeneous structs to fixed vectors (opt-in, after memcpyopt)

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 AggregateToVector = false);
 
 //===----------------------------------------------------------------------===//
 //

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

@@ -21,15 +21,25 @@ namespace llvm {
 
 class Function;
 
-enum class SROAOptions : bool { ModifyCFG, PreserveCFG };
+struct SROAOptions {
+  enum CFGOption { ModifyCFG, PreserveCFG };
+
+  CFGOption CFG;
+  bool AggregateToVector;
+
+  SROAOptions(CFGOption CFG = PreserveCFG, bool AggregateToVector = false)
+      : CFG(CFG), AggregateToVector(AggregateToVector) {}
+};
 
 class SROAPass : public OptionalPassInfoMixin<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.
-  LLVM_ABI SROAPass(SROAOptions PreserveCFG);
+  /// If \p AggregateToVector is set, then the pass will try to convert
+  /// allocas of homogeneous structs into vector allocas.
+  LLVM_ABI SROAPass(SROAOptions Options);
 
   /// Run the pass over the function.
   LLVM_ABI PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

llvm/lib/Passes/PassBuilder.cpp

@@ -1433,16 +1433,36 @@ 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 Result(SROAOptions::ModifyCFG);
+  bool SawCFGOption = false;
+  while (!Params.empty()) {
+    StringRef ParamName;
+    std::tie(ParamName, Params) = Params.split(';');
+
+    if (ParamName == "modify-cfg") {
+      if (SawCFGOption)
+        return make_error<StringError>("multiple SROA CFG options specified",
+                                       inconvertibleErrorCode());
+      Result.CFG = SROAOptions::ModifyCFG;
+      SawCFGOption = true;
+    } else if (ParamName == "preserve-cfg") {
+      if (SawCFGOption)
+        return make_error<StringError>("multiple SROA CFG options specified",
+                                       inconvertibleErrorCode());
+      Result.CFG = SROAOptions::PreserveCFG;
+      SawCFGOption = true;
+    } else if (ParamName == "aggregate-to-vector") {
+      Result.AggregateToVector = true;
+    } else {
+      return make_error<StringError>(
+          formatv("invalid SROA pass parameter '{}' (expected preserve-cfg, "
+                  "modify-cfg, or aggregate-to-vector)",
+                  ParamName)
+              .str(),
+          inconvertibleErrorCode());
+    }
+  }
+  return Result;
 }
 
 Expected<StackLifetime::LivenessType>

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -1372,7 +1372,16 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     // NOTE: we are very late in the pipeline, and we don't have any LICM
     // or SimplifyCFG passes scheduled after us, that would cleanup
     // the CFG mess this may created if allowed to modify CFG, so forbid that.
-    FPM.addPass(SROAPass(SROAOptions::PreserveCFG));
+
+    // We also turn on struct to vector canonicalization here, which allows
+    // converting allocas of homogeneous structs into vector allocas when the
+    // allocas' users are all memory intrinsics. This allows promotion in some
+    // cases because structs cannot promote to SSA values, but vectors can. We
+    // only turn this on after memcpyopt runs because this might hinder
+    // memcpyopt's optimizations if done before. Look at the documentation for
+    // `tryCanonicalizeStructToVector` in SROA.cpp to see why.
+    FPM.addPass(SROAPass(SROAOptions(SROAOptions::PreserveCFG,
+                                     /*AggregateToVector=*/true)));
   }
 
   if (!isFullLTOPostLink(LTOPhase)) {
@@ -1464,7 +1473,16 @@ void PassBuilder::addVectorPasses(OptimizationLevel Level,
     // NOTE: we are very late in the pipeline, and we don't have any LICM
     // or SimplifyCFG passes scheduled after us, that would cleanup
     // the CFG mess this may created if allowed to modify CFG, so forbid that.
-    FPM.addPass(SROAPass(SROAOptions::PreserveCFG));
+
+    // We also turn on struct to vector canonicalization here, which allows
+    // converting allocas of homogeneous structs into vector allocas when the
+    // allocas' users are all memory intrinsics. This allows promotion in some
+    // cases because structs cannot promote to SSA values, but vectors can. We
+    // only turn this on after memcpyopt runs because this might hinder
+    // memcpyopt's optimizations if done before. Look at the documentation for
+    // `tryCanonicalizeStructToVector` in SROA.cpp to see why.
+    FPM.addPass(SROAPass(SROAOptions(SROAOptions::PreserveCFG,
+                                     /*AggregateToVector=*/true)));
   }
 
   FPM.addPass(InferAlignmentPass());

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -300,7 +300,8 @@ void NVPTXPassConfig::addEarlyCSEOrGVNPass() {
 void NVPTXPassConfig::addAddressSpaceInferencePasses() {
   // NVPTXLowerArgs emits alloca for byval parameters which can often
   // be eliminated by SROA.
-  addPass(createSROAPass());
+  addPass(createSROAPass(/*PreserveCFG=*/true,
+                         /*AggregateToVector=*/true));
   addPass(createNVPTXLowerAllocaPass());
   // TODO: Consider running InferAddressSpaces during opt, earlier in the
   // compilation flow.
@@ -391,7 +392,8 @@ void NVPTXPassConfig::addIRPasses() {
     addEarlyCSEOrGVNPass();
     if (!DisableLoadStoreVectorizer)
       addPass(createLoadStoreVectorizerPass());
-    addPass(createSROAPass());
+    addPass(createSROAPass(/*PreserveCFG=*/true,
+                           /*AggregateToVector=*/true));
     addPass(createNVPTXTagInvariantLoadsPass());
     if (!DisableNVPTXIRPeephole)
       addPass(createNVPTXIRPeepholePass());