Enable options for openvino backend in torch.compile (#20946)

* Added backend options for openvino torch.compile

* Added config options to enable openvino configs

* Fixed linter issues

* Addressed review comments

---------

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py

@@ -20,7 +20,8 @@ from openvino.runtime import Core, Type, PartialShape
 from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
 from openvino.frontend.pytorch.torchdynamo.execute import execute, execute_cached
-from openvino.frontend.pytorch.torchdynamo.compile import cached_model_name, cache_root_path, get_device, openvino_compile_cached_model
+from openvino.frontend.pytorch.torchdynamo.compile import cached_model_name, openvino_compile_cached_model
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device, _get_model_caching
 
 from openvino.runtime import Core, Type, PartialShape
 
@@ -44,8 +45,8 @@ log = logging.getLogger(__name__)
 
 @register_backend
 @fake_tensor_unsupported
-def openvino(subgraph, example_inputs):
-    return fx_openvino(subgraph, example_inputs)
+def openvino(subgraph, example_inputs, options=None):
+    return fx_openvino(subgraph, example_inputs, options)
 
 @register_backend
 @fake_tensor_unsupported
@@ -111,18 +112,19 @@ def ts_openvino(subgraph, example_inputs):
         return compile_fx(subgraph, example_inputs)
 
 
-def fx_openvino(subgraph, example_inputs):
+def fx_openvino(subgraph, example_inputs, options):
     try:
         executor_parameters = None
         inputs_reversed = False
-        if os.getenv("OPENVINO_TORCH_MODEL_CACHING") is not None:
+        openvino_model_caching = _get_model_caching(options)
+        if openvino_model_caching is not None:
             # Create a hash to be used for caching
             model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest()
             executor_parameters = {"model_hash_str": model_hash_str}
             # Check if the model was fully supported and already cached
             example_inputs.reverse()
             inputs_reversed = True
-            maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, cache_root_path())
+            maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", _get_device(options), example_inputs, _get_cache_dir(options))
             if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"):
                 # Model is fully supported and already cached. Run the cached OV model directly.
                 compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs)
@@ -146,7 +148,7 @@ def fx_openvino(subgraph, example_inputs):
 
         def _call(*args):
             res = execute(compiled_model, *args, executor="openvino",
-                          executor_parameters=executor_parameters)
+                          executor_parameters=executor_parameters, options=options)
             return res
         return _call
     except Exception as e:

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend_utils.py

@@ -0,0 +1,60 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+from typing import Optional, Any
+import os
+from openvino.runtime import Core
+
+
+def _get_device(options) -> Optional[Any]:
+    core = Core()
+    device = "CPU"
+
+    if "device" in options:
+        device = options["device"]
+    else:
+        device = os.getenv("OPENVINO_TORCH_BACKEND_DEVICE")
+
+    if device is not None:
+        assert device in core.available_devices, (
+            "Specified device "
+            + device
+            + " is not in the list of OpenVINO Available Devices"
+        )
+
+    return device
+
+def _is_cache_dir_in_config(options) -> Optional[Any]:
+    if "config" in options:
+        cfg = options["config"]
+        if "CACHE_DIR" in cfg:
+            return True
+    return False
+
+
+def _get_cache_dir(options) -> Optional[Any]:
+    cache_dir = "./cache"
+    if options is not None and "cache_dir" in options:
+        cache_dir = options["cache_dir"]
+    if _is_cache_dir_in_config(options):
+        cache_dir = options["config"]["CACHE_DIR"]
+    else:
+        cache_dir_env = os.getenv("OPENVINO_TORCH_CACHE_DIR")
+        if cache_dir_env is not None:
+            cache_dir = cache_dir_env
+    return cache_dir
+
+
+def _get_model_caching(options) -> Optional[Any]:
+    if options is not None and "model_caching" in options:
+        return options["model_caching"]
+    else:
+        return os.getenv("OPENVINO_TORCH_MODEL_CACHING")
+
+
+def _get_config(options) -> Optional[Any]:
+    if options is not None and "config" in options:
+        return options["config"]

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py

@@ -15,6 +15,7 @@ from torch.fx import GraphModule
 from openvino.frontend import FrontEndManager
 from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
 from openvino.runtime import Core, Type, PartialShape, serialize
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device, _get_config, _is_cache_dir_in_config
 
 from typing import Callable, Optional
 
@@ -32,33 +33,17 @@ def cached_model_name(model_hash_str, device, args, cache_root, reversed = False
         return None
 
     inputs_str = ""
-    for idx, input_data in enumerate(args):  
+    for idx, input_data in enumerate(args):
         if reversed:
             inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str
         else:
             inputs_str += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "")
     inputs_str = sha256(inputs_str.encode('utf-8')).hexdigest()
     file_name += inputs_str
-    
+
     return file_name
 
