| /* |
| * 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 elemwise_binary_broadcast_op_basic.cc |
| * \brief CPU Implementation of basic functions for elementwise binary broadcast operator. |
| */ |
| #include "operator/tensor/elemwise_unary_op.h" |
| #include "operator/tensor/elemwise_binary_op-inl.h" |
| #include "operator/tensor/elemwise_binary_broadcast_op.h" |
| #if MXNET_USE_ONEDNN == 1 |
| #include "operator/nn/dnnl/dnnl_binary-inl.h" |
| #include "operator/nn/dnnl/dnnl_sum-inl.h" |
| #endif // MXNET_USE_ONEDNN == 1 |
| |
| namespace mxnet { |
| namespace op { |
| |
| #if MXNET_USE_ONEDNN == 1 |
| template <dnnl::algorithm alg> |
| void DNNLBinaryOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| const mxnet::TShape& input_0_shape = inputs[0].shape(); |
| const mxnet::TShape& input_1_shape = inputs[1].shape(); |
| const mxnet::TShape& output_0_shape = outputs[0].shape(); |
| // We can use more efficient sum kernel, when there is no broadcast - when shapes are the |
| // same. |
| const bool same_shape = (input_0_shape == input_1_shape); |
| |
| if (same_shape && alg == dnnl::algorithm::binary_add) { |
| DNNLSumFwd& fwd = DNNLSumFwd::GetCached(inputs, outputs); |
| fwd.Execute(ctx, inputs, req, outputs); |
| } else { |
| mxnet::TShape new_lshape, new_rshape, new_oshape; |
| int ndim_diff = BinaryBroadcastShapeCompact( |
| input_0_shape, input_1_shape, output_0_shape, &new_lshape, &new_rshape, &new_oshape); |
| std::vector<NDArray> new_inputs; |
| std::vector<NDArray> new_outputs; |
| if (ndim_diff) { |
| new_inputs = {inputs[0].Reshape(new_lshape), inputs[1].Reshape(new_rshape)}; |
| new_outputs = {outputs[0].Reshape(new_oshape)}; |
| } else if (input_0_shape.Size() == 1 && input_1_shape.Size() == 1) { |
| // BinaryBroadcastShapeCompact function doesn't reshape tensors of size (1,1,...,1) |
| // into shape (1). It is mandatory for oneDNN primitive to have this reshape done. |
| mxnet::TShape one_shape = mxnet::TShape(1, 1); |
| new_inputs = {inputs[0].Reshape(one_shape), inputs[1].Reshape(one_shape)}; |
| new_outputs = {outputs[0].Reshape(one_shape)}; |
| } else { |
| new_inputs = {inputs[0], inputs[1]}; |
| new_outputs = {outputs[0]}; |
| } |
| |
| DNNLBinaryOpFwd& fwd = DNNLBinaryOpFwd::GetBinaryOpForward<alg>(new_inputs, new_outputs); |
| fwd.Execute(new_inputs, req, new_outputs); |
| } |
| } |
| #endif |
| |
| template <typename OP> |
| static void BinaryOperatorComputeExCPU(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| #if MXNET_USE_ONEDNN == 1 |
| if (common::ContainsOnlyStorage(inputs, kDefaultStorage)) { |
| if (SupportDNNLBinary(inputs, outputs)) { |
| const dnnl::algorithm alg = DNNLAlgorithm<OP>::value; |
| DNNLRun(DNNLBinaryOpForward<alg>, attrs, ctx, inputs, req, outputs); |
| } else { |
| FallBackCompute(BinaryBroadcastCompute<cpu, OP>, attrs, ctx, inputs, req, outputs); |
| } |
| return; |
| } |
| #endif // MXNET_USE_ONEDNN == 1 |
| if (std::is_same<OP, op::mshadow_op::plus>::value || |
| std::is_same<OP, op::mshadow_op::minus>::value) { |
| BinaryBroadcastComputeDenseEx<cpu, OP>(attrs, ctx, inputs, req, outputs); |
| } else if (std::is_same<OP, op::mshadow_op::mul>::value || |
| std::is_same<OP, op::mshadow_op::div>::value) { |
| BinaryBroadcastComputeSparseEx<cpu, OP>(attrs, ctx, inputs, req, outputs); |
| } |
| } |
| |
| MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(broadcast_add) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_add) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_plus) |
| .add_alias("broadcast_plus") |
| .describe(R"code(Returns element-wise sum of the input arrays with broadcasting. |
| |
| `broadcast_plus` is an alias to the function `broadcast_add`. |
| |
| Example:: |
| |
| x = [[ 1., 1., 1.], |
| [ 1., 1., 1.]] |
| |
| y = [[ 0.], |
| [ 1.]] |
| |
| broadcast_add(x, y) = [[ 1., 1., 1.], |
| [ 2., 2., 2.]] |
| |
| broadcast_plus(x, y) = [[ 1., 1., 1.], |
| [ 2., 2., 2.]] |
| |
| Supported sparse operations: |
| |
| broadcast_add(csr, dense(1D)) = dense |
| broadcast_add(dense(1D), csr) = dense |
| |
| )code" ADD_FILELINE) |
| .set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, op::mshadow_op::plus>) |
| .set_attr<FComputeEx>("FComputeEx<cpu>", BinaryOperatorComputeExCPU<op::mshadow_op::plus>) |
| .set_attr<FInferStorageType>("FInferStorageType", BinaryBroadcastAddStorageType) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_broadcast_add"}); |
| |
| NNVM_REGISTER_OP(_backward_broadcast_add) |
| .set_num_inputs(1) |
| .set_num_outputs(2) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int> >{{0, 0}, {0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>( |
| "FCompute<cpu>", |
| BinaryBroadcastBackwardUseNone<cpu, mshadow_op::identity, mshadow_op::identity>); |
| |
| MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(broadcast_sub) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_sub) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_minus) |
| .add_alias("broadcast_minus") |
| .describe(R"code(Returns element-wise difference of the input arrays with broadcasting. |
| |
| `broadcast_minus` is an alias to the function `broadcast_sub`. |
| |
| Example:: |
| |
| x = [[ 1., 1., 1.], |
| [ 1., 1., 1.]] |
| |
| y = [[ 0.], |
| [ 1.]] |
| |
| broadcast_sub(x, y) = [[ 1., 1., 1.], |
| [ 0., 0., 0.]] |
| |
| broadcast_minus(x, y) = [[ 1., 1., 1.], |
| [ 0., 0., 0.]] |
| |
| Supported sparse operations: |
| |
| broadcast_sub/minus(csr, dense(1D)) = dense |
| broadcast_sub/minus(dense(1D), csr) = dense |
| |
| )code" ADD_FILELINE) |
| .set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, op::mshadow_op::minus>) |
| .set_attr<FComputeEx>("FComputeEx<cpu>", BinaryOperatorComputeExCPU<op::mshadow_op::minus>) |
| .set_attr<FInferStorageType>("FInferStorageType", BinaryBroadcastAddStorageType) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_broadcast_sub"}); |
| |
| NNVM_REGISTER_OP(_backward_broadcast_sub) |
| .set_num_inputs(1) |
| .set_num_outputs(2) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int> >{{0, 0}, {0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>( |
| "FCompute<cpu>", |
| BinaryBroadcastBackwardUseNone<cpu, mshadow_op::identity, mshadow_op::negation>); |
| |
| MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(broadcast_mul) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_mul) |
| .describe(R"code(Returns element-wise product of the input arrays with broadcasting. |
| |
| Example:: |
| |
| x = [[ 1., 1., 1.], |
| [ 1., 1., 1.]] |
| |
| y = [[ 0.], |
| [ 1.]] |
| |
| broadcast_mul(x, y) = [[ 0., 0., 0.], |
| [ 1., 1., 1.]] |
| |
| Supported sparse operations: |
| |
| broadcast_mul(csr, dense(1D)) = csr |
| |
| )code" ADD_FILELINE) |
| .set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, op::mshadow_op::mul>) |
| .set_attr<FComputeEx>("FComputeEx<cpu>", BinaryOperatorComputeExCPU<op::mshadow_op::mul>) |
| .set_attr<FInferStorageType>("FInferStorageType", BinaryBroadcastMulStorageType) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_broadcast_mul"}); |
| |
| NNVM_REGISTER_OP(_backward_broadcast_mul) |
| .set_num_inputs(3) |
| .set_num_outputs(2) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int> >{{0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>("FCompute<cpu>", |
| BinaryBroadcastBackwardUseIn<cpu, mshadow_op::right, mshadow_op::left>); |
| |
| MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(broadcast_div) |
| MXNET_ADD_SPARSE_OP_ALIAS(broadcast_div) |
| .describe(R"code(Returns element-wise division of the input arrays with broadcasting. |
| |
| Example:: |
| |
| x = [[ 6., 6., 6.], |
| [ 6., 6., 6.]] |
| |
| y = [[ 2.], |
| [ 3.]] |
| |
| broadcast_div(x, y) = [[ 3., 3., 3.], |
| [ 2., 2., 2.]] |
| |
| Supported sparse operations: |
| |
| broadcast_div(csr, dense(1D)) = csr |
| |
| )code" ADD_FILELINE) |
| .set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, op::mshadow_op::div>) |
| .set_attr<FComputeEx>("FComputeEx<cpu>", BinaryOperatorComputeExCPU<op::mshadow_op::div>) |
| .set_attr<FInferStorageType>("FInferStorageType", BinaryBroadcastMulStorageType) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_broadcast_div"}); |
| |
| NNVM_REGISTER_OP(_backward_broadcast_div) |
| .set_num_inputs(3) |
| .set_num_outputs(2) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int> >{{0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>( |
| "FCompute<cpu>", |
| BinaryBroadcastBackwardUseIn<cpu, mshadow_op::div_grad, mshadow_op::div_rgrad>); |
| |
| MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(broadcast_mod) |
| .describe(R"code(Returns element-wise modulo of the input arrays with broadcasting. |
| |
| Example:: |
| |
| x = [[ 8., 8., 8.], |
| [ 8., 8., 8.]] |
| |
| y = [[ 2.], |
| [ 3.]] |
| |
| broadcast_mod(x, y) = [[ 0., 0., 0.], |
| [ 2., 2., 2.]] |
| |
| )code" ADD_FILELINE) |
| .set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, mshadow_op::mod>) |
| .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_broadcast_mod"}); |
| |
| NNVM_REGISTER_OP(_backward_broadcast_mod) |
| .set_num_inputs(3) |
| .set_num_outputs(2) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>("FInplaceOption", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::pair<int, int> >{{0, 1}}; |
| }) |
| .set_attr<FResourceRequest>("FResourceRequest", |
| [](const NodeAttrs& attrs) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| .set_attr<FCompute>( |
| "FCompute<cpu>", |
| BinaryBroadcastBackwardUseIn<cpu, mshadow_op::mod_grad, mshadow_op::mod_rgrad>); |
| |
| } // namespace op |
| } // namespace mxnet |
