Refine legalization cleanup policy

Narrow legalize-time cleanup further by separating SSA rewrite modes and by avoiding blanket cleanup where the producer does not require it.

With the companion DXC branch this drops the 58k LoC path tracer payload from 4.702 s to 2.464 s on the same machine while full CodeGenSPIRV still passes.

Author: AnastaZIuk

include/spirv-tools/optimizer.hpp

@@ -32,6 +32,13 @@ class Pass;
 struct DescriptorSetAndBinding;
 }  // namespace opt
 
+enum class SSARewriteMode {
+  None,
+  All,
+  OpaqueOnly,
+  SpecialTypes,
+};
+
 // C++ interface for SPIR-V optimization functionalities. It wraps the context
 // (including target environment and the corresponding SPIR-V grammar) and
 // provides methods for registering optimization passes and optimizing.
@@ -127,7 +134,7 @@ class SPIRV_TOOLS_EXPORT Optimizer {
   Optimizer& RegisterLegalizationPasses(bool preserve_interface);
   Optimizer& RegisterLegalizationPasses(bool preserve_interface,
                                         bool include_loop_unroll,
-                                        bool include_ssa_rewrite);
+                                        SSARewriteMode ssa_rewrite_mode);
 
   // Register passes specified in the list of |flags|.  Each flag must be a
   // string of a form accepted by Optimizer::FlagHasValidForm().
@@ -648,11 +655,6 @@ Optimizer::PassToken CreateLoopPeelingPass();
 // Works best after LICM and local multi store elimination pass.
 Optimizer::PassToken CreateLoopUnswitchPass();
 
-// Creates a pass to legalize multidimensional arrays for Vulkan.
-// This pass will replace multidimensional arrays of resources with a single
-// dimensional array. Combine-access-chains should be run before this pass.
-Optimizer::PassToken CreateLegalizeMultidimArrayPass();
-
 // Create global value numbering pass.
 // This pass will look for instructions where the same value is computed on all
 // paths leading to the instruction.  Those instructions are deleted.
@@ -712,7 +714,8 @@ Optimizer::PassToken CreateLoopUnrollPass(bool fully_unroll, int factor = 0);
 // operations on SSA IDs.  This allows SSA optimizers to act on these variables.
 // Only variables that are local to the function and of supported types are
 // processed (see IsSSATargetVar for details).
-Optimizer::PassToken CreateSSARewritePass();
+Optimizer::PassToken CreateSSARewritePass(
+    SSARewriteMode mode = SSARewriteMode::All);
 
 // Create pass to convert relaxed precision instructions to half precision.
 // This pass converts as many relaxed float32 arithmetic operations to half as

source/opt/mem_pass.cpp