-def cache_root_path():
-    cache_root = "./cache/"
-    if os.getenv("OPENVINO_TORCH_CACHE_DIR") is not None:
-        cache_root = os.getenv("OPENVINO_TORCH_CACHE_DIR")
-    return cache_root
-
-def get_device():
-    core = Core()
-    device = "CPU"
-
-    if os.getenv("OPENVINO_TORCH_BACKEND_DEVICE") is not None:
-        device = os.getenv("OPENVINO_TORCH_BACKEND_DEVICE")
-        assert device in core.available_devices, "Specified device " + device + " is not in the list of OpenVINO Available Devices"
-
-    return device
-
-def openvino_compile_cached_model(cached_model_path, *example_inputs):
+def openvino_compile_cached_model(cached_model_path, *example_inputs, options):
     core = Core()
     om = core.read_model(cached_model_path + ".xml")
 
@@ -78,17 +63,22 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
 
-    core.set_property({'CACHE_DIR': cache_root_path() + '/blob'})
+    config = {}
 
-    compiled_model = core.compile_model(om, get_device())
+    if _is_cache_dir_in_config(options):
+        config = _get_config(options)
+    else:
+        config["CACHE_DIR"] = _get_cache_dir(options)
+
+    compiled_model = core.compile_model(om, _get_device(options), config)
 
     return compiled_model
 
-def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None):
+def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, options=None):
     core = Core()
 
-    device = get_device()
-    cache_root = cache_root_path()
+    device = _get_device(options)
+    cache_root = _get_cache_dir(options)
     file_name = cached_model_name(model_hash_str, device, args, cache_root)
 
     if file_name is not None and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin"):
@@ -99,7 +89,7 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None):
 
         input_shapes = []
         input_types = []
-        for idx, input_data in enumerate(args):  
+        for idx, input_data in enumerate(args):
             input_types.append(input_data.type())
             input_shapes.append(input_data.size())
 
@@ -128,8 +118,13 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None):
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
 
-    if model_hash_str is not None:
-        core.set_property({'CACHE_DIR': cache_root + '/blob'})
+    config = {}
 
-    compiled = core.compile_model(om, device)
+    if model_hash_str is not None:
+        if _is_cache_dir_in_config(options):
+            config = _get_config(options)
+        else:
+            config["CACHE_DIR"] = cache_root
+
+    compiled = core.compile_model(om, device, config)
     return compiled

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py

@@ -21,8 +21,9 @@ from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
 from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
 from openvino.frontend.pytorch.torchdynamo.compile import openvino_compile
 from openvino.runtime import Core, Type, PartialShape
+from openvino.frontend.pytorch.torchdynamo.backend_utils import _get_cache_dir, _get_device
 
-from typing import Callable, Optional
+from typing import Callable, Optional, Any
 
 from torch.fx.experimental.proxy_tensor import make_fx, wrapper_and_args_for_make_fx
 
@@ -44,9 +45,10 @@ def execute(
     *args,
     executor: str = "openvino",
     executor_parameters: Optional[dict] = None,
+    options: Optional[Any] = None,
 ):
     if executor == "openvino":
-        return openvino_execute_partitioned(gm, *args, executor_parameters=executor_parameters)
+        return openvino_execute_partitioned(gm, *args, executor_parameters=executor_parameters, options=options)
     elif executor == "strictly_openvino":
         return openvino_execute(gm, *args, executor_parameters=executor_parameters)
 
@@ -65,7 +67,7 @@ def execute_cached(compiled_model, *args):
     return result
 
 
-def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id):
+def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id, options):
 
     executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
 
@@ -86,7 +88,7 @@ def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition
     if use_cache and (partition_id in compiled_cache):
         compiled = compiled_cache[partition_id]
     else:
-        compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str)
+        compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str, options=options)
         compiled_cache[partition_id] = compiled
 
     flat_args, _ = tree_flatten(args)
@@ -101,20 +103,21 @@ def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition
 
 
 class OpenVINOGraphModule(torch.nn.Module):
-    def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None):
+    def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, options=None):
         super().__init__()
         self.gm = gm
         self.partition_id = partition_id
         self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache,
                                     "model_hash_str": model_hash_str}
         self.perm_fallback = False
+        self.options = options
 
     def __call__(self, *args):
         if self.perm_fallback:
             return self.gm(*args)
 
         try:
-            result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id)
+            result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id, options=self.options)
         except Exception:
             self.perm_fallback = True
             return self.gm(*args)
@@ -122,7 +125,7 @@ class OpenVINOGraphModule(torch.nn.Module):
         return result
 
 
-def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None):
+def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None, options=None):
     global max_openvino_partitions
     partition_id = max_openvino_partitions
     for node in gm.graph.nodes:
@@ -133,7 +136,7 @@ def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_s
             gm.add_submodule(
                 node.target,
                 OpenVINOGraphModule(openvino_submodule, partition_id, use_python_fusion_cache,
-                                    model_hash_str=model_hash_str),
+                                    model_hash_str=model_hash_str, options=options),
             )
             partition_id = partition_id + 1
 
@@ -142,7 +145,7 @@ def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_s
     return gm
 
 
-def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=None):
+def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=None, options=None):
     executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
 
     global partitioned_modules
@@ -162,7 +165,7 @@ def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=Non
 
     if signature not in partitioned_modules:
         partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache,
-                                                         model_hash_str=model_hash_str)
+                                                         model_hash_str=model_hash_str, options=options)
     return partitioned_modules[signature](*args)