From 3416e0ae44032a94e198a2b533522d651a7bccfe Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Fri, 28 Jan 2022 13:04:08 +0100 Subject: [PATCH 1/4] identity fuse --- .../mkldnn/mkldnn_identity_property.h | 173 ++++++++++++++++++ .../mkldnn/mkldnn_subgraph_base-inl.h | 2 +- .../mkldnn/mkldnn_subgraph_property.cc | 3 + 3 files changed, 177 insertions(+), 1 deletion(-) create mode 100644 src/operator/subgraph/mkldnn/mkldnn_identity_property.h diff --git a/src/operator/subgraph/mkldnn/mkldnn_identity_property.h b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h new file mode 100644 index 000000000000..6473c9357dc2 --- /dev/null +++ b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h @@ -0,0 +1,173 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file mkldnn_identity_property.cc + * \brief Graph property for removing identity operators + */ + +#ifndef MXNET_OPERATOR_SUBGRAPH_MKLDNN_MKLDNN_IDENTITY_PROPERTY_H_ +#define MXNET_OPERATOR_SUBGRAPH_MKLDNN_MKLDNN_IDENTITY_PROPERTY_H_ +#if MXNET_USE_MKLDNN == 1 + +#include +#include +#include + +#include "../common.h" +#include "../../nn/dropout-inl.h" +#include "mkldnn_subgraph_base-inl.h" + +namespace mxnet { +namespace op { + +class SgMKLDNNIdentitySelector : public SubgraphSelectorV2 { + private: + std::vector matched_list_; + bool pattern_found = false; + + public: + bool Select(const BiDirectedNode& seed_node, + const std::shared_ptr& node_attr) override { + bool status = false; + if (seed_node.node->op() == Op::Get("_copy")) { + status = true; + } + + if (seed_node.node->op() == Op::Get("Dropout")) { + auto const& dropout_param = nnvm::get(seed_node.node->attrs.parsed); + if (dropout_param.mode == dropout::kTraining) { + status = true; + } + } + + if (status) { + matched_list_.clear(); + matched_list_.emplace_back(&seed_node); + return true; + } + return false; + } + + bool SelectInput(const BiDirectedNode& n, const BiDirectedNode& input_node) override { + if (pattern_found || input_node.node->is_variable()) { + return false; + } else if (input_node.node->op()) { + matched_list_.emplace_back(&input_node); + pattern_found = true; + return true; + } + return false; + } + + bool SelectOutput(const BiDirectedNode& n, const BiDirectedNode& output_node) override { + return false; + } + + std::vector Filter(const std::vector& candidates) override { + // candidates should contain only two nodes - custom node and identity node + if (pattern_found && candidates.size() == matched_list_.size()) { + CHECK_EQ(candidates.size(), 2); + return candidates; + } else { + return std::vector(0); + } + } + + void Reset() override { + CHECK_GE(matched_list_.size(), 1); + auto new_selector = SgMKLDNNIdentitySelector(); + new_selector.Select(*matched_list_[0], nullptr); + *this = new_selector; + } +}; + +inline bool IsIdentityNode(const nnvm::ObjectPtr node) { + return node->op() && (node->op() == Op::Get("_copy") || node->op() == Op::Get("Dropout")); +} + +class SgMKLDNNIdentityProperty : public SubgraphProperty { + public: + SgMKLDNNIdentityProperty() {} + + static SubgraphPropertyPtr Create() { + static const std::string& name = "MKLDNN Identity optimization passs"; + auto property = std::make_shared(); + property->SetAttr("property_name", name); + property->SetAttr("inference_only", true); + return property; + } + + nnvm::ObjectPtr CreateSubgraphNode(const nnvm::Symbol& sym, + const int subgraph_id = 0) const override { + nnvm::NodeEntry identity_node_entry; + for (auto entry : sym.outputs) { + if (IsIdentityNode(entry.node)) { + identity_node_entry = entry; + } + } + + auto last_node = identity_node_entry.node; + nnvm::Symbol new_sym; + new_sym.outputs.emplace_back(last_node); + + nnvm::ObjectPtr org_node; + DFSVisit(new_sym.outputs, [&](const nnvm::ObjectPtr& node) { + if (!IsIdentityNode(node)) { + org_node = node; + } + }); + + // Create copy of original node + nnvm::ObjectPtr n = nnvm::Node::Create(); + n->attrs = org_node->attrs; + if (n->op() && n->op()->attr_parser) { + n->op()->attr_parser(&(n->attrs)); + } + + return n; + } + + void ConnectSubgraphOutputs(const nnvm::ObjectPtr n, + std::vector* output_entries) const override { + // output of identity must be connected as output of operator before identity + // e.g. for: /--index 0--> custom_op + // (n) slice + // \--index 1--> Dropout --index 0--> OUT_NODE + // for OUT_NODE index 0 must be changed to index 1 + for (int i = 0; i < output_entries->size(); ++i) { + auto out_node = output_entries->at(i)->node; + if (IsIdentityNode(out_node)) { + output_entries->at(i)->index = out_node->inputs[0].index; + } + output_entries->at(i)->node = n; + } + } + + SubgraphSelectorV2Ptr CreateSubgraphSelectorV2() const override { + auto selector = std::make_shared(); + return selector; + } +}; + +} // namespace op +} // namespace mxnet + +#endif // if MXNET_USE_MKLDNN == 1 +#endif // MXNET_OPERATOR_SUBGRAPH_MKLDNN_MKLDNN_IDENTITY_PROPERTY_H_ \ No newline at end of file diff --git a/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h b/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h index 6436852ee96a..9e0f9bccaf4b 100644 --- a/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h +++ b/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h @@ -31,7 +31,7 @@ static inline bool SupportMKLDNNAttr(const std::shared_ptr& node_attr) return (node_attr->dispatch_mode == DispatchMode::kFComputeEx) && (node_attr->itype[0] == mshadow::kFloat32 || node_attr->itype[0] == mshadow::kBfloat16) && - (ndim == 1 || ndim == 2 || ndim == 4 || ndim == 5); + (ndim >= 1 && ndim <= 5); } else { return true; } diff --git a/src/operator/subgraph/mkldnn/mkldnn_subgraph_property.cc b/src/operator/subgraph/mkldnn/mkldnn_subgraph_property.cc index 2701e0daa1a9..bc9040777525 100644 --- a/src/operator/subgraph/mkldnn/mkldnn_subgraph_property.cc +++ b/src/operator/subgraph/mkldnn/mkldnn_subgraph_property.cc @@ -24,6 +24,7 @@ #include "mkldnn_fc_post_quantize_property.h" #include "mkldnn_fc_property.h" #include "mkldnn_fc_sum_fuse.h" +#include "mkldnn_identity_property.h" #include "mkldnn_post_quantize_align_scale_property.h" #include "mkldnn_post_quantize_property.h" #include "mkldnn_transformer_post_quantize_property.h" @@ -36,6 +37,7 @@ MXNET_REGISTER_SUBGRAPH_BACKEND(MKLDNN) .set_attr("enable", MKLDNNEnvSet()) .set_attr("context", Context::CPU()); +MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN, SgMKLDNNIdentityProperty); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN, SgMKLDNNConvProperty); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN, SgMKLDNNFCProperty); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN, SgMKLDNNTransformerProperty); @@ -43,6 +45,7 @@ MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN, SgMKLDNNFCSumFuseProperty); MXNET_REGISTER_SUBGRAPH_BACKEND(MKLDNN_QUANTIZE).set_attr("context", Context::CPU()); +MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN_QUANTIZE, SgMKLDNNIdentityProperty); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN_QUANTIZE, SgMKLDNNConvProperty).set_attr("quantize", true); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN_QUANTIZE, SgMKLDNNFCProperty).set_attr("quantize", true); MXNET_REGISTER_SUBGRAPH_PROPERTY(MKLDNN_QUANTIZE, SgMKLDNNTransformerProperty); From 2c6bba4495ea2692aab486421a5815f5f7b71d5e Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 3 Feb 2022 12:39:46 +0100 Subject: [PATCH 2/4] rewrite test --- tests/python/mkl/test_subgraph.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/python/mkl/test_subgraph.py b/tests/python/mkl/test_subgraph.py index a6af10f3d8cf..a948e8c66f5a 100644 --- a/tests/python/mkl/test_subgraph.py +++ b/tests/python/mkl/test_subgraph.py @@ -710,6 +710,28 @@ def fc_eltwise(no_bias, data_shape, flatten=True, alg='relu'): return sym, attr +def fc_identity_eltwise(data_shape, identity_node): + attrs = {'sg_mkldnn_fully_connected_eltwise_0' : {'with_eltwise': 'true'}, + 'sg_mkldnn_fully_connected_eltwise_1' : {'with_eltwise': 'true'}} + data, fc1_weight = head_symbol(data_shape) + fc2_weight = mx.symbol.Variable('fc2_weight', dtype='float32') + + sym = mx.symbol.FullyConnected(name='fc1', data=data, weight=fc1_weight, num_hidden=64, + no_bias=True, flatten=True) + if identity_node == 'copy': + sym = mx.symbol.identity(sym) + else: + sym = mx.symbol.Dropout(sym) + sym = mx.symbol.Activation(sym, act_type='relu') + sym = mx.symbol.FullyConnected(name='fc2', data=sym, weight=fc2_weight, num_hidden=64, + no_bias=True, flatten=True) + if identity_node == 'copy': + sym = mx.symbol.identity(sym) + else: + sym = mx.symbol.Dropout(sym) + sym = mx.symbol.Activation(sym, act_type='relu') + return sym, attrs + def single_selfatt_qk(data_shape, nheads=16): attr = {'selfatt_qk': {}} data = mx.symbol.Variable('data', shape=data_shape, dtype='float32') @@ -1000,6 +1022,12 @@ def test_single_fc(): else: check_fusion(syms, dshape, attrs, check_quantization=False) +@with_seed() +def test_fc_eltwise_identity(): + for dshape, identity_node in itertools.product(DATA_SHAPE, ['copy', 'dropout']): + syms, attrs = fc_identity_eltwise(dshape, identity_node) + check_fusion(syms, dshape, attrs, check_quantization=False) + @with_seed() def test_fc_eltwise(): for dshape, no_bias, flatten, alg in itertools.product(DATA_SHAPE, From c2bfb709f2250e0713211ac0f31bd8754fe4372b Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 3 Feb 2022 13:56:38 +0100 Subject: [PATCH 3/4] fix sanity --- src/operator/subgraph/mkldnn/mkldnn_identity_property.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/operator/subgraph/mkldnn/mkldnn_identity_property.h b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h index 6473c9357dc2..12b80a11bce7 100644 --- a/src/operator/subgraph/mkldnn/mkldnn_identity_property.h +++ b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h @@ -170,4 +170,4 @@ class SgMKLDNNIdentityProperty : public SubgraphProperty { } // namespace mxnet #endif // if MXNET_USE_MKLDNN == 1 -#endif // MXNET_OPERATOR_SUBGRAPH_MKLDNN_MKLDNN_IDENTITY_PROPERTY_H_ \ No newline at end of file +#endif // MXNET_OPERATOR_SUBGRAPH_MKLDNN_MKLDNN_IDENTITY_PROPERTY_H_ From e38281a809babd5dbeacefbefb6a738a8e69478d Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 3 Feb 2022 18:11:45 +0100 Subject: [PATCH 4/4] remove clang warning --- src/operator/subgraph/mkldnn/mkldnn_identity_property.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/operator/subgraph/mkldnn/mkldnn_identity_property.h b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h index 12b80a11bce7..00f949974343 100644 --- a/src/operator/subgraph/mkldnn/mkldnn_identity_property.h +++ b/src/operator/subgraph/mkldnn/mkldnn_identity_property.h @@ -151,7 +151,7 @@ class SgMKLDNNIdentityProperty : public SubgraphProperty { // (n) slice // \--index 1--> Dropout --index 0--> OUT_NODE // for OUT_NODE index 0 must be changed to index 1 - for (int i = 0; i < output_entries->size(); ++i) { + for (size_t i = 0; i < output_entries->size(); ++i) { auto out_node = output_entries->at(i)->node; if (IsIdentityNode(out_node)) { output_entries->at(i)->index = out_node->inputs[0].index;