[BYOC] [TPAT] [TensorRT] Add the ability to automatically generate TensorRT plugins using TVM

python/tvm/contrib/graph_executor.py

@@ -178,6 +178,12 @@ def __init__(self, module):
         self._load_params = module["load_params"]
         self._share_params = module["share_params"]
 
+        self._get_workspace_dtype = module["get_workspace_dtype"]
+        self._get_workspace_size = module["get_workspace_size"]
+        self._get_function_list = module["get_function_list"]
+        self._get_storageid = module["get_storageid"]
+        self._get_output_eid = module["get_output_eid"]
+
     def set_input(self, key=None, value=None, **params):
         """Set inputs to the module via kwargs
 
@@ -512,3 +518,49 @@ def benchmark(
             cooldown_interval_ms=cooldown_interval_ms,
             repeats_to_cooldown=repeats_to_cooldown,
         )()
+
+    def get_workspace_dtype(self):
+        """Get the dtype of workspace to the graph
+
+        Returns
+        -------
+        dtype : str
+            The dtypes of workspace.
+        """
+        return self._get_workspace_dtype()
+
+    def get_workspace_size(self):
+        """Get the dtype of workspace to the graph
+
+        Returns
+        -------
+        dtype : int
+            The bytes size of workspace.
+        """
+        return self._get_workspace_size()
+
+    def get_function_list(self):
+        """Get the Host Function execute order
+
+        Returns
+        -------
+        dtype : str
+            The Host function execute order
+        """
+        return self._get_function_list()
+
+    def get_storageid(self):
+        return self._get_storageid()
+
+    def get_output_eid(self, index):
+        """Get index-th output to out
+
+        Parameters
+        ----------
+        index : int
+            The output index
+
+        out : NDArray
+            The output array container
+        """
+        return self._get_output_eid(index)

python/tvm/relay/backend/executor_factory.py

@@ -180,6 +180,7 @@ def __init__(
         libmod_name,
         params,
         function_metadata,
+        constant_params = None
     ):
         assert isinstance(graph_json_str, string_types)
         fcreate = get_global_func("tvm.graph_executor_factory.create")
@@ -199,6 +200,12 @@ def __init__(
         self.iter_cnt = 0
         self.function_metadata = function_metadata
 
+        self.constant_params = constant_params
+        self.device_function_list = get_global_func("tir.transform.retrieve_device_function_list")
+        self.device_function_thread_config = get_global_func("runtime.module.retrieve_device_function_thread_config")
+        self.device_memory_size = get_global_func("tir.transform.retrieve_device_memory_size")
+
+
     def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
         return self.module.export_library(file_name, fcompile, addons, **kwargs)
 
@@ -216,3 +223,15 @@ def get_executor_config(self):
 
     def get_lib(self):
         return self.lib
+
+    def get_constant_params(self):
+        return self.constant_params
+
+    def get_device_function_list(self):
+        return self.device_function_list()
+
+    def get_grid_block_thread_config(self):
+        return self.device_function_thread_config()
+
+    def get_device_memory_size(self):
+        return self.device_memory_size()

python/tvm/relay/build_module.py

@@ -70,6 +70,7 @@ def __init__(self):
         self._get_executor_codegen_metadata = self.mod["get_executor_codegen_metadata"]
         self._get_devices = self.mod["get_devices"]
         self._get_irmodule = self.mod["get_irmodule"]
+        self._get_constant_params = self.mod["get_constant_params"]
 
     def build(
         self,
@@ -249,6 +250,14 @@ def get_params(self):
             ret[key] = value.data
         return ret
 
+    def get_constant_params(self):
+        """Return the constant params."""
+        params = self._get_constant_params()
+        ret = {}
+        for key, value in params.items():
+            ret[key] = value.data.asnumpy()
+        return ret
+
     def get_irmodule(self):
         """Returns the TargetIRModule's post-lowering"""
         return self._get_irmodule()
@@ -372,6 +381,7 @@ def build(
             mod_name=mod_name,
         )
         func_metadata = bld_mod.get_function_metadata()
+        constant_params = bld_mod.get_constant_params()
         devices = bld_mod.get_devices()
         lowered_ir_mods = bld_mod.get_irmodule()
         executor_codegen_metadata = bld_mod.get_executor_codegen_metadata()
@@ -400,6 +410,7 @@ def build(
                 mod_name,
                 params,
                 func_metadata,
+                constant_params=constant_params
             )
         else:
             assert False, "Executor " + executor + " not supported"

python/tvm/tpat/__init__.py

@@ -0,0 +1,18 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from . import cuda

python/tvm/tpat/cuda/__init__.py

@@ -0,0 +1,18 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from .pipeline import pipeline