[BYOC] Is multiple outputs really supported in DNNL codegen?

masahi · May 8, 2020, 1:44am

Hi, since the PR https://github.com/apache/incubator-tvm/pull/5288, I was under impression that DNNL codegen supports multiple outputs.

But there is this assertion that seems to be saying that C codegen-based backend allows only a single output:

github.com

apache/incubator-tvm/blob/master/src/relay/backend/contrib/codegen_c/codegen_c.h#L196


 * \return The emitted code string.
 */
std::string JitImpl(const std::string& ext_func_id, const Array<Var>& args,
                    const std::vector<std::string>& buf_decl,
                    const std::vector<std::string>& body,
                    const std::vector<Output>& out) {
  // Create the signature. For example, it could be:
  // extern "C" void dnnl_0_(float* input0, float* input1, float* out, int M, int N) {}
  code_stream_ << "extern \"C\" void " << ext_func_id << "_(";


  CHECK_EQ(out.size(), 1U) << "Internal error: only single output is support.";


  for (const auto& arg : args) {
    const auto& dtype_str = GetDtypeString(arg);
    code_stream_ << dtype_str << "* " << arg->name_hint() << ", ";
  }
  code_stream_ << out[0].dtype << "* out) {\n";
  this->EnterScope();


  // Function body
  for (auto decl : buf_decl) {

I’m a bit confused now. What is the situation? @comaniac @zhiics

comaniac · May 8, 2020, 2:12am

Ahh…if I remember correctly, I refined the DNNL codegen in the PR you pointed out, so DNNL codegen now should support the op with multiple outputs such as batch norm as the first summary point in the PR description. However, we didn’t support the general C codegen accordingly, meaning that the subgraph with multiple outputs wasn’t supported in that PR. I remember there was an issue when I was trying to implement multiple output support for subgraphs, but I don’t remember the details.

comaniac · May 26, 2020, 10:22pm

@masahi finally get some time to look into the confusing message. Here is a complete explanation:

First, I checked the DNNL codegen and confirmed that it should be able to support a subgraph with multiple outputs whatever the subgraph has one or multiple operators:

github.com

apache/incubator-tvm/blob/master/src/relay/backend/contrib/dnnl/codegen.cc#L284


// Make function call with input buffers when visiting arguments
CHECK_GT(func_args.size(), 0);
std::ostringstream decl_stream;
decl_stream << "(" << func_args[0];
for (size_t i = 1; i < func_args.size(); ++i) {
  decl_stream << ", " << func_args[i];
}


// Analyze the output buffers
std::vector<Type> out_types;
if (root_call->checked_type()->IsInstance<TupleTypeNode>()) {
  auto type_node = root_call->checked_type().as<TupleTypeNode>();
  for (auto field : type_node->fields) {
    CHECK(field->IsInstance<TensorTypeNode>());
    out_types.push_back(field);
  }
} else if (root_call->checked_type()->IsInstance<TensorTypeNode>()) {
  CHECK(root_call->checked_type()->IsInstance<TensorTypeNode>());
  out_types.push_back(root_call->checked_type());
} else {
  LOG(FATAL) << "Unrecognized type node: " << AsText(root_call->checked_type(), false);

With the above statement, I revisited the assertion you pointed out, and realized that we should refine the message. Currently the way we support multiple outputs is to create a new tuple node of them as a single output. That’s why our unit tests are all passed.

On the other hand, “creating a new tuple node of all outputs” is performed by partition graph pass. In order to make sure users run the pass in advance, we still need this checker. Sorry for the confusion, and we should definitely improve the message.

mak · May 27, 2020, 5:20am

@comaniac

Could you please provide example code to creating a new tuple node of all outputs from partition?