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 <yan3.zh...@intel.com> 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)