lhutton1 commented on a change in pull request #10375:
URL: https://github.com/apache/tvm/pull/10375#discussion_r816144228
##########
File path: src/relay/backend/contrib/cmsisnn/scalar_to_tensor_constant.cc
##########
@@ -177,14 +173,22 @@ class ScalarToTensorConstantMutator : public
MixedModeMutator {
};
IRModule ScalarToTensorConstant(const IRModule& mod) {
- auto mutator = ScalarToTensorConstantMutator(mod);
- Function main_func = Downcast<Function>(mod->Lookup("main"));
- auto new_main_body = mutator.VisitExpr(main_func->body);
- if (!new_main_body.same_as(main_func->body)) {
- auto main_var = mod->GetGlobalVar("main");
- auto new_main_func = Function(main_func->params, new_main_body,
main_func->ret_type,
- main_func->type_params, main_func->attrs);
- mod->Update(main_var, new_main_func);
+ for (auto gv : mod->GetGlobalVars()) {
+ Function func = Downcast<Function>(mod->Lookup(gv));
+
+ // only mutate CMSIS-NN external functions
+ auto compiler_name = func->GetAttr<String>(attr::kCompiler);
+ if (!compiler_name.defined() || compiler_name != "cmsis-nn") {
Review comment:
I believe the two approaches here are slightly different. With the
approach mentioned above it seems like we still visit the main function,
including the call nodes within it. Its true that we won't visit non CMSIS-NN
external functions, but it seems like its possible to still mutate some
operators in main - is this okay for this pass? Meanwhile, with the approach in
the PR it is guaranteed that only CMSIS-NN functions are mutated and that we
won't accidentally mutate things in main. Please let me know if I missed
something though as I'm unfamiliar with these passes :)
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
