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

[MKLDNN] Enable convolution fusion. #12308

Closed
ZhennanQin wants to merge 1 commit into from
Closed

[MKLDNN] Enable convolution fusion. #12308

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/operator/nn/convolution-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ enum ConvolutionOpInputs {kData, kWeight, kBias};
enum ConvolutionOpOutputs {kOut};
enum ConvolutionOpResource {kTempSpace};
enum ConvolutionOpCudnnTune {kOff, kLimited, kFastest};
}
} // namespace conv

struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
TShape kernel;
Expand Down Expand Up @@ -129,6 +129,10 @@ void ConvolutionParamParser(nnvm::NodeAttrs* attrs);

typedef ParamOpSign<ConvolutionParam> ConvSignature;

static inline size_t GetInShapeSize(const ConvolutionParam &param_) {
return 2 + (param_.no_bias ? 0 : 1);
}

} // namespace op
} // namespace mxnet

Expand Down Expand Up @@ -176,8 +180,7 @@ class ConvolutionOp {
using namespace mshadow;
using namespace mshadow::expr;
CHECK_EQ(req[conv::kOut], kWriteTo);
size_t expected = param_.no_bias ? 2 : 3;
CHECK_EQ(in_data.size(), expected);
CHECK_EQ(in_data.size(), GetInShapeSize(param_));
CHECK_EQ(out_data.size(), 1U);
CHECK_EQ(req[conv::kOut], kWriteTo);
LayerSetUp(in_data[conv::kData].shape_, out_data[conv::kOut].shape_);
Expand Down
28 changes: 15 additions & 13 deletions src/operator/nn/convolution.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "../elemwise_op_common.h"
#include "./mkldnn/mkldnn_ops-inl.h"
#include "./mkldnn/mkldnn_base-inl.h"
#include "./mkldnn/mkldnn_convolution-inl.h"
#if MXNET_USE_NNPACK == 1
#include "../nnpack/nnpack_pooling-inl.h"
#endif // MXNET_USE_NNPACK
Expand All @@ -41,11 +42,19 @@ static inline index_t AddPad(index_t dsize, index_t pad) {
}

static inline std::vector<std::string> ListArguments(const ConvolutionParam& param_) {
if (!param_.no_bias) {
if (!param_.no_bias)
return {"data", "weight", "bias"};
} else {
else
return {"data", "weight"};
}

static inline std::string PrintArguments(const ConvolutionParam& param_) {
auto args = ListArguments(param_);
std::string str = "[";
Copy link
Copy Markdown
Contributor

@lebeg lebeg Aug 24, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's better to use std::stringstream to compose such string, otherwise it's a lot of redundant copying internally.

for (const auto &arg : args) {
str += arg + ", ";
}
return str.substr(0, str.size() - 2) + "]";
}

#if MXNET_USE_MKLDNN == 1
Expand Down Expand Up @@ -85,11 +94,7 @@ static bool ConvolutionShape(const nnvm::NodeAttrs& attrs,
std::vector<TShape> *out_shape) {
using namespace mshadow;
const ConvolutionParam& param_ = nnvm::get<ConvolutionParam>(attrs.parsed);
if (!param_.no_bias) {
CHECK_EQ(in_shape->size(), 3U) << "Input:[data, weight, bias]";
} else {
CHECK_EQ(in_shape->size(), 2U) << "Input:[data, weight]";
}
CHECK_EQ(in_shape->size(), GetInShapeSize(param_)) << "Input:" << PrintArguments(param_);
// CHECK_EQ(out_shape->size(), 1) << "Output: [output]";
out_shape->resize(1, TShape());
const TShape &dshp = (*in_shape)[conv::kData];
Expand Down Expand Up @@ -294,7 +299,7 @@ inline static bool ConvStorageType(const nnvm::NodeAttrs& attrs,
std::vector<int> *in_attrs,
std::vector<int> *out_attrs) {
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
uint32_t in_expected = param.no_bias ? 2 : 3;
uint32_t in_expected = GetInShapeSize(param);
CHECK_EQ(in_attrs->size(), in_expected);
CHECK_EQ(out_attrs->size(), 1);

Expand Down Expand Up @@ -470,17 +475,14 @@ There are other options to tune the performance.
)code" ADD_FILELINE)
.set_num_inputs([](const NodeAttrs& attrs) {
const ConvolutionParam& params = nnvm::get<ConvolutionParam>(attrs.parsed);
return params.no_bias ? 2 : 3;
return GetInShapeSize(params);
})
.set_num_outputs(1)
.set_attr_parser(ConvolutionParamParser)
.set_attr<nnvm::FListInputNames>("FListInputNames",
[](const NodeAttrs& attrs) {
const ConvolutionParam& params = nnvm::get<ConvolutionParam>(attrs.parsed);
if (params.no_bias)
return std::vector<std::string>{"data", "weight"};
else
return std::vector<std::string>{"data", "weight", "bias"};
return ListArguments(params);
})
.set_attr<nnvm::FListOutputNames>("FListOutputNames",
[](const NodeAttrs& attrs) {
Expand Down
41 changes: 35 additions & 6 deletions src/operator/nn/mkldnn/mkldnn_convolution-inl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,19 +35,48 @@
namespace mxnet {
namespace op {

struct ConvFusionParam : public dmlc::Parameter<ConvFusionParam> {
// When adding more members into this clss, please double check GetHash()
// won't overflow.
bool with_bn;
bool with_relu;
bool with_sum;
bool with_postsum_relu;
DMLC_DECLARE_PARAMETER(ConvFusionParam) {
DMLC_DECLARE_FIELD(with_bn).set_default(false)
.describe("Add post batchnorm.");
DMLC_DECLARE_FIELD(with_relu).set_default(false)
.describe("Add post relu");
DMLC_DECLARE_FIELD(with_sum).set_default(false)
.describe("Add post sum");
DMLC_DECLARE_FIELD(with_postsum_relu).set_default(false)
.describe("Add post relu after sum");
}
const int GetHash() const {
Copy link
Copy Markdown
Contributor

@lebeg lebeg Aug 24, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can std::string be used here instead?

int hash = 0;
hash = hash * 2 + this->with_bn ? 1 : 0;
Copy link
Copy Markdown
Contributor

@marcoabreu marcoabreu Aug 23, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possible hash collision: with_bn=0 and with_relu=1 equals BN=1 and relu0. Consider using bitflags

Copy link
Copy Markdown
Contributor

@larroy larroy Aug 24, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not std::hash?

hash = hash * 2 + this->with_relu ? 1 : 0;
hash = hash * 2 + this->with_sum ? 1 : 0;
hash = hash * 2 + this->with_postsum_relu ? 1 : 0;
return hash;
}
};

mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
const ConvolutionParam& param, const bool is_train, const NDArray &data,
const NDArray &weights, const NDArray *bias, const NDArray &output);
const ConvolutionParam &param, const ConvFusionParam &fusion_param,
const bool is_train, const NDArray &data, const NDArray &weights,
const NDArray *bias, const NDArray &output);

class MKLDNNConvForward {
public:
mkldnn::convolution_forward::primitive_desc fwd_pd;

MKLDNNConvForward(const ConvolutionParam& param, const bool is_train,
MKLDNNConvForward(const ConvolutionParam &param,
const ConvFusionParam &fusion_param, const bool is_train,
const NDArray &data, const NDArray &weights,
const NDArray *bias, const NDArray &output): fwd_pd(
GetConvFwdImpl(param, is_train, data, weights, bias, output)) {
}
const NDArray *bias, const NDArray &output)
: fwd_pd(GetConvFwdImpl(param, fusion_param, is_train, data, weights,
bias, output)) {}

void SetNewMem(const mkldnn::memory &data, const mkldnn::memory &weight,
const mkldnn::memory *bias, const mkldnn::memory &output);
Expand Down
92 changes: 69 additions & 23 deletions src/operator/nn/mkldnn/mkldnn_convolution.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,18 @@
namespace mxnet {
namespace op {

DMLC_REGISTER_PARAMETER(ConvFusionParam);

bool SupportMKLDNNConv(const ConvolutionParam& params, const NDArray &input) {
if (params.kernel.ndim() != 2)
return false;
return input.dtype() == mshadow::kFloat32 && input.shape().ndim() == 4;
}

mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
const ConvolutionParam& param, const bool is_train, const NDArray &data,
const NDArray &weights, const NDArray *bias, const NDArray &output) {
const ConvolutionParam &param, const ConvFusionParam &fusion_param,
const bool is_train, const NDArray &data, const NDArray &weights,
const NDArray *bias, const NDArray &output) {
auto prop = is_train ? mkldnn::prop_kind::forward_training : mkldnn::prop_kind::forward_scoring;
auto data_md = GetMemDesc(data);
auto weight_md = GetWeightDesc(weights, param.num_group);
Expand All @@ -57,16 +60,36 @@ mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
mkldnn::memory::dims padding{0, 0};
padding[0] = param.pad[0];
padding[1] = param.pad[1];
mkldnn::primitive_attr attr;
mkldnn::post_ops ops;
if (fusion_param.with_relu) {
float scale = 1.0f; // for fp32, scale is 1.
float alpha = 0.0f; // negative slope for mkldnn_eltwise_relu.
float beta = 1.0f; // ignored for mkldnn_eltwise_relu.
ops.append_eltwise(scale, eltwise_relu, alpha, beta);

}
if (fusion_param.with_sum) {
float scale = 1.0f;
ops.append_sum(scale);
}
if (fusion_param.with_postsum_relu) {
float scale = 1.0f; // for fp32, scale is 1.
float alpha = 0.0f; // negative slope for mkldnn_eltwise_relu.
float beta = 1.0f; // ignored for mkldnn_eltwise_relu.
ops.append_eltwise(scale, eltwise_relu, alpha, beta);
}
attr.set_post_ops(ops);
if (param.dilate.ndim() == 0 && bias == nullptr) {
mkldnn::convolution_forward::desc desc(prop, mkldnn::algorithm::convolution_direct,
data_md, weight_md, out_md, strides, padding, padding, mkldnn::padding_kind::zero);
return mkldnn::convolution_forward::primitive_desc(desc, engine);
return mkldnn::convolution_forward::primitive_desc(desc, attr, engine);
} else if (param.dilate.ndim() == 0) {
auto bias_md = GetMemDesc(*bias);
mkldnn::convolution_forward::desc desc(prop, mkldnn::algorithm::convolution_direct,
data_md, weight_md, bias_md, out_md, strides, padding, padding,
mkldnn::padding_kind::zero);
return mkldnn::convolution_forward::primitive_desc(desc, engine);
return mkldnn::convolution_forward::primitive_desc(desc, attr, engine);
} else {
mkldnn::memory::dims dilates{0, 0};
dilates[0] = param.dilate[0] - 1;
Expand All @@ -75,14 +98,14 @@ mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
mkldnn::convolution_forward::desc desc(prop, mkldnn::algorithm::convolution_direct,
data_md, weight_md, out_md, strides, dilates, padding, padding,
mkldnn::padding_kind::zero);
return mkldnn::convolution_forward::primitive_desc(desc, engine);
return mkldnn::convolution_forward::primitive_desc(desc, attr, engine);
} else {
auto bias_md = GetMemDesc(*bias);
mkldnn::convolution_forward::desc desc(prop, mkldnn::algorithm::convolution_direct,
data_md, weight_md, bias_md, out_md, strides,
dilates, padding, padding,
mkldnn::padding_kind::zero);
return mkldnn::convolution_forward::primitive_desc(desc, engine);
return mkldnn::convolution_forward::primitive_desc(desc, attr, engine);
}
}
}
Expand Down Expand Up @@ -207,16 +230,20 @@ void MKLDNNConvForward::SetNewMem(const mkldnn::memory &data,
}
}

MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs& attrs, const bool is_train,
const NDArray &data, const NDArray &weights,
const NDArray *bias, const NDArray &output) {
MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs &attrs,
const bool is_train, const NDArray &data,
const NDArray &weights, const NDArray *bias,
const NDArray &output) {
#if DMLC_CXX11_THREAD_LOCAL
static thread_local std::unordered_map<MKLDNNConvSignature, MKLDNNConvForward, OpHash> fwds;
#else
static MX_THREAD_LOCAL std::unordered_map<MKLDNNConvSignature, MKLDNNConvForward, OpHash> fwds;
#endif
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
ConvFusionParam fusion_param;
fusion_param.Init(attrs.dict, dmlc::parameter::kAllowUnknown);
MKLDNNConvSignature key(param);
key.AddSign(fusion_param.GetHash());
key.AddSign(is_train);
// Here we can sign the conv op with NDArray because conv primitive will
// decide the right layout for the, so we only need to get the shape and the
Expand All @@ -227,9 +254,10 @@ MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs& attrs, const bool is_train,
if (bias)
key.AddSign(*bias);


auto it = fwds.find(key);
if (it == fwds.end()) {
MKLDNNConvForward fwd(param, is_train, data, weights, bias, output);
MKLDNNConvForward fwd(param, fusion_param, is_train, data, weights, bias, output);
auto ins_ret = fwds.insert(
std::pair<MKLDNNConvSignature, MKLDNNConvForward>(key, fwd));
CHECK(ins_ret.second);
Expand All @@ -238,15 +266,20 @@ MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs& attrs, const bool is_train,
return it->second;
}

void MKLDNNConvolutionForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
const std::vector<NDArray> &in_data,
const std::vector<OpReqType> &req,
const std::vector<NDArray> &out_data) {
void MKLDNNConvolutionForward(const nnvm::NodeAttrs &attrs,
const OpContext &ctx,
const std::vector<NDArray> &in_data,
const std::vector<OpReqType> &req,
const std::vector<NDArray> &out_data) {
TmpMemMgr::Get()->Init(ctx.requested[conv::kTempSpace]);
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
const ConvolutionParam &param = nnvm::get<ConvolutionParam>(attrs.parsed);
ConvFusionParam fusion_param;
fusion_param.Init(attrs.dict, dmlc::parameter::kAllowUnknown);
NDArray weight = in_data[conv::kWeight];
MKLDNNConvForward &fwd = GetConvFwd(attrs, ctx.is_train, in_data[conv::kData], weight,
param.no_bias ? nullptr : &in_data[conv::kBias], out_data[conv::kOut]);
bool no_bias = param.no_bias && !fusion_param.with_bn;
MKLDNNConvForward &fwd = GetConvFwd(
attrs, ctx.is_train, in_data[conv::kData], weight,
no_bias ? nullptr : &in_data[conv::kBias], out_data[conv::kOut]);

auto data_mem = in_data[conv::kData].GetMKLDNNDataReorder(fwd.fwd_pd.src_primitive_desc());
const mkldnn::memory *weight_mem;
Expand All @@ -271,11 +304,21 @@ void MKLDNNConvolutionForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx
CHECK(weight_mem->get_primitive_desc() == fwd.fwd_pd.weights_primitive_desc());
}
}
auto out_mem = CreateMKLDNNMem(out_data[conv::kOut], fwd.fwd_pd.dst_primitive_desc(),
req[conv::kOut]);
mkldnn_output_t out_mem;
if (fusion_param.with_sum) {
out_mem = mkldnn_output_t(
OutDataOp::Noop,
const_cast<mkldnn::memory *>(out_data[conv::kOut].GetMKLDNNDataReorder(
fwd.fwd_pd.dst_primitive_desc())));
} else {
out_mem = CreateMKLDNNMem(out_data[conv::kOut],
fwd.fwd_pd.dst_primitive_desc(), req[conv::kOut]);
}

const mkldnn::memory *bias_mem = nullptr;
if (!param.no_bias)
bias_mem = in_data[conv::kBias].GetMKLDNNDataReorder(fwd.fwd_pd.bias_primitive_desc());
if (!no_bias) {
bias_mem = in_data[conv::kBias].GetMKLDNNData();
}
fwd.SetNewMem(*data_mem, *weight_mem, bias_mem, *out_mem.second);
MKLDNNStream::Get()->RegisterPrim(fwd.GetFwd());

Expand All @@ -290,8 +333,11 @@ void MKLDNNConvolutionBackward(const nnvm::NodeAttrs& attrs, const OpContext &ct
TmpMemMgr::Get()->Init(ctx.requested[conv::kTempSpace]);
const std::vector<NDArray> &in_grad = outputs;
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
mkldnn::convolution_forward::primitive_desc fwd_pd = GetConvFwdImpl(param, ctx.is_train,
inputs[conv::kData + 1], inputs[conv::kWeight + 1],
ConvFusionParam fusion_param;
fusion_param.Init(attrs.dict, dmlc::parameter::kAllowUnknown);
mkldnn::convolution_forward::primitive_desc fwd_pd = GetConvFwdImpl(
param, fusion_param, ctx.is_train, inputs[conv::kData + 1],
inputs[conv::kWeight + 1],
param.no_bias ? nullptr : &inputs[conv::kBias + 1], inputs[conv::kOut]);

CHECK_NE(req[conv::kWeight], kWriteInplace) << "cannot write weight inplace";
Expand Down
Loading