From cd24f0794a22523298d9a5cab489a7a2f6ccc128 Mon Sep 17 00:00:00 2001
From: Kyushick Lee <kyushick@gmail.com>
Date: Fri, 20 Jan 2023 07:30:00 -0800
Subject: [PATCH] Extend ort_backend.py for another ep (#14349)

### Description
<!-- Describe your changes. -->

This PR extends OrtBackend to allow for configuring an EP based on the
name, and fallbacks to existing mechanism that infers the EP based on
tensor affinity if nothing is provided.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Currently OrtBackend needs `get_ort_device()` with the device tag
inferred from torch.Tensor, but ort device is not yet supported for
dort. The change allows run dort with a supported EP, by configuring
dort with a desired EP and letting the dort (ort InferenceSession) take
CPU-affined pytorch Tensors as inputs then inject data transfer nodes
internally.
---
 .../training/torchdynamo/ort_backend.py       | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/orttraining/orttraining/python/training/torchdynamo/ort_backend.py b/orttraining/orttraining/python/training/torchdynamo/ort_backend.py
index 3a0c119c22..bd397edf67 100644
--- a/orttraining/orttraining/python/training/torchdynamo/ort_backend.py
+++ b/orttraining/orttraining/python/training/torchdynamo/ort_backend.py
@@ -337,12 +337,16 @@ def _create_onnx_model(onnx_proto):
     return onnx.ModelProto.FromString(onnx_proto)
 
 
-def _create_onnx_session(onnx_proto, device):
+def _create_onnx_session(onnx_proto, ep: str):
     # TODO(wechi): Add more EPs per PyTorch device types.
     # TODO(wechi): enable external allocators.
+    return onnxruntime.InferenceSession(onnx_proto, providers=[ep])
+
+
+def _infer_ep_from_device(device):
     if device.type == "cuda":
-        return onnxruntime.InferenceSession(onnx_proto, providers=["CUDAExecutionProvider"])
-    return onnxruntime.InferenceSession(onnx_proto, providers=["CPUExecutionProvider"])
+        return "CUDAExecutionProvider"
+    return "CPUExecutionProvider"
 
 
 def _get_onnx_devices(values: Tuple[torch.Tensor, ...]) -> Tuple[ORTC.OrtDevice, ...]:  # type: ignore
@@ -453,13 +457,15 @@ class OrtBackend:
         3. Inside _ort_accelerated_call, it creates onnxruntime.InferenceSession and calls it to execute the sub-graph.
     """
 
-    def __init__(self):
+    def __init__(self, ep: str = ""):
         self._supported_ops = OrtOperatorSupport()
         # TODO: this is a naive implementation of cache without proper guard
         self._partitioner_cache: Dict[torch.fx.GraphModule, torch.fx.GraphModule] = {}
         # TODO: this is a naive implementation of cache without proper guard, this will only work for identical inputs
         self._ort_execution_info = OrtExecutionInfo()
 
+        self.ep = ep
+
     def _ort_acclerated_call(self, graph_module: torch.fx.GraphModule, *args, **kwargs):
         if graph_module in self._ort_execution_info.sessions:
             # We have seen this graph before, so we can use cached objects including session.
@@ -493,10 +499,13 @@ class OrtBackend:
                 _decorate_script_module(script_module, args, (prim_outputs,))
             # Generate ONNX ModelProto from torch._C.Graph.
             onnx_proto = _create_onnx_proto(script_module)
+
             # Initialize a ORT session to execute this ONNX model.
             # TorchDynamo assumes all inputs/outputs are on the same device,
             # so we add execution provider only based on the first input's device.
-            onnx_session = _create_onnx_session(onnx_proto, args[0].device)
+            ep = self.ep if self.ep else _infer_ep_from_device(args[0].device)
+
+            onnx_session = _create_onnx_session(onnx_proto, ep)
             # Cache ORT session. It's reused for the same "graph_module".
             self._ort_execution_info.sessions[graph_module] = onnx_session
             # Generate ONNX model and extract its input and output names.