Narrow blanket SPIR-V legalization work in optimizer recipes

BUILD.gn

@@ -630,6 +630,8 @@ static_library("spvtools_opt") {
     "source/opt/tree_iterator.h",
     "source/opt/trim_capabilities_pass.cpp",
     "source/opt/trim_capabilities_pass.h",
+    "source/opt/trim_variable_pointers_capabilities_pass.cpp",
+    "source/opt/trim_variable_pointers_capabilities_pass.h",
     "source/opt/type_manager.cpp",
     "source/opt/type_manager.h",
     "source/opt/types.cpp",

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.
@@ -102,6 +109,8 @@ class SPIRV_TOOLS_EXPORT Optimizer {
   // interface are considered live and are not eliminated.
   Optimizer& RegisterPerformancePasses();
   Optimizer& RegisterPerformancePasses(bool preserve_interface);
+  Optimizer& RegisterPerformancePassesFastCompile();
+  Optimizer& RegisterPerformancePassesFastCompile(bool preserve_interface);
 
   // Registers passes that attempt to improve the size of generated code.
   // This sequence of passes is subject to constant review and will change
@@ -125,6 +134,10 @@ class SPIRV_TOOLS_EXPORT Optimizer {
   // interface are considered live and are not eliminated.
   Optimizer& RegisterLegalizationPasses();
   Optimizer& RegisterLegalizationPasses(bool preserve_interface);
+  Optimizer& RegisterLegalizationPassesFastCompile();
+  Optimizer& RegisterLegalizationPassesFastCompile(
+      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().
@@ -710,6 +723,7 @@ Optimizer::PassToken CreateLoopUnrollPass(bool fully_unroll, int factor = 0);
 // 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);
 
 // Create pass to convert relaxed precision instructions to half precision.
 // This pass converts as many relaxed float32 arithmetic operations to half as
@@ -949,6 +963,12 @@ Optimizer::PassToken CreateFixFuncCallArgumentsPass();
 // the unknown capability interacts with one of the trimmed capabilities.
 Optimizer::PassToken CreateTrimCapabilitiesPass();
 
+// Creates a pass that trims unused VariablePointers capabilities.
+// This pass is intended for targeted call-sites that need to remove stale
+// VariablePointers / VariablePointersStorageBuffer declarations left after
+// optimization when the final module no longer requires them.
+Optimizer::PassToken CreateTrimVariablePointersCapabilitiesPass();
+
 // Creates a struct-packing pass.
 // This pass re-assigns all offset layout decorators to tightly pack
 // the struct with OpName matching `structToPack` according to the given packing

source/opt/CMakeLists.txt

@@ -128,6 +128,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   switch_descriptorset_pass.h
   tree_iterator.h
   trim_capabilities_pass.h
+  trim_variable_pointers_capabilities_pass.h
   type_manager.h
   types.h
   unify_const_pass.h
@@ -249,6 +250,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   struct_packing_pass.cpp
   switch_descriptorset_pass.cpp
   trim_capabilities_pass.cpp
+  trim_variable_pointers_capabilities_pass.cpp
   type_manager.cpp
   types.cpp
   unify_const_pass.cpp

source/opt/dead_variable_elimination.cpp

@@ -14,6 +14,7 @@
 
 #include "source/opt/dead_variable_elimination.h"
 
+#include <unordered_set>
 #include <vector>
 
 #include "source/opt/ir_context.h"
@@ -77,9 +78,43 @@ Pass::Status DeadVariableElimination::Process() {
       DeleteVariable(result_id);
     }
   }
+
+  ids_to_remove.clear();
+  for (auto& function : *get_module()) {
+    if (function.IsDeclaration()) continue;
+
+    auto& entry = *function.begin();
+    for (auto inst = entry.begin(); inst != entry.end(); ++inst) {
+      if (inst->opcode() != spv::Op::OpVariable) break;
+      if (!IsFunctionLocalVariable(&*inst)) continue;
+      if (IsLiveVar(inst->result_id())) continue;
+      ids_to_remove.push_back(inst->result_id());
+    }
+  }
+
+  if (!ids_to_remove.empty()) {
+    modified = true;
+    for (auto result_id : ids_to_remove) {
+      DeleteLocalVariable(result_id);
+    }
+  }
+
   return (modified ? Status::SuccessWithChange : Status::SuccessWithoutChange);
 }
 
+bool DeadVariableElimination::IsFunctionLocalVariable(
+    const Instruction* inst) const {
+  if (inst->opcode() != spv::Op::OpVariable) return false;
+
+  const Instruction* type_inst = get_def_use_mgr()->GetDef(inst->type_id());
+  if (type_inst == nullptr || type_inst->opcode() != spv::Op::OpTypePointer) {
+    return false;
+  }
+
+  return spv::StorageClass(type_inst->GetSingleWordInOperand(0)) ==
+         spv::StorageClass::Function;
+}
+
 void DeadVariableElimination::DeleteVariable(uint32_t result_id) {
   Instruction* inst = get_def_use_mgr()->GetDef(result_id);
   assert(inst->opcode() == spv::Op::OpVariable &&
@@ -108,5 +143,19 @@ void DeadVariableElimination::DeleteVariable(uint32_t result_id) {
   }
   context()->KillDef(result_id);
 }
+
+void DeadVariableElimination::DeleteLocalVariable(uint32_t result_id) {
+  std::queue<Instruction*> dead_stores;
+  std::unordered_set<Instruction*> processed;
+  AddStores(result_id, &dead_stores);
+  while (!dead_stores.empty()) {
+    Instruction* inst = dead_stores.front();
+    dead_stores.pop();
+    if (!processed.insert(inst).second) continue;
+    DCEInst(inst, nullptr);
+  }
+
+  context()->KillDef(result_id);
+}
 }  // namespace opt
 }  // namespace spvtools

source/opt/dead_variable_elimination.h

@@ -37,6 +37,8 @@ class DeadVariableElimination : public MemPass {
  private:
   // Deletes the OpVariable instruction who result id is |result_id|.
   void DeleteVariable(uint32_t result_id);
+  void DeleteLocalVariable(uint32_t result_id);
+  bool IsFunctionLocalVariable(const Instruction* inst) const;
 
   // Keeps track of the number of references of an id.  Once that value is 0, it
   // is safe to remove the corresponding instruction.

source/opt/local_single_store_elim_pass.cpp

@@ -308,6 +308,18 @@ bool LocalSingleStoreElimPass::RewriteLoads(
   else
     stored_id = store_inst->GetSingleWordInOperand(kVariableInitIdInIdx);
 
+  const auto get_image_pointer_id = [this](uint32_t value_id) {
+    Instruction* value_inst = context()->get_def_use_mgr()->GetDef(value_id);
+    while (value_inst && value_inst->opcode() == spv::Op::OpCopyObject) {
+      value_id = value_inst->GetSingleWordInOperand(0);
+      value_inst = context()->get_def_use_mgr()->GetDef(value_id);
+    }
+    if (!value_inst || value_inst->opcode() != spv::Op::OpLoad) {
+      return uint32_t{0};
+    }
+    return value_inst->GetSingleWordInOperand(0);
+  };
+
   *all_rewritten = true;
   bool modified = false;
   for (Instruction* use : uses) {
@@ -322,6 +334,17 @@ bool LocalSingleStoreElimPass::RewriteLoads(
       context()->KillNamesAndDecorates(use->result_id());
       context()->ReplaceAllUsesWith(use->result_id(), stored_id);
       context()->KillInst(use);
+    } else if (use->opcode() == spv::Op::OpImageTexelPointer &&
+               dominator_analysis->Dominates(store_inst, use)) {
+      const uint32_t image_ptr_id = get_image_pointer_id(stored_id);
+      if (image_ptr_id == 0) {
+        *all_rewritten = false;
+        continue;
+      }
+      modified = true;
+      context()->ForgetUses(use);
+      use->SetInOperand(0, {image_ptr_id});
+      context()->AnalyzeUses(use);
     } else {
       *all_rewritten = false;
     }

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:
+      switch (typeInst->opcode()) {
+        case spv::Op::OpTypePointer:
+        case spv::Op::OpTypeUntypedPointerKHR:
+        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)))) {
@@ -198,7 +218,7 @@ bool MemPass::IsLiveVar(uint32_t varId) const {
 void MemPass::AddStores(uint32_t ptr_id, std::queue<Instruction*>* insts) {
   get_def_use_mgr()->ForEachUser(ptr_id, [this, insts](Instruction* user) {
     spv::Op op = user->opcode();
-    if (IsNonPtrAccessChain(op)) {
+    if (IsNonPtrAccessChain(op) || op == spv::Op::OpCopyObject) {
       AddStores(user->result_id(), insts);
     } else if (op == spv::Op::OpStore) {
       insts->push(user);
@@ -243,6 +263,9 @@ 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) {
     auto dbg_op = user->GetCommonDebugOpcode();

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"
@@ -69,6 +70,7 @@ class MemPass : public Pass {
 
  protected:
   MemPass();
+  explicit MemPass(SSARewriteMode ssa_rewrite_mode);
 
   // Returns true if |typeInst| is a scalar type
   // or a vector or matrix
@@ -133,7 +135,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