Narrow blanket SPIR-V legalization work in optimizer recipes

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.
@@ -125,6 +132,9 @@ class SPIRV_TOOLS_EXPORT Optimizer {
   // interface are considered live and are not eliminated.
   Optimizer& RegisterLegalizationPasses();
   Optimizer& RegisterLegalizationPasses(bool preserve_interface);
+  Optimizer& RegisterLegalizationPasses(bool preserve_interface,
+                                        bool include_loop_unroll,
+                                        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().
@@ -645,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.
@@ -709,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

@@ -120,8 +120,10 @@ Optimizer& Optimizer::RegisterPass(PassToken&& p) {
 // The legalization problem is essentially a very general copy propagation
 // problem.  The optimization we use are all used to either do copy propagation
 // or enable more copy propagation.
-Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface) {
-  return
+Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface,
+                                                 bool include_loop_unroll,
+                                                 SSARewriteMode ssa_rewrite_mode) {
+  auto& optimizer =
       // Wrap OpKill instructions so all other code can be inlined.
       RegisterPass(CreateWrapOpKillPass())
           // Remove unreachable block so that merge return works.
@@ -130,87 +132,93 @@ 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))
-          .RegisterPass(CreateLocalMultiStoreElimPass())
-          .RegisterPass(CreateCombineAccessChainsPass())
-          .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-          .RegisterPass(CreateLegalizeMultidimArrayPass())
-          // Propagate constants to get as many constant conditions on branches
-          // as possible.
-          .RegisterPass(CreateCCPPass())
-          .RegisterPass(CreateLoopUnrollPass(true))
-          .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
-          // or unused references to unbound external objects
-          .RegisterPass(CreateVectorDCEPass())
-          .RegisterPass(CreateDeadInsertElimPass())
-          .RegisterPass(CreateReduceLoadSizePass())
-          .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-          .RegisterPass(CreateRemoveUnusedInterfaceVariablesPass())
-          .RegisterPass(CreateInterpolateFixupPass())
-          .RegisterPass(CreateInvocationInterlockPlacementPass())
-          .RegisterPass(CreateOpExtInstWithForwardReferenceFixupPass());
+          .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
+      // 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
+      // May need loop unrolling here see
+      // https://github.com/Microsoft/DirectXShaderCompiler/pull/930
+      // Get rid of unused code that contain traces of illegal code
+      // or unused references to unbound external objects
+      .RegisterPass(CreateVectorDCEPass())
+      .RegisterPass(CreateDeadInsertElimPass())
+      .RegisterPass(CreateReduceLoadSizePass())
+      .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
+      .RegisterPass(CreateRemoveUnusedInterfaceVariablesPass())
+      .RegisterPass(CreateInterpolateFixupPass())
+      .RegisterPass(CreateInvocationInterlockPlacementPass())
+      .RegisterPass(CreateOpExtInstWithForwardReferenceFixupPass());
 }
 
 Optimizer& Optimizer::RegisterLegalizationPasses() {
-  return RegisterLegalizationPasses(false);
+  return RegisterLegalizationPasses(false, true, SSARewriteMode::All);
+}
+
+Optimizer& Optimizer::RegisterLegalizationPasses(bool preserve_interface) {
+  return RegisterLegalizationPasses(preserve_interface, true,
+                                    SSARewriteMode::All);
 }
 
 Optimizer& Optimizer::RegisterPerformancePasses(bool preserve_interface) {
-  return RegisterPass(CreateWrapOpKillPass())
+  auto& optimizer = RegisterPass(CreateWrapOpKillPass())
       .RegisterPass(CreateDeadBranchElimPass())
       .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))
-      .RegisterPass(CreateLoopUnrollPass(true))
-      .RegisterPass(CreateDeadBranchElimPass())
-      .RegisterPass(CreateRedundancyEliminationPass())
-      .RegisterPass(CreateCombineAccessChainsPass())
+      .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())
       .RegisterPass(CreateSimplificationPass())
       .RegisterPass(CreateScalarReplacementPass(0))
       .RegisterPass(CreateLocalAccessChainConvertPass())
       .RegisterPass(CreateLocalSingleBlockLoadStoreElimPass())
       .RegisterPass(CreateLocalSingleStoreElimPass())
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateSSARewritePass())
+      .RegisterPass(CreateSSARewritePass(SSARewriteMode::SpecialTypes))
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
       .RegisterPass(CreateVectorDCEPass())
       .RegisterPass(CreateDeadInsertElimPass())
@@ -220,9 +228,9 @@ Optimizer& Optimizer::RegisterPerformancePasses(bool preserve_interface) {
       .RegisterPass(CreateCopyPropagateArraysPass())
       .RegisterPass(CreateReduceLoadSizePass())
       .RegisterPass(CreateAggressiveDCEPass(preserve_interface))
-      .RegisterPass(CreateBlockMergePass())
-      .RegisterPass(CreateRedundancyEliminationPass())
-      .RegisterPass(CreateDeadBranchElimPass())
+      .RegisterPass(CreateBlockMergePass());
+  optimizer.RegisterPass(CreateLocalRedundancyEliminationPass());
+  return optimizer.RegisterPass(CreateDeadBranchElimPass())
       .RegisterPass(CreateBlockMergePass())
       .RegisterPass(CreateSimplificationPass());
 }
@@ -401,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));
@@ -965,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>());
@@ -1019,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;