[vulkan][py] torch.utils.optimize_for_vulkan (pytorch#44903)

Summary: Pull Request resolved: pytorch#44903

Test Plan: Imported from OSS

Reviewed By: kimishpatel

Differential Revision: D23766039

Pulled By: IvanKobzarev

fbshipit-source-id: dbdf484ee7d3a7719aab105efba51b92ebc51568

Author: IvanKobzarev

torch/_C/__init__.pyi.in

@@ -156,6 +156,8 @@ def _jit_get_operation(op_name: str) -> Callable: ...
 def _jit_pass_optimize_for_mobile(module: 'torch.jit.ScriptModule',
                                   optimization_blocklist: Set[MobileOptimizerType],
                                   preserved_methods: List[AnyStr]) -> 'torch.jit.ScriptModule': ...
+def _jit_pass_vulkan_optimize_for_mobile(module: 'torch.jit.ScriptModule',
+                                         preserved_methods: List[AnyStr]) -> 'torch.jit.ScriptModule': ...
 def _jit_pass_inline(Graph) -> None: ...
 def _jit_get_schemas_for_operator(name :str) -> List[FunctionSchema]: ...
 def _jit_can_fuse_on_cpu() -> _bool: ...

torch/csrc/jit/passes/vulkan_rewrite.cpp

@@ -159,12 +159,14 @@ void vulkanFoldPrePackingOps(script::Module& m) {
   PrePackingOpsFolder(m, filter_fn, "prepack_folding");
 }
 
-script::Module vulkanOptimizeForMobile(const script::Module& m) {
+script::Module vulkanOptimizeForMobile(
+    const script::Module& m,
+    const std::vector<std::string>& preserved_methods) {
   auto cloned_module = m.clone();
   cloned_module.eval();
   cloned_module = FoldConvBatchNorm(cloned_module);
   vulkanInsertPrePackedOps(cloned_module);
-  cloned_module = freeze_module(cloned_module);
+  cloned_module = freeze_module(cloned_module, preserved_methods);
   vulkanFusePrePackedConvWithClamp(cloned_module);
   vulkanFoldPrePackingOps(cloned_module);
   removeDropout(cloned_module);
@@ -193,7 +195,9 @@ void vulkanFoldPrePackingOps(script::Module& m) {
       "Vulkan is not enabled. Please build with USE_VULKAN=1");
 }
 
-script::Module vulkanOptimizeForMobile(const script::Module& module) {
+script::Module vulkanOptimizeForMobile(
+    const script::Module& module,
+    const std::vector<std::string>& preserved_methods) {
   TORCH_INTERNAL_ASSERT(
       "Mobile optimizaiton only available with Vulkan at the moment. "
       "Vulkan is not enabled. Please build with USE_VULKAN=1");

torch/csrc/jit/passes/vulkan_rewrite.h

@@ -9,6 +9,8 @@ TORCH_API void vulkanInsertPrePackedOps(std::shared_ptr<Graph>& graph);
 TORCH_API void vulkanInsertPrePackedOps(script::Module& module);
 TORCH_API void vulkanFusePrePackedConvWithClamp(script::Module& module);
 TORCH_API void vulkanFoldPrePackingOps(script::Module& module);
-TORCH_API script::Module vulkanOptimizeForMobile(const script::Module& module);
+TORCH_API script::Module vulkanOptimizeForMobile(
+    const script::Module& module,
+    const std::vector<std::string>& preserved_methods);
 } // namespace jit
 } // namespace torch

torch/csrc/jit/python/init.cpp

@@ -647,8 +647,9 @@ void initJITBindings(PyObject* module) {
           })
       .def(
           "_jit_pass_vulkan_optimize_for_mobile",
-          [](script::Module& module) {
-            return vulkanOptimizeForMobile(module);
+          [](script::Module& module,
+             std::vector<std::string>& preserved_methods) {
+            return vulkanOptimizeForMobile(module, preserved_methods);
           })
       .def(
           "_jit_pass_onnx_unpack_quantized_weights",

torch/utils/mobile_optimizer.py

@@ -16,14 +16,16 @@ class LintCode(Enum):
 def optimize_for_mobile(
         script_module,
         optimization_blocklist: Set[MobileOptimizerType] = None,
-        preserved_methods: List[AnyStr] = None):
+        preserved_methods: List[AnyStr] = None,
+        backend: str = 'CPU'):
     """
     Args:
         script_module: An instance of torch script module with type of ScriptModule.
         optimization_blocklist: A set with type of MobileOptimizerType. When set is not passed,
             optimization method will run all the optimizer pass; otherwise, optimizer
             method will run the optimization pass that is not included inside optimization_blocklist.
-        perserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked.
+        perserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked
+        backend: Device type to use for running the result model ('CPU'(default) or 'Vulkan').
     Returns:
         A new optimized torch script module
     """
@@ -37,7 +39,13 @@ def optimize_for_mobile(
     if preserved_methods is None:
         preserved_methods = []
 
-    optimized_cpp_module = torch._C._jit_pass_optimize_for_mobile(script_module._c, optimization_blocklist, preserved_methods)
+    if backend == 'CPU':
+        optimized_cpp_module = torch._C._jit_pass_optimize_for_mobile(script_module._c, optimization_blocklist, preserved_methods)
+    elif backend == 'Vulkan':
+        optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(script_module._c, preserved_methods)
+    else:
+        raise TypeError("Unknown backend, must be one of 'CPU', 'Vulkan'")
+
     return torch.jit._recursive.wrap_cpp_module(optimized_cpp_module)