comaniac commented on a change in pull request #4741: [External codegen] Add
test cases for fused ops with manual annotation
URL: https://github.com/apache/incubator-tvm/pull/4741#discussion_r368362837
##########
File path: src/relay/backend/contrib/dnnl/codegen.cc
##########
@@ -50,82 +51,109 @@ class CodegenDNNL : public ExprVisitor, public
CodegenCBase {
out_.push_back({node->name_hint(), 0});
}
- void VisitExpr_(const TupleGetItemNode* op) final {
- // Do nothing
- }
-
void VisitExpr_(const CallNode* call) final {
- std::ostringstream decl_stream;
- std::ostringstream buf_stream;
- // Args: ID
- std::vector<std::string> args;
+ struct Output {
+ std::string decl, buf;
+ int out_size = 1;
+ std::string out;
+ };
+
+ auto generate_body = [=](const CallNode* root_call, const std::string&
func_name,
+ const std::vector<std::string>& args,
+ const std::vector<std::string>& fused_func_args) {
+ // Make function call with input buffers when visiting arguments
+ bool first = true;
+ std::ostringstream arg_stream;
+ arg_stream << "(";
+ for (size_t i = 0; i < root_call->args.size(); ++i) {
+ VisitExpr(root_call->args[i]);
+ for (auto out : out_) {
+ if (!first) {
+ arg_stream << ", ";
+ }
+ first = false;
+ arg_stream << out.first;
+ }
+ }
+
+ for (auto arg_name : fused_func_args) {
+ arg_stream << ", " << arg_name;
+ }
+
+ // Analyze the output buffer
+ auto type_node = root_call->checked_type().as<TensorTypeNode>();
+ CHECK(type_node != nullptr && runtime::TypeMatch(type_node->dtype,
kDLFloat, 32))
+ << "Only support single output tensor with float type";
+
+ auto out_shape = GetShape(root_call->checked_type());
+
+ Output ret;
+ ret.out = "buf_" + std::to_string(buf_idx_++);
+ ret.out_size = std::accumulate(out_shape.begin(), out_shape.end(), 1,
std::multiplies<int>());
+
+ this->PrintIndents();
+
+ std::ostringstream buf_stream;
+ buf_stream << "float* " << ret.out << " = (float*)std::malloc(4 * " <<
ret.out_size << ");";
+ ret.buf = buf_stream.str();
- // Get the arguments for various DNNL kernels.
- if (IsOp(call, "nn.conv2d")) {
- decl_stream << "dnnl_conv2d";
- args = Conv2d(call);
+ arg_stream << ", " << ret.out;
+ // Attach attribute arguments
+ for (size_t i = 0; i < args.size(); ++i) {
+ arg_stream << ", " << args[i];
+ }
+ arg_stream << ");";
+ ret.decl = func_name + arg_stream.str();
+
+ return ret;
+ };
+
+ Output ret;
+ if (auto conv_call = DetectFusedConv2DBiasReLU(call)) {
Review comment:
While I agree with you that it's fine to handle fusion in this DNNL codegen,
I also agree with @zhiics that the current implementation is a bit too ad-hoc
even it's only used for demo purpose for now. As you have implemented, MKL DNN
uses `set_post_ops` to attach ops to be fused. I think this part could be more
general. For example:
```
if call == "relu":
visit(arg)
if this->curr_layer == "conv2d":
generate_post_ops(call)
else:
generate_a_layer(call)
```
In this way, the codegen is able to deal with all MKL DNN supported conv2d
fusion (conv2d, conv2d+add, conv2d+add+relu). We could still put heuristic
pattern annotations to the annotator and improve it gradually. I like the one
you made for conv2d+bias+relu in this PR, for instance.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
