[FEATURE] Add oneDNN support for numpy concatenate operator (#20652) · apache/mxnet@94fc557 · GitHub
Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 94fc557

Browse files
agrabowagrabow
authored
[FEATURE] Add oneDNN support for numpy concatenate operator (#20652)
* Integrate oneDNN support for numpy concatenate operator * Fix sanity check errors * Remove redundant npi_concatenate nnvm_register_op for gpu * Fix sanity check * Fix gpu register op and dmlc::optional problems * Fix GPU build
1 parent 1a1464f commit 94fc557

14 files changed

Lines changed: 85 additions & 202 deletions

File tree

‎cpp-package/example/charRNN.cpp‎

Lines changed: 2 additions & 1 deletion

‎python/mxnet/amp/lists/symbol_fp16.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -646,7 +646,6 @@
646646
'_mod',
647647
'_not_equal',
648648
'_npi_column_stack',
649-
'_npi_concatenate',
650649
'_npi_copysign',
651650
'_npi_cross',
652651
'_npi_dot',

‎src/api/operator/numpy/np_matrix_op.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -139,17 +139,17 @@ MXNET_REGISTER_API("_npi.concatenate")
139139
using namespace runtime;
140140
const nnvm::Op* op = Op::Get("_npi_concatenate");
141141
nnvm::NodeAttrs attrs;
142-
op::NumpyConcatenateParam param;
142+
op::ConcatParam param;
143143
int arg_size = args.num_args;
144144
param.num_args = arg_size - 2;
145145
if (args[arg_size - 2].type_code() == kNull) {
146-
param.axis = dmlc::nullopt;
146+
param.dim = dmlc::nullopt;
147147
} else {
148-
param.axis = args[arg_size - 2].operator int();
148+
param.dim = args[arg_size - 2].operator int();
149149
}
150150
attrs.parsed = param;
151151
attrs.op = op;
152-
SetAttrDict<op::NumpyConcatenateParam>(&attrs);
152+
SetAttrDict<op::ConcatParam>(&attrs);
153153
int num_inputs = arg_size - 2;
154154
std::vector<NDArray*> inputs;
155155
inputs.reserve(num_inputs);

‎src/operator/nn/concat-inl.h‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,12 @@ enum ConcatOpOutputs { kOut };
4747

4848
struct ConcatParam : public dmlc::Parameter<ConcatParam> {
4949
int num_args;
50-
int dim;
50+
dmlc::optional<int> dim;
5151
DMLC_DECLARE_PARAMETER(ConcatParam) {
5252
DMLC_DECLARE_FIELD(num_args).set_lower_bound(1).describe("Number of inputs to be concated.");
53-
DMLC_DECLARE_FIELD(dim).set_default(1).describe("the dimension to be concated.");
53+
DMLC_DECLARE_FIELD(dim)
54+
.set_default(dmlc::optional<int>(1))
55+
.describe("the dimension to be concated.");
5456
}
5557
void SetAttrDict(std::unordered_map<std::string, std::string>* dict) {
5658
std::ostringstream num_args_s, dim_s;
@@ -66,7 +68,7 @@ class ConcatOp {
6668
public:
6769
void Init(const ConcatParam& param) {
6870
this->size_ = param.num_args;
69-
this->dimension_ = param.dim;
71+
this->dimension_ = param.dim.has_value() ? param.dim.value() : 0;
7072
}
7173

7274
void Forward(const OpContext& ctx,
@@ -140,10 +142,18 @@ void ConcatCompute(const nnvm::NodeAttrs& attrs,
140142
const std::vector<OpReqType>& req,
141143
const std::vector<TBlob>& outputs) {
142144
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
143-
MSHADOW_TYPE_SWITCH(inputs[concat_enum::kData0].type_flag_, DType, {
145+
std::vector<TBlob> in_data(param.num_args);
146+
for (int i = 0; i < param.num_args; i++) {
147+
if (!param.dim.has_value()) {
148+
in_data[i] = inputs[i].reshape(mxnet::TShape(1, inputs[i].shape_.Size()));
149+
} else {
150+
in_data[i] = inputs[i];
151+
}
152+
}
153+
MSHADOW_TYPE_SWITCH_WITH_BOOL(in_data[concat_enum::kData0].type_flag_, DType, {
144154
ConcatOp<xpu, DType> op;
145155
op.Init(param);
146-
op.Forward(ctx, inputs, req, outputs);
156+
op.Forward(ctx, in_data, req, outputs);
147157
});
148158
}
149159

@@ -209,10 +219,18 @@ void ConcatGradCompute(const nnvm::NodeAttrs& attrs,
209219
const std::vector<OpReqType>& req,
210220
const std::vector<TBlob>& outputs) {
211221
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
212-
MSHADOW_TYPE_SWITCH(inputs[concat_enum::kOut].type_flag_, DType, {
222+
std::vector<TBlob> out_data(param.num_args);
223+
for (int i = 0; i < param.num_args; i++) {
224+
if (!param.dim.has_value()) {
225+
out_data[i] = outputs[i].reshape(mxnet::TShape(1, outputs[i].shape_.Size()));
226+
} else {
227+
out_data[i] = outputs[i];
228+
}
229+
}
230+
MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[concat_enum::kOut].type_flag_, DType, {
213231
ConcatOp<xpu, DType> op;
214232
op.Init(param);
215-
op.Backward(ctx, inputs[concat_enum::kOut], req, outputs);
233+
op.Backward(ctx, inputs[concat_enum::kOut], req, out_data);
216234
});
217235
}
218236

@@ -318,7 +336,7 @@ void ConcatCSRImpl(const nnvm::NodeAttrs& attrs,
318336
using namespace csr;
319337
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
320338
int num_args = param.num_args;
321-
int concat_dim = param.dim;
339+
int concat_dim = param.dim.has_value() ? param.dim.value() : 0;
322340
CHECK_EQ(inputs.size(), num_args);
323341
CHECK_EQ(outputs.size(), 1);
324342
int axis = CheckAxis(concat_dim, inputs[0].shape().ndim());

‎src/operator/nn/concat.cc‎

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,18 @@ bool ConcatShape(const nnvm::NodeAttrs& attrs,
3939
CHECK_EQ(in_shape->size(), static_cast<size_t>(param_.num_args));
4040
mxnet::TShape dshape;
4141
dim_t size = 0;
42+
int param_dim = param_.dim.has_value() ? param_.dim.value() : 0;
4243
bool has_unknown_dim_size = false;
4344
int axis = -1;
45+
if (!param_.dim.has_value()) {
46+
for (int i = 0; i < param_.num_args; ++i) {
47+
(*in_shape)[i] = Shape1((*in_shape)[i].Size());
48+
}
49+
}
4450
for (int i = 0; i < param_.num_args; ++i) {
4551
mxnet::TShape tmp = (*in_shape)[i];
4652
if (tmp.ndim() > 0) {
47-
axis = CheckAxis(param_.dim, tmp.ndim());
53+
axis = CheckAxis(param_dim, tmp.ndim());
4854
has_unknown_dim_size = !mxnet::dim_size_is_known(tmp, axis) || has_unknown_dim_size;
4955
size += tmp[axis];
5056
tmp[axis] = -1;
@@ -54,7 +60,7 @@ bool ConcatShape(const nnvm::NodeAttrs& attrs,
5460

5561
mxnet::TShape tmp = (*out_shape)[0];
5662
if (tmp.ndim() > 0) {
57-
axis = CheckAxis(param_.dim, tmp.ndim());
63+
axis = CheckAxis(param_dim, tmp.ndim());
5864
tmp[axis] = -1;
5965
shape_assign(&dshape, tmp);
6066
}
@@ -89,11 +95,12 @@ static bool RNNParamConcatShape(const nnvm::NodeAttrs& attrs,
8995
mxnet::TShape dshape;
9096
index_t size = 0;
9197
std::vector<int> zero_indices;
92-
int axis = -1;
98+
int axis = -1;
99+
int param_dim = param_.dim.has_value() ? param_.dim.value() : 0;
93100
for (int i = 0; i < param_.num_args; ++i) {
94101
mxnet::TShape tmp = (*in_shape)[i];
95102
if (tmp.ndim() > 0) {
96-
axis = CheckAxis(param_.dim, tmp.ndim());
103+
axis = CheckAxis(param_dim, tmp.ndim());
97104
if (!mxnet::dim_size_is_known(tmp, axis)) {
98105
zero_indices.emplace_back(i);
99106
} else {
@@ -107,7 +114,7 @@ static bool RNNParamConcatShape(const nnvm::NodeAttrs& attrs,
107114

108115
mxnet::TShape tmp = (*out_shape)[0];
109116
if (tmp.ndim() > 0) {
110-
axis = CheckAxis(param_.dim, tmp.ndim());
117+
axis = CheckAxis(param_dim, tmp.ndim());
111118
tmp[axis] = -1;
112119
shape_assign(&dshape, tmp);
113120
}
@@ -193,13 +200,14 @@ inline static bool ConcatForwardInferStorageType(const nnvm::NodeAttrs& attrs,
193200
auto& out_stype = out_attrs->at(0);
194201
bool dispatched = false;
195202
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
196-
if (!dispatched && common::ContainsOnlyStorage(*in_attrs, kCSRStorage) && param.dim == 0) {
203+
int param_dim = param.dim.has_value() ? param.dim.value() : 0;
204+
if (!dispatched && common::ContainsOnlyStorage(*in_attrs, kCSRStorage) && param_dim == 0) {
197205
dispatched =
198206
storage_type_assign(&out_stype, kCSRStorage, dispatch_mode, DispatchMode::kFComputeEx);
199207
}
200208
#if MXNET_USE_ONEDNN == 1
201209
if (!dispatched && dev_mask == mshadow::cpu::kDevMask &&
202-
common::ContainsOnlyStorage(*in_attrs, kDefaultStorage) && param.dim > 0) {
210+
common::ContainsOnlyStorage(*in_attrs, kDefaultStorage)) {
203211
dispatched =
204212
storage_type_assign(&out_stype, kDefaultStorage, dispatch_mode, DispatchMode::kFComputeEx);
205213
}
@@ -225,10 +233,8 @@ inline static bool BackwardConcatStorageType(const nnvm::NodeAttrs& attrs,
225233
std::vector<int>* out_attrs) {
226234
DispatchMode wanted_mode;
227235
#if MXNET_USE_ONEDNN == 1
228-
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
229236
CHECK_EQ(out_attrs->size(), in_attrs->size() - 1);
230-
if (dev_mask == mshadow::cpu::kDevMask &&
231-
common::ContainsOnlyStorage(*in_attrs, kDefaultStorage) && param.dim > 0)
237+
if (dev_mask == mshadow::cpu::kDevMask && common::ContainsOnlyStorage(*in_attrs, kDefaultStorage))
232238
wanted_mode = DispatchMode::kFComputeEx;
233239
else
234240
#endif // MXNET_USE_ONEDNN == 1
@@ -251,8 +257,9 @@ bool SupportDNNLConcat(const std::vector<NDArray>& arrs) {
251257
return false;
252258
int ndim = arr.shape().ndim();
253259
const int dnnl_ndims = arr.GetDNNLData()->get_desc().data.ndims;
254-
if (!(ndim == 2 || ndim == 4) || ndim != dnnl_ndims)
260+
if ((ndim != 2 && ndim != 4) || ndim != dnnl_ndims) {
255261
return false;
262+
}
256263
}
257264
return true;
258265
}
@@ -347,12 +354,14 @@ DMLC_REGISTER_PARAMETER(ConcatParam);
347354
NNVM_REGISTER_OP(Concat)
348355
MXNET_ADD_SPARSE_OP_ALIAS(concat)
349356
.add_alias("concat")
357+
.add_alias("_npi_concatenate")
350358
.describe(R"code(Joins input arrays along a given axis.
351359
352360
.. note:: `Concat` is deprecated. Use `concat` instead.
353361
354362
The dimensions of the input arrays should be the same except the axis along
355-
which they will be concatenated.
363+
which they will be concatenated. With dimension parameter ``None`` input
364+
arrays are flattened before concatenating them along axis 0.
356365
The dimension of the output array along the concatenated axis will be equal
357366
to the sum of the corresponding dimensions of the input arrays.
358367
@@ -376,6 +385,11 @@ Example::
376385
[ 7., 7.],
377386
[ 8., 8.]]
378387
388+
concat(x,y,z,dim=None) = [1., 1., 2., 2.,
389+
3., 3., 4., 4.,
390+
5., 5., 6., 6.,
391+
7., 7., 8., 8.]
392+
379393
Note that you cannot concat x,y,z along dimension 1 since dimension
380394
0 is not the same for all the input arrays.
381395
@@ -397,6 +411,7 @@ Example::
397411
.add_arguments(ConcatParam::__FIELDS__());
398412

399413
NNVM_REGISTER_OP(_backward_Concat)
414+
.add_alias("_backward_np_concat")
400415
.set_num_inputs([](const NodeAttrs& attrs) {
401416
#if MXNET_USE_ONEDNN == 1
402417
const ConcatParam& params = nnvm::get<ConcatParam>(attrs.parsed);

‎src/operator/nn/concat.cu‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,14 +47,17 @@ static void ConcatComputeExGPU(const nnvm::NodeAttrs& attrs,
4747
}
4848

4949
NNVM_REGISTER_OP(Concat)
50+
.add_alias("_npi_concatenate")
5051
.set_attr<FCompute>("FCompute<gpu>", ConcatCompute<gpu>)
5152
.set_attr<FComputeEx>("FComputeEx<gpu>", ConcatComputeExGPU);
5253

5354
NNVM_REGISTER_OP(_rnn_param_concat)
5455
.set_attr<FCompute>("FCompute<gpu>", ConcatCompute<gpu>)
5556
.set_attr<FComputeEx>("FComputeEx<gpu>", ConcatComputeExGPU);
5657

57-
NNVM_REGISTER_OP(_backward_Concat).set_attr<FCompute>("FCompute<gpu>", ConcatGradCompute<gpu>);
58+
NNVM_REGISTER_OP(_backward_Concat)
59+
.add_alias("_backward_np_concat")
60+
.set_attr<FCompute>("FCompute<gpu>", ConcatGradCompute<gpu>);
5861

5962
} // namespace op
6063
} // namespace mxnet

‎src/operator/nn/dnnl/dnnl_concat.cc‎

Lines changed: 5 additions & 3 deletions

0 commit comments

Comments
 (0)