[SROA][TTI][DirectX] Add support for struct alloca decomposition

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -1977,6 +1977,11 @@ class TargetTransformInfo {
   /// target.
   LLVM_ABI bool allowVectorElementIndexingUsingGEP() const;
 
+  /// \returns True if the target does not support struct allocas and therefore
+  /// requires struct alloca instructions to be scalarized / decomposed into
+  /// its components.
+  LLVM_ABI bool shouldDecomposeStructAllocas() const;
+
 private:
   std::unique_ptr<const TargetTransformInfoImplBase> TTIImpl;
 };

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -1163,6 +1163,8 @@ class TargetTransformInfoImplBase {
 
   virtual bool allowVectorElementIndexingUsingGEP() const { return true; }
 
+  virtual bool shouldDecomposeStructAllocas() const { return false; }
+
 protected:
   // Obtain the minimum required size to hold the value (without the sign)
   // In case of a vector it returns the min required size for one element.

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/Analysis/TargetTransformInfo.cpp

@@ -1506,6 +1506,10 @@ bool TargetTransformInfo::allowVectorElementIndexingUsingGEP() const {
   return TTIImpl->allowVectorElementIndexingUsingGEP();
 }
 
+bool TargetTransformInfo::shouldDecomposeStructAllocas() const {
+  return TTIImpl->shouldDecomposeStructAllocas();
+}
+
 TargetTransformInfoImplBase::~TargetTransformInfoImplBase() = default;
 
 TargetIRAnalysis::TargetIRAnalysis() : TTICallback(&getDefaultTTI) {}

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/Target/DirectX/DirectXTargetMachine.cpp

@@ -48,6 +48,7 @@
 #include "llvm/Target/TargetLoweringObjectFile.h"
 #include "llvm/Transforms/IPO/GlobalDCE.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/SROA.h"
 #include "llvm/Transforms/Scalar/Scalarizer.h"
 #include <optional>
 
@@ -107,6 +108,10 @@ class DirectXPassConfig : public TargetPassConfig {
 
   FunctionPass *createTargetRegisterAllocator(bool) override { return nullptr; }
   void addCodeGenPrepare() override {
+    // Clang does not apply SROA with -O0, but it is required for DXIL. So we
+    // add SROA here when -O0 is given.
+    if (getOptLevel() == CodeGenOptLevel::None)
+      addPass(createSROAPass(/*PreserveCFG=*/true, /*DecomposeStructs=*/true));
     addPass(createDXILFinalizeLinkageLegacyPass());
     addPass(createGlobalDCEPass());
     addPass(createDXILResourceAccessLegacyPass());

llvm/lib/Target/DirectX/DirectXTargetTransformInfo.cpp

@@ -65,3 +65,5 @@ bool DirectXTTIImpl::isTargetIntrinsicTriviallyScalarizable(
     return false;
   }
 }
+
+bool DirectXTTIImpl::shouldDecomposeStructAllocas() const { return true; }

llvm/lib/Target/DirectX/DirectXTargetTransformInfo.h

@@ -39,6 +39,7 @@ class DirectXTTIImpl final : public BasicTTIImplBase<DirectXTTIImpl> {
                                           unsigned ScalarOpdIdx) const override;
   bool isTargetIntrinsicWithOverloadTypeAtArg(Intrinsic::ID ID,
                                               int OpdIdx) const override;
+  bool shouldDecomposeStructAllocas() const override;
 };
 } // namespace llvm