alexbooth commented on a change in pull request #5272: [BYOC] Add example of
Composite + Annotate for DNNL fused op
URL: https://github.com/apache/incubator-tvm/pull/5272#discussion_r405809268
##########
File path: src/relay/backend/contrib/dnnl/codegen.cc
##########
@@ -30,14 +30,108 @@
#include <tvm/runtime/registry.h>
#include <fstream>
+#include <numeric>
#include <sstream>
+#include "../../utils.h"
#include "../codegen_c/codegen_c.h"
namespace tvm {
namespace relay {
namespace contrib {
+using namespace backend;
+
+const CallNode* GetRootConv2DCall(const CallNode* relu_call) {
+ CHECK(relu_call && IsOp(relu_call, "nn.relu"));
+ const auto relu_arg = relu_call->args[0];
+ const CallNode* add_call = relu_arg.as<CallNode>();
+ CHECK(add_call && IsOp(add_call, "add"));
+ const auto add_arg = add_call->args[0];
+ const CallNode* conv_call = add_arg.as<CallNode>();
+ CHECK(conv_call && IsOp(conv_call, "nn.conv2d"));
+ return conv_call;
+}
Review comment:
Maybe this can be more generic, might get tedious to implement this kind of
function if there are many different patterns. Could we just pass in a vector
of op names like ["nn.relu", "add", "nn.conv2d"] along with the CallNode?
----------------------------------------------------------------
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
