[GitHub] [tvm] mbs-octoml commented on a change in pull request #8423: Implementation of relay_to_tir target hook

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

mbs-octoml commented on a change in pull request #8423:
URL: https://github.com/apache/tvm/pull/8423#discussion_r688563079



##########
File path: src/target/target_kind.cc
##########
@@ -431,6 +431,9 @@ TVM_REGISTER_TARGET_KIND("hybrid", kDLCPU)  // line break
 
 TVM_REGISTER_TARGET_KIND("composite", kDLCPU).add_attr_option<Array<Target>>("devices");
 
+TVM_REGISTER_TARGET_KIND("test", kDLCPU)
+    .set_attr<String>("relay_to_tir", "target.test.tir_lowering");
+

Review comment:
       I think we're stuck with it this way so long as the units tests are driven from the python side with a standard tvm lib.

##########
File path: src/relay/backend/te_compiler.cc
##########
@@ -185,16 +192,25 @@ class TECompilerImpl : public TECompilerNode {
     }
     cur_ccache_key_ = key;
 
-    // No need to lower external functions for now. We will invoke the external
-    // codegen tool once and lower all functions together.
     if (key->source_func->GetAttr<String>(attr::kCompiler).defined()) {
       auto ir_module = IRModule();
       const auto name_node = key->source_func->GetAttr<String>(tvm::attr::kGlobalSymbol);
       ICHECK(name_node.defined()) << "External function has not been attached a name yet.";
+
       auto func_name = GetUniqueName(name_node.value(), &name_map_);
       auto target = Target("ext_dev");
       auto global_var = GlobalVar(func_name);
       global_var->checked_type_ = key->source_func->checked_type();
+
+      auto code_gen_name = key->source_func->GetAttr<String>(attr::kCompiler).value();
+      auto target_kind = tvm::TargetKind::Get(code_gen_name);
+      auto custom_lowering_to_tir =
+          target_kind ? target_kind.value().GetRegisteredHook("relay_to_tir") : nullptr;
+      if (custom_lowering_to_tir != nullptr) {
+        target = key->target;
+        ir_module->Update((*custom_lowering_to_tir)(key->source_func, key->target));
+      }
+

Review comment:
       Between Andrew, Jared and myself we've been thinking the ideal world would be something like:
    - we can recover a Target from any prim call node. They get there using a combination of registered target tags and targets, some way of annotating targets for a lexical scope, and a transform to push the consequences of those annotations down to every node.
    - the relay_to_tir packed func would be an attribute on the TargetKind just as you've done, but can be retrieved directly from the Target thanks to the TargetKind-to-Target attribute copying (if I read that correctly).
    - the BYOC compiler attribute is retired in favor of these Target attributes.
   
   So that's all well and good but is a refactor I'm loath to foist on you now.
   
   So for right now:
    - Can we make the relay_to_tir a TypedPackedFunc so the contract is clear? And a comment in one place explaining what it is supposed to be.
    - Try to bring the compiler-name-to-relay-to-tir lookup code into one place, but continue with the TargetKind lookup as you've written it.
   
   

##########
File path: src/relay/backend/te_compiler.cc
##########
@@ -185,16 +192,25 @@ class TECompilerImpl : public TECompilerNode {
     }
     cur_ccache_key_ = key;
 
-    // No need to lower external functions for now. We will invoke the external
-    // codegen tool once and lower all functions together.
     if (key->source_func->GetAttr<String>(attr::kCompiler).defined()) {
       auto ir_module = IRModule();
       const auto name_node = key->source_func->GetAttr<String>(tvm::attr::kGlobalSymbol);
       ICHECK(name_node.defined()) << "External function has not been attached a name yet.";
+
       auto func_name = GetUniqueName(name_node.value(), &name_map_);
       auto target = Target("ext_dev");
       auto global_var = GlobalVar(func_name);
       global_var->checked_type_ = key->source_func->checked_type();
+
+      auto code_gen_name = key->source_func->GetAttr<String>(attr::kCompiler).value();
+      auto target_kind = tvm::TargetKind::Get(code_gen_name);
+      auto custom_lowering_to_tir =
+          target_kind ? target_kind.value().GetRegisteredHook("relay_to_tir") : nullptr;
+      if (custom_lowering_to_tir != nullptr) {
+        target = key->target;
+        ir_module->Update((*custom_lowering_to_tir)(key->source_func, key->target));
+      }
+

Review comment:
       lgtm, only nit is to make "relay_to_tir" a constexpr with a comment explaining the intended signature. Unless I'm missing it currently that's implied only by the call? thanks.




-- 
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