How to customize node names?

gasgallo · March 1, 2021, 4:48am

Hi everybody,

I’m looking into how to use custom node names into TVM models.

For example I’d like to add a prefix or a suffix, e.g. fused_nn_contrib_conv2d_NCHWc_add_nn_relu_9 becomes prefix_fused_nn_contrib_conv2d_NCHWc_add_nn_relu_9_postfix.

Can anybody point me to the right piece of code that I can modify to make this possible?

Many thanks.

zhuochenKIDD · March 3, 2021, 5:27am

perhaps try here:

github.com

apache/tvm/blob/main/src/relay/backend/compile_engine.cc#L130


      const auto* ttype = field.as<TensorTypeNode>();
      // TODO(@icemelon): Allow recursive tuple
      ICHECK(ttype != nullptr);
      tvm::te::Tensor tensor = tvm::te::placeholder(GetShape(ttype->shape), ttype->dtype);
      cache_node->inputs.push_back(tensor);
      inputs.push_back(tensor);
    }
  }
  memo_[param] = inputs;
}
readable_name_stream_ << "fused";
cache_node->outputs = this->VisitExpr(prim_func->body);
auto candidate_name = readable_name_stream_.str();
constexpr static size_t kMaxFuncNameLength = 80;
if (candidate_name.size() > kMaxFuncNameLength) {
  std::stringstream truncated_name;
  truncated_name << candidate_name.substr(0, kMaxFuncNameLength);
  truncated_name << "_" << std::hash<std::string>{}(candidate_name) << "_";
  candidate_name = truncated_name.str();
}
cache_node->func_name = candidate_name;