@@ -53,7 +53,27 @@ bool MemPass::IsBaseTargetType(const Instruction* typeInst) const {
 }
 
 bool MemPass::IsTargetType(const Instruction* typeInst) const {
-  if (IsBaseTargetType(typeInst)) return true;
+  switch (ssa_rewrite_mode_) {
+    case SSARewriteMode::None:
+      return false;
+    case SSARewriteMode::OpaqueOnly:
+      if (typeInst->IsOpaqueType()) return true;
+      break;
+    case SSARewriteMode::SpecialTypes:
+      if (typeInst->IsOpaqueType()) return true;
+      switch (typeInst->opcode()) {
+        case spv::Op::OpTypePointer:
+        case spv::Op::OpTypeCooperativeMatrixNV:
+        case spv::Op::OpTypeCooperativeMatrixKHR:
+          return true;
+        default:
+          break;
+      }
+      break;
+    case SSARewriteMode::All:
+      if (IsBaseTargetType(typeInst)) return true;
+      break;
+  }
   if (typeInst->opcode() == spv::Op::OpTypeArray) {
     if (!IsTargetType(
             get_def_use_mgr()->GetDef(typeInst->GetSingleWordOperand(1)))) {
@@ -72,8 +92,7 @@ bool MemPass::IsTargetType(const Instruction* typeInst) const {
 
 bool MemPass::IsNonPtrAccessChain(const spv::Op opcode) const {
   return opcode == spv::Op::OpAccessChain ||
-         opcode == spv::Op::OpInBoundsAccessChain ||
-         opcode == spv::Op::OpUntypedAccessChainKHR;
+         opcode == spv::Op::OpInBoundsAccessChain;
 }
 
 bool MemPass::IsPtr(uint32_t ptrId) {
@@ -89,14 +108,11 @@ bool MemPass::IsPtr(uint32_t ptrId) {
     ptrInst = get_def_use_mgr()->GetDef(varId);
   }
   const spv::Op op = ptrInst->opcode();
-  if (op == spv::Op::OpVariable || op == spv::Op::OpUntypedVariableKHR ||
-      IsNonPtrAccessChain(op))
-    return true;
+  if (op == spv::Op::OpVariable || IsNonPtrAccessChain(op)) return true;
   const uint32_t varTypeId = ptrInst->type_id();
   if (varTypeId == 0) return false;
   const Instruction* varTypeInst = get_def_use_mgr()->GetDef(varTypeId);
-  return varTypeInst->opcode() == spv::Op::OpTypePointer ||
-         varTypeInst->opcode() == spv::Op::OpTypeUntypedPointerKHR;
+  return varTypeInst->opcode() == spv::Op::OpTypePointer;
 }
 
 Instruction* MemPass::GetPtr(uint32_t ptrId, uint32_t* varId) {
@@ -106,13 +122,11 @@ Instruction* MemPass::GetPtr(uint32_t ptrId, uint32_t* varId) {
 
   switch (ptrInst->opcode()) {
     case spv::Op::OpVariable:
-    case spv::Op::OpUntypedVariableKHR:
     case spv::Op::OpFunctionParameter:
       varInst = ptrInst;
       break;
     case spv::Op::OpAccessChain:
     case spv::Op::OpInBoundsAccessChain:
-    case spv::Op::OpUntypedAccessChainKHR:
     case spv::Op::OpPtrAccessChain:
     case spv::Op::OpInBoundsPtrAccessChain:
     case spv::Op::OpImageTexelPointer:
@@ -125,8 +139,7 @@ Instruction* MemPass::GetPtr(uint32_t ptrId, uint32_t* varId) {
       break;
   }
 
-  if (varInst->opcode() == spv::Op::OpVariable ||
-      varInst->opcode() == spv::Op::OpUntypedVariableKHR) {
+  if (varInst->opcode() == spv::Op::OpVariable) {
     *varId = varInst->result_id();
   } else {
     *varId = 0;
@@ -241,7 +254,8 @@ void MemPass::DCEInst(Instruction* inst,
   }
 }
 
-MemPass::MemPass() {}
+MemPass::MemPass(SSARewriteMode ssa_rewrite_mode)
+    : ssa_rewrite_mode_(ssa_rewrite_mode) {}
 
 bool MemPass::HasOnlySupportedRefs(uint32_t varId) {
   return get_def_use_mgr()->WhileEachUser(varId, [this](Instruction* user) {

source/opt/mem_pass.h

@@ -25,6 +25,7 @@
 #include <unordered_set>
 #include <utility>
 
+#include "spirv-tools/optimizer.hpp"
 #include "source/opt/basic_block.h"
 #include "source/opt/def_use_manager.h"
 #include "source/opt/dominator_analysis.h"
@@ -68,7 +69,7 @@ class MemPass : public Pass {
   void CollectTargetVars(Function* func);
 
  protected:
-  MemPass();
+  explicit MemPass(SSARewriteMode ssa_rewrite_mode = SSARewriteMode::All);
 
   // Returns true if |typeInst| is a scalar type
   // or a vector or matrix
@@ -133,7 +134,9 @@ class MemPass : public Pass {
   // Cache of verified non-target vars
   std::unordered_set<uint32_t> seen_non_target_vars_;
 
- private:
+private:
+  SSARewriteMode ssa_rewrite_mode_ = SSARewriteMode::All;
+
   // Return true if all uses of |varId| are only through supported reference
   // operations ie. loads and store. Also cache in supported_ref_vars_.
   // TODO(dnovillo): This function is replicated in other passes and it's

source/opt/optimizer.cpp

@@ -122,7 +122,7 @@ Optimizer& Optimizer::RegisterPass(PassToken&& p) {
 // or enable more copy propagation.
 Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface,
                                                  bool include_loop_unroll,
-                                                 bool include_ssa_rewrite) {
+                                                 SSARewriteMode ssa_rewrite_mode) {
   auto& optimizer =
       // Wrap OpKill instructions so all other code can be inlined.
       RegisterPass(CreateWrapOpKillPass())
@@ -132,45 +132,39 @@ Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface,
           .RegisterPass(CreateMergeReturnPass())
           // Make sure uses and definitions are in the same function.
           .RegisterPass(CreateInlineExhaustivePass())
-          // Make private variable function scope
-          .RegisterPass(CreateEliminateDeadFunctionsPass())
-          .RegisterPass(CreatePrivateToLocalPass())
-          // Fix up the storage classes that DXC may have purposely generated
-          // incorrectly.  All functions are inlined, and a lot of dead code has
-          // been removed.
-          .RegisterPass(CreateFixStorageClassPass())
-          // Propagate the value stored to the loads in very simple cases.
-          .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
-          .RegisterPass(CreateLocalSingleStoreElimPass())
-          .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-          // Split up aggregates so they are easier to deal with.
-          .RegisterPass(CreateScalarReplacementPass(0))
-          // Remove loads and stores so everything is in intermediate values.
-          // Takes care of copy propagation of non-members.
-          .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
-          .RegisterPass(CreateLocalSingleStoreElimPass())
-          .RegisterPass(CreateAggressiveDCEPass(preserve_interface));
-  if (include_ssa_rewrite) {
-    optimizer.RegisterPass(CreateLocalMultiStoreElimPass());
+          .RegisterPass(CreateEliminateDeadFunctionsPass());
+  optimizer.RegisterPass(CreatePrivateToLocalPass());
+  // Fix up the storage classes that DXC may have purposely generated
+  // incorrectly.  All functions are inlined, and a lot of dead code has
+  // been removed.
+  optimizer.RegisterPass(CreateFixStorageClassPass());
+  // Propagate the value stored to the loads in very simple cases.
+  optimizer.RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
+      .RegisterPass(CreateLocalSingleStoreElimPass())
+      .RegisterPass(CreateAggressiveDCEPass(preserve_interface));
+  optimizer
+      // Split up aggregates so they are easier to deal with.
+      .RegisterPass(CreateScalarReplacementPass(0));
+  // Remove loads and stores so everything is in intermediate values.
+  // Takes care of copy propagation of non-members.
+  optimizer.RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
+      .RegisterPass(CreateLocalSingleStoreElimPass())
+      .RegisterPass(CreateAggressiveDCEPass(preserve_interface));
+  if (ssa_rewrite_mode != SSARewriteMode::None) {
+    optimizer.RegisterPass(CreateSSARewritePass(ssa_rewrite_mode));
   }
-  optimizer.RegisterPass(CreateCombineAccessChainsPass());
-  if (include_ssa_rewrite) {
-    optimizer.RegisterPass(CreateAggressiveDCEPass(preserve_interface));
-  }
-  optimizer.RegisterPass(CreateLegalizeMultidimArrayPass())
-          // Propagate constants to get as many constant conditions on branches
-          // as possible.
-          .RegisterPass(CreateCCPPass());
+  optimizer
+      // Propagate constants to get as many constant conditions on branches
+      // as possible.
+      .RegisterPass(CreateCCPPass());
   if (include_loop_unroll) {
     optimizer.RegisterPass(CreateLoopUnrollPass(true));
   }
+  optimizer.RegisterPass(CreateDeadBranchElimPass())
+      // Copy propagate members. Cleans up code sequences generated by scalar
+      // replacement. Also important for removing OpPhi nodes.
+      .RegisterPass(CreateSimplificationPass());
   return optimizer
-      .RegisterPass(CreateDeadBranchElimPass())
-      // Copy propagate members.  Cleans up code sequences generated by
-      // scalar replacement.  Also important for removing OpPhi nodes.
-      .RegisterPass(CreateSimplificationPass())
-      .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateCopyPropagateArraysPass())
       // May need loop unrolling here see
       // https://github.com/Microsoft/DirectXShaderCompiler/pull/930
       // Get rid of unused code that contain traces of illegal code
@@ -186,11 +180,12 @@ Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface,
 }
 
 Optimizer& Optimizer::RegisterLegalizationPasses() {
-  return RegisterLegalizationPasses(false, true, true);
+  return RegisterLegalizationPasses(false, true, SSARewriteMode::All);
 }
 
 Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface) {
-  return RegisterLegalizationPasses(preserve_interface, true, true);
+  return RegisterLegalizationPasses(preserve_interface, true,
+                                    SSARewriteMode::All);
 }
 
 Optimizer& Optimizer::RegisterPerformancePasses(bool preserve_interface) {
@@ -199,20 +194,21 @@ Optimizer& Optimizer::RegisterPerformancePasses(bool preserve_interface) {
       .RegisterPass(CreateMergeReturnPass())
       .RegisterPass(CreateInlineExhaustivePass())
       .RegisterPass(CreateEliminateDeadFunctionsPass())
-      .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
       .RegisterPass(CreatePrivateToLocalPass())
       .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
       .RegisterPass(CreateLocalSingleStoreElimPass())
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
       .RegisterPass(CreateScalarReplacementPass(0))
-      .RegisterPass(CreateLocalAccessChainConvertPass())
-      .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
+      .RegisterPass(CreateLocalAccessChainConvertPass());
+  optimizer.RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
       .RegisterPass(CreateLocalSingleStoreElimPass())
-      .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateLocalMultiStoreElimPass())
-      .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateCCPPass())
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface));
+  optimizer.RegisterPass(CreateCCPPass())
+      .RegisterPass(CreateAggressiveDCEPass(preserve_interface));
+  // Preserve LoopControl::Unroll in the IR instead of always materializing
+  // it here. The optimizer-side full unroll is very costly on large modules
+  // with many tiny [unroll]-annotated loops, while the hint remains available
+  // to downstream consumers in the final SPIR-V.
   optimizer.RegisterPass(CreateDeadBranchElimPass());
   optimizer.RegisterPass(CreateLocalRedundancyEliminationPass());
   optimizer.RegisterPass(CreateCombineAccessChainsPass())
@@ -222,7 +218,7 @@ Optimizer& Optimizer::RegisterPerformancePasses(bool preserve_interface) {
       .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
       .RegisterPass(CreateLocalSingleStoreElimPass())
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateSSARewritePass())
+      .RegisterPass(CreateSSARewritePass(SSARewriteMode::SpecialTypes))
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
       .RegisterPass(CreateVectorDCEPass())
       .RegisterPass(CreateDeadInsertElimPass())
@@ -413,8 +409,6 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag,
     RegisterPass(CreateFoldSpecConstantOpAndCompositePass());
   } else if (pass_name == "loop-unswitch") {
     RegisterPass(CreateLoopUnswitchPass());
-  } else if (pass_name == "legalize-multidim-array") {
-    RegisterPass(CreateLegalizeMultidimArrayPass());
   } else if (pass_name == "scalar-replacement") {
     if (pass_args.size() == 0) {
       RegisterPass(CreateScalarReplacementPass(0));
@@ -977,11 +971,6 @@ Optimizer::PassToken CreateLoopUnswitchPass() {
       MakeUnique<opt::LoopUnswitchPass>());
 }
 
-Optimizer::PassToken CreateLegalizeMultidimArrayPass() {
-  return MakeUnique<Optimizer::PassToken::Impl>(
-      MakeUnique<opt::LegalizeMultidimArrayPass>());
-}
-
 Optimizer::PassToken CreateRedundancyEliminationPass() {
   return MakeUnique<Optimizer::PassToken::Impl>(
       MakeUnique<opt::RedundancyEliminationPass>());
@@ -1031,9 +1020,9 @@ Optimizer::PassToken CreateLoopUnrollPass(bool fully_unroll, int factor) {
       MakeUnique<opt::LoopUnroller>(fully_unroll, factor));
 }
 
-Optimizer::PassToken CreateSSARewritePass() {
+Optimizer::PassToken CreateSSARewritePass(SSARewriteMode mode) {
   return MakeUnique<Optimizer::PassToken::Impl>(
-      MakeUnique<opt::SSARewritePass>());
+      MakeUnique<opt::SSARewritePass>(mode));
 }
 
 Optimizer::PassToken CreateCopyPropagateArraysPass() {

source/opt/ssa_rewrite_pass.h

@@ -294,7 +294,8 @@ class SSARewriter {
 
 class SSARewritePass : public MemPass {
  public:
-  SSARewritePass() = default;
+  explicit SSARewritePass(SSARewriteMode mode = SSARewriteMode::All)
+      : MemPass(mode) {}
 
   const char* name() const override { return "ssa-rewrite"; }
   Status Process() override;