[GitHub] [tvm] manupa-arm commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

[GitBox <gi...@apache.org>]

manupa-arm commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749197497



##########
File path: python/tvm/relay/backend/executor_factory.py
##########
@@ -75,6 +75,8 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
     ----------
     ir_mod : :py:class:`~tvm.IRModule`
         The IR module to build.
+    built_ir_mods : dict[Target, IRModule]
+        The IR modules built per Target.

Review comment:
       What does IR modules "built" mean here ?  When IRModule are built wont they become runtime.Modules ?

##########
File path: python/tvm/relay/build_module.py
##########
@@ -376,10 +381,11 @@ def build(ir_mod, target=None, target_host=None, params=None, mod_name="default"
         )
         func_metadata = bld_mod.get_function_metadata()
         devices = bld_mod.get_devices()
+        final_ir_mods = bld_mod.get_irmodule()
 
         if executor == "aot":
             executor_factory = _executor_factory.AOTExecutorFactoryModule(
-                ir_mod, target, runtime_mod, mod_name, params, func_metadata, devices
+                ir_mod, final_ir_mods, target, runtime_mod, mod_name, params, func_metadata, devices

Review comment:
       We should use a better term and possibly change the other ir_mod to disambiguate, though Im not sure what to call the second argument here.

##########
File path: python/tvm/relay/build_module.py
##########
@@ -249,6 +250,10 @@ def get_params(self):
             ret[key] = value.data
         return ret
 
+    def get_irmodule(self):
+        """Returns the Target IRModule's from code generation"""

Review comment:
       Do we want to say something like per target primitive IRModule that is the input for code generation ?
   IRModule from code generation does not sound right. WDYT ?

##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -587,8 +633,12 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     dict_attrs.Set("global_symbol", run_func_name);
     dict_attrs.Set("runner_function", Bool(true));
 
+    tir::Stmt device_activations = GenerateAllDeviceHook("Activate");
+    tir::Stmt device_deactivations = GenerateAllDeviceHook("Deactivate");
+    tir::Stmt final_body = tir::SeqStmt({device_activations, body, device_deactivations});

Review comment:
       Should we just re-use body here ?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org