This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6eb3a1fc36 [BYOC-DNNL] suppport more dnnl ops (#11823)
6eb3a1fc36 is described below
commit 6eb3a1fc3688bede69efd7f14b313b35497ccf02
Author: Ivy Zhang <[email protected]>
AuthorDate: Mon Jul 25 15:16:16 2022 +0800
[BYOC-DNNL] suppport more dnnl ops (#11823)
* support dnnl.global_avg_pooling2d
* fuse pad-avg_pool2d
* fix lint
---
python/tvm/relay/op/contrib/dnnl.py | 14 ++++++++++
src/runtime/contrib/dnnl/dnnl_json_runtime.cc | 36 +++++++++++++++++--------
tests/python/contrib/test_dnnl.py | 25 +++++++++++++++++
tests/python/relay/test_pass_partition_graph.py | 4 +--
4 files changed, 66 insertions(+), 13 deletions(-)
diff --git a/python/tvm/relay/op/contrib/dnnl.py
b/python/tvm/relay/op/contrib/dnnl.py
index 228619e0ef..fa98ed002c 100644
--- a/python/tvm/relay/op/contrib/dnnl.py
+++ b/python/tvm/relay/op/contrib/dnnl.py
@@ -89,6 +89,7 @@ _register_external_op_helper("nn.conv3d_transpose")
_register_external_op_helper("nn.dense")
_register_external_op_helper("nn.max_pool2d")
_register_external_op_helper("nn.avg_pool2d")
+_register_external_op_helper("nn.global_avg_pool2d")
_register_external_op_helper("nn.max_pool3d")
_register_external_op_helper("nn.avg_pool3d")
_register_external_op_helper("abs")
@@ -459,6 +460,18 @@ def tag2layout(input_data, is_weight=False,
conv_type="Conv1D"):
return res
+def legalize_pad_avg_pool(attrs, inputs, types):
+ """Legalize pad->avg_pool2d pattern.
+ Fuse this pattern into one avg_pool2d with padding = (1, 1),
+ and count_include_pad = True"""
+ data = inputs[0]
+ new_attrs = dict(attrs)
+ if isinstance(data, relay.expr.Call) and data.op.name == "nn.pad":
+ new_attrs["padding"] = (1, 1)
+ new_attrs["count_include_pad"] = True
+ return relay.nn.avg_pool2d(data.args[0], **new_attrs)
+
+
def legalize_group_conv(attrs, inputs, types):
"""Legalize group conv / conv_transpose calculation.
Alter weight layout from OIHW to GOIHW / IOHW to GIOHW"""
@@ -575,6 +588,7 @@ class IsComputeIntensiveGraph(ExprVisitor):
"nn.dense",
"nn.layer_norm",
"nn.batch_matmul",
+ "nn.global_avg_pool2d",
]
)
if isinstance(call.op, tvm.tir.op.Op):
diff --git a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
index 93c53dda16..dcf1a86785 100644
--- a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
+++ b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
@@ -624,32 +624,46 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
void Pooling(const size_t& nid, dnnl::algorithm algo) {
auto node = nodes_[nid];
+ auto op_name = node.GetOpName();
+ bool is_global = op_name.find("global") != std::string::npos;
auto src_tr = GetInput(nid, 0);
auto dst_tr = GetOutput(nid, 0);
- // Setup attributes.
- auto strides = GetNodeAttr<std::vector<int64_t>>(node, "strides");
- auto dilates = GetNodeAttr<std::vector<int64_t>>(node, "dilation");
- auto padding = GetNodeAttr<std::vector<int64_t>>(node, "padding");
- std::vector<int64_t> padding_l(padding.begin(), padding.begin() +
padding.size() / 2);
- std::vector<int64_t> padding_r(padding.begin() + padding.size() / 2,
padding.end());
- auto kernel = GetNodeAttr<std::vector<int64_t>>(node, "pool_size");
+ // Get layout.
auto src_layout = GetNodeAttr<std::string>(node, "layout");
auto dst_layout = GetNodeAttr<std::string>(node, "out_layout");
// dst_layout == "" means to use data_layout
if (dst_layout.empty()) dst_layout = src_layout;
- // Minus one for DNNL representation. No dilation for DNNL is 0, for relay
is 1.
- for (auto& d : dilates) d--;
-
// Take into account provided layout strings
src_tr = src_tr.TreatAs(src_layout);
dst_tr = dst_tr.TreatAs(dst_layout);
+ ICHECK(src_tr.dims().size() > 2);
+
+ std::vector<int64_t> feature_size;
+ for (size_t i = 2; i < src_tr.dims().size(); i++) {
+ feature_size.push_back(int64_t(src_tr.dims()[i]));
+ }
+
+ // Set attributes.
+ auto kernel = is_global ? feature_size :
GetNodeAttr<std::vector<int64_t>>(node, "pool_size");
+ auto strides = is_global ? std::vector<int64_t>(src_tr.dims().size() - 2,
1)
+ : GetNodeAttr<std::vector<int64_t>>(node,
"strides");
+ auto dilates = is_global ? std::vector<int64_t>(src_tr.dims().size() - 2,
1)
+ : GetNodeAttr<std::vector<int64_t>>(node,
"dilation");
+ auto padding = is_global ? std::vector<int64_t>((src_tr.dims().size() - 2)
* 2, 0)
+ : GetNodeAttr<std::vector<int64_t>>(node,
"padding");
+ std::vector<int64_t> padding_l(padding.begin(), padding.begin() +
padding.size() / 2);
+ std::vector<int64_t> padding_r(padding.begin() + padding.size() / 2,
padding.end());
+
+ // Minus one for DNNL representation. No dilation for DNNL is 0, for relay
is 1.
+ for (auto& d : dilates) d--;
+
// Attributes related to AvgPool
- if (algo == dnnl::algorithm::pooling_avg) {
+ if (!is_global && algo == dnnl::algorithm::pooling_avg) {
auto include_pad = GetNodeAttr<bool>(node, "count_include_pad");
algo = include_pad ? dnnl::algorithm::pooling_avg_include_padding
: dnnl::algorithm::pooling_avg_exclude_padding;
diff --git a/tests/python/contrib/test_dnnl.py
b/tests/python/contrib/test_dnnl.py
index 1bf8068b2e..e744cab6e9 100755
--- a/tests/python/contrib/test_dnnl.py
+++ b/tests/python/contrib/test_dnnl.py
@@ -991,6 +991,15 @@ def test_pool2d(run_module, dtype="float32"):
)
+def test_global_avg_pooling2d(run_module, dtype="float32"):
+ x_shape = (1, 3, 32, 32)
+ x = relay.var("x", shape=(x_shape), dtype=dtype)
+ out = relay.nn.global_avg_pool2d(x)
+ out = tvm.IRModule.from_expr(out)
+ config = out, {"x": x_shape}, []
+ run_and_verify_func(config, run_module=run_module)
+
+
def test_pool3d(run_module, dtype="float32"):
def get_graph(
op,
@@ -1196,6 +1205,22 @@ def test_resnetv1_rewrite(run_module, dtype="float32"):
run_and_verify_func(config, run_module=run_module, dtype=dtype)
+def test_fuse_pad_avg_pool(run_module, dtype="float32"):
+ def get_graph():
+ data_shape = (1, 768, 17, 17)
+ x = relay.var("x", shape=data_shape, dtype=dtype)
+ out = relay.nn.pad(x, pad_width=[[0, 0], [0, 0], [1, 1], [1, 1]])
+ out = relay.nn.avg_pool2d(out, pool_size=[3, 3])
+ dic = {"x": data_shape}
+ param_lst = []
+ return out, dic, param_lst
+
+ net, dic, param_lst = get_graph()
+ net = tvm.IRModule.from_expr(net)
+ config = net, dic, param_lst
+ run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
def permute_shape(shape, l_from="", l_to=""):
res_shape = []
for label in l_to:
diff --git a/tests/python/relay/test_pass_partition_graph.py
b/tests/python/relay/test_pass_partition_graph.py
index 4b7ac92136..31aa4e4fe2 100644
--- a/tests/python/relay/test_pass_partition_graph.py
+++ b/tests/python/relay/test_pass_partition_graph.py
@@ -1055,8 +1055,8 @@ def test_dnnl_fuse():
def test_partition_mobilenet():
mod, params = relay.testing.mobilenet.get_workload()
mod = get_partitoned_mod(mod, params, dnnl_patterns)
- # 27 fused conv + bn + relu, one dense and one softmax
- assert len(mod.functions) - 1 == 29 # -1 for main
+ # 27 fused conv + bn + relu, one dense, one softmax and one
global_avg_pooling
+ assert len(mod.functions) - 1 == 30 # -1 for main
def test_exec(mod, params, ref_mod, ref_params, out_shape):
ishape = (1, 3, 224, 224)
