[Vulkan] Add support for Optimization Blocklist to Vulkan Rewrite (#87431)

Optimization Blocklist will be used in a future diff (D40315730) to make the rewrite to transfer input/output backends optional

Differential Revision: [D40315729](https://our.internmc.facebook.com/intern/diff/D40315729/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/87431
Approved by: https://github.com/mcr229, https://github.com/digantdesai
This commit is contained in:
Salil Desai
2022-10-30 20:30:55 -07:00
committed by PyTorch MergeBot
parent f717986f93
commit bc68625151
9 changed files with 35 additions and 19 deletions

View File

@ -52,6 +52,7 @@
#include <torch/csrc/jit/passes/lower_graph.h>
#include <torch/csrc/jit/passes/lower_tuples.h>
#include <torch/csrc/jit/passes/metal_rewrite.h>
#include <torch/csrc/jit/passes/mobile_optimizer_type.h>
#include <torch/csrc/jit/passes/normalize_ops.h>
#include <torch/csrc/jit/passes/peephole.h>
#include <torch/csrc/jit/passes/peephole_list_idioms.h>
@ -1081,8 +1082,10 @@ void initJITBindings(PyObject* module) {
.def(
"_jit_pass_vulkan_optimize_for_mobile",
[](script::Module& module,
std::set<MobileOptimizerType>& optimization_blocklist,
std::vector<std::string>& preserved_methods) {
return vulkanOptimizeForMobile(module, preserved_methods);
return vulkanOptimizeForMobile(
module, optimization_blocklist, preserved_methods);
})
.def(
"_jit_pass_metal_insert_prepacked_ops",