Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
[v1.6.x] Backport [MKL-DNN] Integrate Conv3d and Pool3d/1d (#17884) a…
Browse files Browse the repository at this point in the history 
…nd Fix Sanity pipeline in 1.6.x (#18206)

* [MKL-DNN] Integrate Conv3d and Pool3d/1d (#17884)

* Integrate MKl-DNN conv3d and pool3d/1d

* fix UT & address comments

* clean code

* rebase against latest master

* pylint astroid sanity issue

* astroid and pylint versions only supported in py3

* remove kBFloat16 as its not supported int 1.6

* added missing definition GetPaddingSizeFull

* Remove dilation restriction for conv3d (#17491)

* Remove conv3d dilation restriction

* Remove comment

* fix unix-gpu test for num_outputs and inputs

Co-authored-by: Wuxun Zhang <[email protected]>
Co-authored-by: reminisce <[email protected]>
  • Loading branch information
@ChaiBapchya @wuxun-zhang @reminisce
3 people committed May 6, 2020
1 parent 1e129fa commit b748c45
Show file tree
Hide file tree
Showing 19 changed files with 515 additions and 296 deletions.
1 change: 1 addition & 0 deletions ci/docker/install/requirements
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ nose==1.3.7
nose-timer==0.7.3
numpy>1.16.0,<1.18.0
pylint==2.3.1; python_version >= '3.0'
astroid==2.3.3; python_version >= '3.0'
requests<2.19.0,>=2.18.4
scipy==1.2.1
six==1.11.0
4 changes: 0 additions & 4 deletions src/operator/nn/convolution.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,8 +223,6 @@ static bool ConvolutionShape(const nnvm::NodeAttrs& attrs,
SHAPE_ASSIGN_CHECK(*in_shape, conv::kBias, Shape1(param_.num_filter));
}

// Note: 3D dilation currently not supported.
// Calculations below done to preserve symmetry with 1D/2D code.
const index_t dilated_ksize_d = param_.DilatedKernelSize(0);
const index_t dilated_ksize_y = param_.DilatedKernelSize(1);
const index_t dilated_ksize_x = param_.DilatedKernelSize(2);
Expand All @@ -239,8 +237,6 @@ static bool ConvolutionShape(const nnvm::NodeAttrs& attrs,
<< "incorrect stride size: " << param_.stride;
CHECK_GT(param_.dilate.Size(), 0U) \
<< "incorrect dilate size: " << param_.dilate;
CHECK_EQ(param_.dilate.Size(), 1U)
<< "Dilate is not supported in 3d convolution";
Shape<5> oshape;
oshape[0] = dshape[0];
oshape[1] = param_.num_filter;
Expand Down
12 changes: 6 additions & 6 deletions src/operator/nn/mkldnn/mkldnn_act.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,10 @@ bool SupportMKLDNNAct(const ActivationParam& param) {
}

bool SupportMKLDNNAct(const ActivationParam& param, const NDArray &input) {
// MKL-DNN Activation supports 1d, 2d, 3d, 4d data layout
// MKL-DNN Activation supports 1d, 2d, 3d, 4d and 5d data layout
if ((input.shape().ndim() < 1) ||
(input.shape().ndim() > 4) ||
(input.dtype() != mshadow::kFloat32))
(input.shape().ndim() > 5) ||
!(input.dtype() == mshadow::kFloat32))
return false;
return SupportMKLDNNAct(param);
}
Expand All @@ -62,10 +62,10 @@ bool SupportMKLDNNLeakyRelu(const LeakyReLUParam& param) {
}

bool SupportMKLDNNLeakyRelu(const LeakyReLUParam& param, const NDArray &input) {
// MKL-DNN Activation supports 1d, 2d, 3d, 4d data layout
// MKL-DNN Activation supports 1d, 2d, 3d, 4d and 5d data layout
if ((input.shape().ndim() < 1) ||
(input.shape().ndim() > 4) ||
(input.dtype() != mshadow::kFloat32))
(input.shape().ndim() > 5) ||
!(input.dtype() == mshadow::kFloat32))
return false;
return SupportMKLDNNLeakyRelu(param);
}
Expand Down
49 changes: 27 additions & 22 deletions src/operator/nn/mkldnn/mkldnn_base-inl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,15 +129,8 @@ static inline bool SupportMKLDNN(int dtype, const mxnet::TShape &shape) {
// MKLDNN currently does not support 0-dim Tensor and 0-size Tensor
return false;
}
return dtype == mshadow::kFloat32 && (ndim == 1 || ndim == 2 || ndim == 4);
}

static inline bool SupportMKLDNNRnn(const NDArray &input) {
if (input.dtype() == mshadow::kFloat32 && input.shape().ndim() == 3
&& dmlc::GetEnv("MXNET_USE_MKLDNN_RNN", 1)) {
return true;
}
return false;
return (dtype == mshadow::kFloat32) &&
(ndim == 1 || ndim == 2 || ndim == 4);
}

static inline bool SupportMKLDNNQuantize(int dtype) {
Expand Down Expand Up @@ -302,20 +295,32 @@ inline static mkldnn::memory::desc GetWeightDesc(const NDArray &arr,
if (num_groups == 1) {
return GetMemDesc(arr, dtype);
} else {
auto ndim = arr.shape().ndim();
CHECK((ndim == 3) || (ndim == 4))
<< "MKL-DNN weight currectly supports 3d and 4d layout";
const auto ndim = arr.shape().ndim();
CHECK((ndim == 3) || (ndim == 4) || (ndim == 5))
<< "MKL-DNN weight currently supports 3d or 4d or 5d layout";
auto tz = mkldnn::memory::dims{0};
const int N = 0, H = 2, W = 3, C = 1;
if (ndim == 3) {
tz = mkldnn::memory::dims{
num_groups, static_cast<int>(arr.shape()[N] / num_groups),
static_cast<int>(arr.shape()[C]), static_cast<int>(arr.shape()[H])};
} else {
tz = mkldnn::memory::dims{
num_groups, static_cast<int>(arr.shape()[N] / num_groups),
static_cast<int>(arr.shape()[C]), static_cast<int>(arr.shape()[H]),
static_cast<int>(arr.shape()[W])};
int N = 0, C = 1, H = 2, W = 3;
int D = -1;
if (ndim == 5) {
D = 2;
H = 3;
W = 4;
}
switch (ndim) {
case 3:
tz = mkldnn::memory::dims{
num_groups, arr.shape()[N] / num_groups,
arr.shape()[C], arr.shape()[H]};
break;
case 4:
tz = mkldnn::memory::dims{
num_groups, arr.shape()[N] / num_groups,
arr.shape()[C], arr.shape()[H], arr.shape()[W]};
break;
case 5:
tz = mkldnn::memory::dims{
num_groups, arr.shape()[N] / num_groups,
arr.shape()[C], arr.shape()[D], arr.shape()[H], arr.shape()[W]};
}
return mkldnn::memory::desc{tz, get_mkldnn_type(dtype), mkldnn::memory::format_tag::any};
}
Expand Down
47 changes: 30 additions & 17 deletions src/operator/nn/mkldnn/mkldnn_base.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -240,31 +240,44 @@ const mkldnn::memory *GetWeights(const NDArray &arr, int num_groups) {
auto tz = mkldnn::memory::dims{0};
auto format_tag = mkldnn::memory::format_tag::undef;
auto engine = CpuEngine::Get()->get_engine();
const int O = 0, I = 1, H = 2, W = 3;
if (arr.shape().ndim() == 2) {
tz = mkldnn::memory::dims{static_cast<int>(arr.shape()[O]), static_cast<int>(arr.shape()[I])};
const int ndim = arr.shape().ndim();
int O = 0, I = 1, H = 2, W = 3;
int D = -1;
if (ndim == 5) {
D = 2;
H = 3;
W = 4;
}
if (ndim == 2) {
tz = mkldnn::memory::dims{arr.shape()[O], arr.shape()[I]};
format_tag = mkldnn::memory::format_tag::oi;
} else if (arr.shape().ndim() == 3) {
} else if (ndim == 3) {
tz = num_groups > 1
? mkldnn::memory::dims{num_groups, static_cast<int>(arr.shape()[O] / num_groups),
static_cast<int>(arr.shape()[I]),
static_cast<int>(arr.shape()[H])}
: mkldnn::memory::dims{static_cast<int>(arr.shape()[O]),
static_cast<int>(arr.shape()[I]),
static_cast<int>(arr.shape()[H])};
? mkldnn::memory::dims{num_groups, arr.shape()[O] / num_groups,
arr.shape()[I], arr.shape()[H]}
: mkldnn::memory::dims{arr.shape()[O],
arr.shape()[I], arr.shape()[H]};
format_tag = num_groups > 1 ? mkldnn::memory::format_tag::goiw
: mkldnn::memory::format_tag::oiw;
} else if (arr.shape().ndim() == 4) {
} else if (ndim == 4) {
tz = num_groups > 1
? mkldnn::memory::dims{num_groups, static_cast<int>(arr.shape()[O] / num_groups),
static_cast<int>(arr.shape()[I]),
static_cast<int>(arr.shape()[H]),
static_cast<int>(arr.shape()[W])}
? mkldnn::memory::dims{num_groups, arr.shape()[O] / num_groups,
arr.shape()[I], arr.shape()[H],
arr.shape()[W]}
: mkldnn::memory::dims{
static_cast<int>(arr.shape()[O]), static_cast<int>(arr.shape()[I]),
static_cast<int>(arr.shape()[H]), static_cast<int>(arr.shape()[W])};
arr.shape()[O], arr.shape()[I], arr.shape()[H], arr.shape()[W]};
format_tag = num_groups > 1 ? mkldnn::memory::format_tag::goihw
: mkldnn::memory::format_tag::oihw;
} else if (ndim == 5) {
tz = num_groups > 1
? mkldnn::memory::dims{num_groups, arr.shape()[O] / num_groups,
arr.shape()[I], arr.shape()[D],
arr.shape()[H], arr.shape()[W]}
: mkldnn::memory::dims{
arr.shape()[O], arr.shape()[I], arr.shape()[D],
arr.shape()[H], arr.shape()[W]};
format_tag = num_groups > 1 ? mkldnn::memory::format_tag::goidhw
: mkldnn::memory::format_tag::oidhw;
} else {
LOG(FATAL) << "The weight array has an unsupported number of dimensions";
}
Expand Down
60 changes: 52 additions & 8 deletions src/operator/nn/mkldnn/mkldnn_convolution.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,13 @@ DMLC_REGISTER_PARAMETER(MKLDNNConvParam);

bool SupportMKLDNNConv(const ConvolutionParam& params, const NDArray &input) {
if ((params.kernel.ndim() != 1) &&
(params.kernel.ndim() != 2))
(params.kernel.ndim() != 2) &&
(params.kernel.ndim() != 3))
return false;
return SupportMKLDNNQuantize(input.dtype()) &&
((input.shape().ndim() == 3) ||
(input.shape().ndim() == 4));
(input.shape().ndim() == 4) ||
(input.shape().ndim() == 5));
}

std::shared_ptr<mkldnn::convolution_forward::primitive_desc> GetConvFwdImpl(
Expand Down Expand Up @@ -77,9 +79,19 @@ std::shared_ptr<mkldnn::convolution_forward::primitive_desc> GetConvFwdImpl(
strides[1] = param.conv_param.stride[1];
padding[0] = param.conv_param.pad[0];
padding[1] = param.conv_param.pad[1];
} else if (param.conv_param.kernel.ndim() == 3) {
CHECK_GE(param.conv_param.stride.ndim(), 3);
CHECK_GE(param.conv_param.pad.ndim(), 3);
CHECK_GE(param.conv_param.dilate.ndim(), 3);
strides[0] = param.conv_param.stride[0];
strides[1] = param.conv_param.stride[1];
strides[2] = param.conv_param.stride[2];
padding[0] = param.conv_param.pad[0];
padding[1] = param.conv_param.pad[1];
padding[2] = param.conv_param.pad[2];
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv kernel size "
<< param.conv_param.kernel.ndim() << ", supporting only 1 or 2.";
<< param.conv_param.kernel.ndim() << ", supporting only 1 or 2 or 3.";
}
mkldnn::primitive_attr attr;
mkldnn::post_ops ops;
Expand Down Expand Up @@ -141,9 +153,13 @@ std::shared_ptr<mkldnn::convolution_forward::primitive_desc> GetConvFwdImpl(
} else if (param.conv_param.dilate.ndim() == 2) {
dilates[0] = param.conv_param.dilate[0] - 1;
dilates[1] = param.conv_param.dilate[1] - 1;
} else if (param.conv_param.dilate.ndim() == 3) {
dilates[0] = param.conv_param.dilate[0] - 1;
dilates[1] = param.conv_param.dilate[1] - 1;
dilates[2] = param.conv_param.dilate[2] - 1;
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv dilate size " << param.conv_param.dilate.ndim()
<< ", supporting only 1 or 2.";
<< ", supporting only 1 or 2 or 3.";
}
if (bias_md_ptr == nullptr) {
mkldnn::convolution_forward::desc desc(prop, mkldnn::algorithm::convolution_direct, data_md,
Expand Down Expand Up @@ -181,9 +197,19 @@ static std::shared_ptr<mkldnn::convolution_backward_data::primitive_desc> GetCon
strides[1] = param.stride[1];
padding[0] = param.pad[0];
padding[1] = param.pad[1];
} else if (param.kernel.ndim() == 3) {
CHECK_GE(param.stride.ndim(), 3);
CHECK_GE(param.pad.ndim(), 3);
CHECK_GE(param.dilate.ndim(), 3);
strides[0] = param.stride[0];
strides[1] = param.stride[1];
strides[2] = param.stride[2];
padding[0] = param.pad[0];
padding[1] = param.pad[1];
padding[2] = param.pad[2];
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv kernel size " << param.kernel.ndim()
<< ", supporting only 1 or 2.";
<< ", supporting only 1 or 2 or 3.";
}

auto GetConvBwdDataPd = [&data, &weight, &output,
Expand Down Expand Up @@ -216,9 +242,13 @@ static std::shared_ptr<mkldnn::convolution_backward_data::primitive_desc> GetCon
} else if (param.dilate.ndim() == 2) {
dilates[0] = param.dilate[0] - 1;
dilates[1] = param.dilate[1] - 1;
} else if (param.dilate.ndim() == 3) {
dilates[0] = param.dilate[0] - 1;
dilates[1] = param.dilate[1] - 1;
dilates[2] = param.dilate[2] - 1;
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv dilate size "
<< param.dilate.ndim() << ", supporting only 1 or 2.";
<< param.dilate.ndim() << ", supporting only 1 or 2 or 3.";
}
mkldnn::convolution_backward_data::desc desc(mkldnn::algorithm::convolution_direct, data_md,
weight_md, out_md, strides, dilates, padding,
Expand Down Expand Up @@ -250,9 +280,19 @@ static std::shared_ptr<mkldnn::convolution_backward_weights::primitive_desc> Get
strides[1] = param.stride[1];
padding[0] = param.pad[0];
padding[1] = param.pad[1];
} else if (param.kernel.ndim() == 3) {
CHECK_GE(param.stride.ndim(), 3);
CHECK_GE(param.pad.ndim(), 3);
CHECK_GE(param.dilate.ndim(), 3);
strides[0] = param.stride[0];
strides[1] = param.stride[1];
strides[2] = param.stride[2];
padding[0] = param.pad[0];
padding[1] = param.pad[1];
padding[2] = param.pad[2];
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv kernel size " << param.kernel.ndim()
<< ", supporting only 1 or 2.";
<< ", supporting only 1 or 2 or 3.";
}

auto GetConvBwdWeightsPd = [&data, &weight, &output,
Expand Down Expand Up @@ -291,9 +331,13 @@ static std::shared_ptr<mkldnn::convolution_backward_weights::primitive_desc> Get
} else if (param.dilate.ndim() == 2) {
dilates[0] = param.dilate[0] - 1;
dilates[1] = param.dilate[1] - 1;
} else if (param.dilate.ndim() == 3) {
dilates[0] = param.dilate[0] - 1;
dilates[1] = param.dilate[1] - 1;
dilates[2] = param.dilate[2] - 1;
} else {
LOG(FATAL) << "Unexpected MKL-DNN Conv dilate size "
<< param.dilate.ndim() << ", supporting only 1 or 2.";
<< param.dilate.ndim() << ", supporting only 1 or 2 or 3.";
}
if (bias == nullptr) {
mkldnn::convolution_backward_weights::desc desc(mkldnn::algorithm::convolution_direct,
Expand Down
Loading

0 comments on commit b748c45

Please sign in to comment.