提交 f2a88042 编写于 作者: M Michal Gallus

Fix style @ concat integration and tests

test=develop
上级 738069e4
...@@ -30,15 +30,15 @@ using platform::to_void_cast; ...@@ -30,15 +30,15 @@ using platform::to_void_cast;
static void EnforceLayouts(const std::vector<const Tensor*> inputs) { static void EnforceLayouts(const std::vector<const Tensor*> inputs) {
for (auto* input : inputs) { for (auto* input : inputs) {
const bool is_layout_correct = input->layout() == DataLayout::kMKLDNN; const bool is_layout_correct = input->layout() == DataLayout::kMKLDNN;
const bool is_format_defined = input->format() != const bool is_format_defined =
memory::format::format_undef; input->format() != memory::format::format_undef;
PADDLE_ENFORCE(is_layout_correct && is_format_defined, PADDLE_ENFORCE(is_layout_correct && is_format_defined,
"Wrong layout/format set for Input tensor"); "Wrong layout/format set for Input tensor");
} }
} }
static memory::primitive_desc CreateMemPrimDesc( static memory::primitive_desc CreateMemPrimDesc(const Tensor& input,
const Tensor& input, const mkldnn::engine& engine) { const mkldnn::engine& engine) {
constexpr auto data_type = mkldnn::memory::f32; constexpr auto data_type = mkldnn::memory::f32;
const auto dims = paddle::framework::vectorize2int(input.dims()); const auto dims = paddle::framework::vectorize2int(input.dims());
const auto format = input.format(); const auto format = input.format();
...@@ -62,8 +62,7 @@ static platform::CPUPlace GetCpuPlace( ...@@ -62,8 +62,7 @@ static platform::CPUPlace GetCpuPlace(
static const mkldnn::engine& GetMKLDNNEngine( static const mkldnn::engine& GetMKLDNNEngine(
const paddle::framework::ExecutionContext& ctx) { const paddle::framework::ExecutionContext& ctx) {
auto& dev_ctx = auto& dev_ctx = ctx.template device_context<platform::MKLDNNDeviceContext>();
ctx.template device_context<platform::MKLDNNDeviceContext>();
return dev_ctx.GetEngine(); return dev_ctx.GetEngine();
} }
...@@ -103,8 +102,8 @@ class ConcatPrimitiveFactory { ...@@ -103,8 +102,8 @@ class ConcatPrimitiveFactory {
for (size_t i = 0; i < multi_input.size(); i++) { for (size_t i = 0; i < multi_input.size(); i++) {
auto mem_prim_desc = CreateMemPrimDesc(*multi_input[i], mkldnn_engine); auto mem_prim_desc = CreateMemPrimDesc(*multi_input[i], mkldnn_engine);
srcs_pd.push_back(mem_prim_desc); srcs_pd.push_back(mem_prim_desc);
srcs.push_back(memory(mem_prim_desc, srcs.push_back(
to_void_cast(multi_input[i]->data<T>()))); memory(mem_prim_desc, to_void_cast(multi_input[i]->data<T>())));
} }
} }
...@@ -134,8 +133,8 @@ class ConcatMKLDNNOpKernel : public paddle::framework::OpKernel<T> { ...@@ -134,8 +133,8 @@ class ConcatMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
int64_t concat_axis = static_cast<int64_t>(ctx.Attr<int>("axis")); int64_t concat_axis = static_cast<int64_t>(ctx.Attr<int>("axis"));
ConcatPrimitiveFactory<T> prim_creator; ConcatPrimitiveFactory<T> prim_creator;
auto concat_pd = prim_creator.CreateConcatPrimDescriptor(multi_input, auto concat_pd = prim_creator.CreateConcatPrimDescriptor(
output, static_cast<int>(concat_axis), mkldnn_engine); multi_input, output, static_cast<int>(concat_axis), mkldnn_engine);
auto concat = prim_creator.CreateConcatPrimitive(concat_pd, output, place); auto concat = prim_creator.CreateConcatPrimitive(concat_pd, output, place);
stream(stream::kind::eager).submit({concat}).wait(); stream(stream::kind::eager).submit({concat}).wait();
......
...@@ -14,9 +14,9 @@ limitations under the License. */ ...@@ -14,9 +14,9 @@ limitations under the License. */
#include "paddle/fluid/operators/concat_op.h" #include "paddle/fluid/operators/concat_op.h"
#include <paddle/fluid/platform/mkldnn_helper.h>
#include <string> #include <string>
#include <vector> #include <vector>
#include <paddle/fluid/platform/mkldnn_helper.h>
namespace paddle { namespace paddle {
namespace operators { namespace operators {
...@@ -63,16 +63,17 @@ class ConcatOp : public framework::OperatorWithKernel { ...@@ -63,16 +63,17 @@ class ConcatOp : public framework::OperatorWithKernel {
protected: protected:
framework::OpKernelType GetExpectedKernelType( framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override { const framework::ExecutionContext &ctx) const override {
auto input_data_type = framework::GetDataTypeOfVar(ctx.MultiInputVar("X")[0]); auto input_data_type =
framework::GetDataTypeOfVar(ctx.MultiInputVar("X")[0]);
#ifdef PADDLE_WITH_MKLDNN #ifdef PADDLE_WITH_MKLDNN
if (platform::CanMKLDNNBeUsed(ctx)) { if (platform::CanMKLDNNBeUsed(ctx)) {
return framework::OpKernelType(input_data_type, ctx.GetPlace(), return framework::OpKernelType(input_data_type, ctx.GetPlace(),
framework::DataLayout::kMKLDNN, framework::DataLayout::kMKLDNN,
framework::LibraryType::kMKLDNN); framework::LibraryType::kMKLDNN);
} }
#endif #endif
return framework::OpKernelType(input_data_type, ctx.GetPlace()); return framework::OpKernelType(input_data_type, ctx.GetPlace());
} }
}; };
...@@ -82,7 +83,8 @@ class ConcatOpMaker : public framework::OpProtoAndCheckerMaker { ...@@ -82,7 +83,8 @@ class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
void Make() override { void Make() override {
AddInput("X", "Input tensors of concat operator.").AsDuplicable(); AddInput("X", "Input tensors of concat operator.").AsDuplicable();
AddOutput("Out", "Output tensor of concat operator."); AddOutput("Out", "Output tensor of concat operator.");
AddAttr<bool>("use_mkldnn", AddAttr<bool>(
"use_mkldnn",
"(bool, default false) Indicates if MKL-DNN kernel will be used") "(bool, default false) Indicates if MKL-DNN kernel will be used")
.SetDefault(false); .SetDefault(false);
AddAttr<int>("axis", AddAttr<int>("axis",
...@@ -101,7 +103,6 @@ Examples: ...@@ -101,7 +103,6 @@ Examples:
[5,6]] [5,6]]
)DOC"); )DOC");
} }
}; };
......
...@@ -29,6 +29,7 @@ class TestMKLDNNConcatOp(TestConcatOp): ...@@ -29,6 +29,7 @@ class TestMKLDNNConcatOp(TestConcatOp):
def init_kernel_type(self): def init_kernel_type(self):
self.use_mkldnn = True self.use_mkldnn = True
class TestMKLDNNConcatOp2(TestConcatOp2): class TestMKLDNNConcatOp2(TestConcatOp2):
def setUp(self): def setUp(self):
super(TestMKLDNNConcatOp2, self).setUp() super(TestMKLDNNConcatOp2, self).setUp()
...@@ -40,6 +41,7 @@ class TestMKLDNNConcatOp2(TestConcatOp2): ...@@ -40,6 +41,7 @@ class TestMKLDNNConcatOp2(TestConcatOp2):
def init_kernel_type(self): def init_kernel_type(self):
self.use_mkldnn = True self.use_mkldnn = True
class TestMKLDNNConcatOp3(TestConcatOp3): class TestMKLDNNConcatOp3(TestConcatOp3):
def setUp(self): def setUp(self):
super(TestMKLDNNConcatOp3, self).setUp() super(TestMKLDNNConcatOp3, self).setUp()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册