diff --git a/inference-engine/include/ie_precision.hpp b/inference-engine/include/ie_precision.hpp index 2c8f3fe8e21..8d13a4bab04 100644 --- a/inference-engine/include/ie_precision.hpp +++ b/inference-engine/include/ie_precision.hpp @@ -224,7 +224,7 @@ public: (precisionInfo.value == Precision::Q78) || (precisionInfo.value == Precision::I16) || (precisionInfo.value == Precision::I8) || (precisionInfo.value == Precision::I32) || (precisionInfo.value == Precision::I64) || (precisionInfo.value == Precision::BIN) || - (precisionInfo.value == Precision::CUSTOM); + (precisionInfo.value == Precision::BF16) || (precisionInfo.value == Precision::CUSTOM); } protected: diff --git a/inference-engine/src/mkldnn_plugin/bf16transformer.h b/inference-engine/src/mkldnn_plugin/bf16transformer.h index 370656eb234..6ff30cdcae3 100644 --- a/inference-engine/src/mkldnn_plugin/bf16transformer.h +++ b/inference-engine/src/mkldnn_plugin/bf16transformer.h @@ -13,7 +13,7 @@ namespace MKLDNNPlugin { class BF16Transformer { const InferenceEngine::details::caseless_set _initbf16 = - { "convolution", "fullyconnected", "innerproduct" }; + { "convolution", "fullyconnected", "innerproduct", "gemm" }; const InferenceEngine::details::caseless_set _complementbf16 = { "relu", "tanh", "elu", "square", "abs", "sqrt", "linear", "bounded_relu", "soft_relu", "logistic", "exp", "gelu", "clamp", "swish", "prelu", "pooling", "norm", "gather", "memory" }; diff --git a/inference-engine/src/mkldnn_plugin/nodes/mkldnn_gemm_node.cpp b/inference-engine/src/mkldnn_plugin/nodes/mkldnn_gemm_node.cpp index 2ee3523e229..f1df0912108 100644 --- a/inference-engine/src/mkldnn_plugin/nodes/mkldnn_gemm_node.cpp +++ b/inference-engine/src/mkldnn_plugin/nodes/mkldnn_gemm_node.cpp @@ -122,8 +122,13 @@ void MKLDNNGemmNode::initSupportedPrimitiveDescriptors() { auto inPrec0 = getCnnLayer()->insData[0].lock()->getPrecision(); auto inPrec1 = getCnnLayer()->insData[1].lock()->getPrecision(); if ((inPrec0 != Precision::U8 && inPrec0 != Precision::I8) || inPrec1 != Precision::I8 || isThreeInputs) { - inPrec0 = Precision::FP32; - inPrec1 = Precision::FP32; + if (inPrec0 == Precision::BF16 || inPrec1 == Precision::BF16) { + inPrec0 = Precision::BF16; + inPrec1 = Precision::BF16; + } else { + inPrec0 = Precision::FP32; + inPrec1 = Precision::FP32; + } } auto inputDataType0 = MKLDNNExtensionUtils::IEPrecisionToDataType(inPrec0); @@ -192,6 +197,11 @@ inline void process_gemm(char transa, char transb, int M, int N, int K, float al mkldnn_sgemm(transa, transb, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc); } +inline void process_gemm(char transa, char transb, int M, int N, int K, float alpha, const uint16_t *A, int lda, + const uint16_t *B, int ldb, float beta, float *C, int ldc) { + mkldnn_gemm_bf16bf16f32(transa, transb, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc); +} + inline void process_gemm(char transa, char transb, int M, int N, int K, float alpha, const uint8_t *A, int lda, const int8_t *B, int ldb, float beta, float *C, int ldc) { const int32_t co = 0; @@ -288,6 +298,9 @@ void MKLDNNGemmNode::execute(mkldnn::stream strm) { case Precision::FP32: process_data(); break; + case Precision::BF16: + process_data(); + break; case Precision::I8: process_data(); break; diff --git a/inference-engine/thirdparty/mkl-dnn b/inference-engine/thirdparty/mkl-dnn index 4f511de56e2..b96a54762ae 160000 --- a/inference-engine/thirdparty/mkl-dnn +++ b/inference-engine/thirdparty/mkl-dnn @@ -1 +1 @@ -Subproject commit 4f511de56e21b417f7c49c3f50cf7217350412ab +Subproject commit b96a54762aedf27711f8e9144d05b2697b03cc40 diff --git a/model-optimizer/automation/package_BOM.txt b/model-optimizer/automation/package_BOM.txt index d41c3d77d25..1e1c23c8d2b 100644 --- a/model-optimizer/automation/package_BOM.txt +++ b/model-optimizer/automation/package_BOM.txt @@ -119,6 +119,7 @@ extensions/front/create_tensor_nodes.py extensions/front/disable_weights_quantize_value_propagation.py extensions/front/div.py extensions/front/eltwise_n.py +extensions/front/EmbeddingBagFuse.py extensions/front/ExpandDimsToUnsqueeze.py extensions/front/FillToBroadcast.py extensions/front/flatten_to_reshape.py diff --git a/model-optimizer/extensions/front/EmbeddingBagFuse.py b/model-optimizer/extensions/front/EmbeddingBagFuse.py new file mode 100644 index 00000000000..f193f94dca4 --- /dev/null +++ b/model-optimizer/extensions/front/EmbeddingBagFuse.py @@ -0,0 +1,91 @@ +""" + Copyright (C) 2020 Intel Corporation + + Licensed 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. +""" + +from extensions.ops.embedding_bag import EmbeddingBagOffsetsSum +from mo.front.common.replacement import FrontReplacementSubgraph +from mo.graph.graph import Graph, rename_nodes + + +class EmbeddingBagFuse(FrontReplacementSubgraph): + enabled = True + + def run_after(self): + from extensions.front.ExpandDimsToUnsqueeze import ExpandDimsToUnsqueeze + from extensions.front.AttributedGatherNormalizer import AttributedGatherNormalizer + return [ExpandDimsToUnsqueeze, AttributedGatherNormalizer] + + def pattern(self): + return dict( + nodes=[ + ('weights', dict(op='Const')), + ('concat_before', dict(op='Concat')), + ('gather_before1_1', dict(op='Gather')), + ('unsqueeze_before1_1', dict(op='Unsqueeze')), + ('gather_before2_1', dict(op='Gather')), + ('unsqueeze_before2_1', dict(op='Unsqueeze')), + ('slice1', dict(op='Slice')), + ('gather_after1', dict(op='Gather')), + ('reduce1', dict(op='ReduceSum')), + ('unsqueeze_after1', dict(op='Unsqueeze')), + ('concat_after', dict(op='Concat')), + ], + edges=[ + ('concat_before', 'gather_before1_1'), + ('concat_before', 'gather_before2_1'), + ('gather_before1_1', 'unsqueeze_before1_1'), + ('gather_before2_1', 'unsqueeze_before2_1'), + ('unsqueeze_before1_1', 'slice1', {'out': 0, 'in': 1}), + ('unsqueeze_before2_1', 'slice1', {'out': 0, 'in': 2}), + ('weights', 'gather_after1', {'out': 0, 'in': 0}), + ('slice1', 'gather_after1', {'out': 0, 'in': 1}), + ('gather_after1', 'reduce1'), + ('reduce1', 'unsqueeze_after1'), + ('unsqueeze_after1', 'concat_after'), + ]) + + def replace_sub_graph(self, graph: Graph, match: dict): + concat_before = match['concat_before'] + gather_after1 = match['gather_after1'] + slice1 = match['slice1'] + concat_after = match['concat_after'] + weights_node = gather_after1.in_port(0).get_source().node + gather_after_axis = gather_after1.in_port(2).get_source().node.soft_get('value') + for dst_port in weights_node.out_port(0).get_destinations(): + node = dst_port.node + if node.op == 'Gather': + # validate that all Gathers have same axis + if node.in_port(2).get_source().node.soft_get('value') != gather_after_axis: + return + dst_port.disconnect() + indices_node = slice1.in_port(0).get_source().node + slice_axis = slice1.in_port(3).get_source().node.soft_get('value') + for dst_port in indices_node.out_port(0).get_destinations(): + node = dst_port.node + if node.op == 'Slice': + # validate that all Slices have same axis + if node.in_port(3).get_source().node.soft_get('value') != slice_axis: + return + dst_port.disconnect() + emb_bag = EmbeddingBagOffsetsSum(graph, {}).create_node() + weights_node.out_port(0).connect(emb_bag.in_port(0)) + indices_node.out_port(0).connect(emb_bag.in_port(1)) + concat_before.in_port(0).get_connection().set_destination(emb_bag.in_port(2)) + concat_after.out_port(0).get_connection().set_source(emb_bag.out_port(0)) + concat_name = concat_after.soft_get('name', concat_after.id) + rename_nodes([(concat_after, concat_name + '/TBD'), (emb_bag, concat_name)]) + + # remove this sub-graph since a lot of matchings will be obsolete + graph.remove_nodes_from(graph.dfs(concat_before.id, set()))