You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2021/08/20 12:30:26 UTC
[GitHub] [incubator-mxnet] bgawrych commented on a change in pull request #19562: [operator] Integrate oneDNN layer normalization implementation
bgawrych commented on a change in pull request #19562:
URL: https://github.com/apache/incubator-mxnet/pull/19562#discussion_r692904008
##########
File path: src/operator/nn/mkldnn/mkldnn_layer_norm.cc
##########
@@ -0,0 +1,261 @@
+/*
+ * 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_layer_norm.cc
+ */
+
+#if MXNET_USE_ONEDNN == 1
+
+#include "./mkldnn_layer_norm-inl.h"
+
+namespace mxnet {
+namespace op {
+
+bool SupportMKLDNNLayerNorm(const LayerNormParam& param, const std::vector<NDArray>& inputs) {
+ const mxnet::TShape& shape = inputs[layernorm::kData].shape();
+
+ // Native implementation (which can be found in function LayerNormCPU) is faster than oneDNN's one
+ // for small tensors. Below is the heuristic based on measurements on clx machine deciding whether
+ // the shape is better for oneDNN or native implementation.
+ auto ShapeBetterForMKLDNN = [](const mxnet::TShape& shape) {
+ constexpr size_t shapeLimit = 1024;
+ return shape.Size() / shape[0] >= shapeLimit && shape[0] >= shapeLimit;
+ };
+
+ return (ShapeBetterForMKLDNN(shape) &&
+ (GetRealAxis(param.axis, shape.ndim()) == shape.ndim() - 1) && (shape.ndim() >= 2) &&
+ (shape.ndim() <= 5) &&
+ (inputs[layernorm::kData].dtype() == mshadow::kFloat32 ||
+ inputs[layernorm::kData].dtype() == mshadow::kBfloat16) &&
+ inputs[layernorm::kGamma].dtype() == mshadow::kFloat32 &&
+ inputs[layernorm::kBeta].dtype() == mshadow::kFloat32);
+}
+
+void MKLDNNLayerNormForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const LayerNormParam& param = nnvm::get<LayerNormParam>(attrs.parsed);
+ const auto& fwd = MKLDNNLayerNormFwd::GetCached(param, ctx, inputs[layernorm::kData]);
+ fwd.Execute(param, ctx, inputs, req[layernorm::kOut], outputs);
+}
+
+MKLDNNLayerNormFwd& MKLDNNLayerNormFwd::GetCached(const LayerNormParam& param,
+ const OpContext& ctx,
+ const NDArray& data) {
+ using layernorm_fwd_map = std::unordered_map<LayerNormSignature, MKLDNNLayerNormFwd, OpHash>;
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local layernorm_fwd_map layer_norm_fwds;
+#else
+ static MX_THREAD_LOCAL layernorm_fwd_map layer_norm_fwds;
+#endif
+
+ LayerNormSignature key(param);
+ key.AddSign(data);
+ key.AddSign(param.eps);
Review comment:
Isn't eps already added by passing param to key?
##########
File path: src/operator/nn/mkldnn/mkldnn_layer_norm.cc
##########
@@ -0,0 +1,261 @@
+/*
+ * 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_layer_norm.cc
+ */
+
+#if MXNET_USE_ONEDNN == 1
+
+#include "./mkldnn_layer_norm-inl.h"
+
+namespace mxnet {
+namespace op {
+
+bool SupportMKLDNNLayerNorm(const LayerNormParam& param, const std::vector<NDArray>& inputs) {
+ const mxnet::TShape& shape = inputs[layernorm::kData].shape();
+
+ // Native implementation (which can be found in function LayerNormCPU) is faster than oneDNN's one
+ // for small tensors. Below is the heuristic based on measurements on clx machine deciding whether
+ // the shape is better for oneDNN or native implementation.
+ auto ShapeBetterForMKLDNN = [](const mxnet::TShape& shape) {
+ constexpr size_t shapeLimit = 1024;
+ return shape.Size() / shape[0] >= shapeLimit && shape[0] >= shapeLimit;
+ };
+
+ return (ShapeBetterForMKLDNN(shape) &&
+ (GetRealAxis(param.axis, shape.ndim()) == shape.ndim() - 1) && (shape.ndim() >= 2) &&
+ (shape.ndim() <= 5) &&
+ (inputs[layernorm::kData].dtype() == mshadow::kFloat32 ||
+ inputs[layernorm::kData].dtype() == mshadow::kBfloat16) &&
+ inputs[layernorm::kGamma].dtype() == mshadow::kFloat32 &&
+ inputs[layernorm::kBeta].dtype() == mshadow::kFloat32);
+}
+
+void MKLDNNLayerNormForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const LayerNormParam& param = nnvm::get<LayerNormParam>(attrs.parsed);
+ const auto& fwd = MKLDNNLayerNormFwd::GetCached(param, ctx, inputs[layernorm::kData]);
+ fwd.Execute(param, ctx, inputs, req[layernorm::kOut], outputs);
+}
+
+MKLDNNLayerNormFwd& MKLDNNLayerNormFwd::GetCached(const LayerNormParam& param,
+ const OpContext& ctx,
+ const NDArray& data) {
+ using layernorm_fwd_map = std::unordered_map<LayerNormSignature, MKLDNNLayerNormFwd, OpHash>;
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local layernorm_fwd_map layer_norm_fwds;
+#else
+ static MX_THREAD_LOCAL layernorm_fwd_map layer_norm_fwds;
+#endif
+
+ LayerNormSignature key(param);
+ key.AddSign(data);
+ key.AddSign(param.eps);
+
+ auto it = layer_norm_fwds.find(key);
+ if (it == layer_norm_fwds.end()) {
+ MKLDNNLayerNormFwd fwd(param, data);
+ it = AddToCache(&layer_norm_fwds, key, fwd);
+ }
+ return it->second;
+}
+
+MKLDNNLayerNormFwd::MKLDNNLayerNormFwd(const LayerNormParam& param, const NDArray& data) {
+ const mkldnn::memory::desc data_md = data.GetMKLDNNData()->get_desc();
+ fwd_pd = CreatePrimitiveDesc(param, data_md);
+ fwd = std::make_shared<mkldnn::layer_normalization_forward>(*fwd_pd);
Review comment:
Sometimes you're using defined types by using, sometimes directly from mkldnn namespace (layernorm_fwd_t) - be consistent ;)
--
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.
To unsubscribe, e-mail: commits-unsubscribe@mxnet.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org