tensor-tang
7 years ago
committed by
GitHub
parent
0151e4eb3d
commit
90780e22ce
@ -1,240 +0,0 @@
|
|||||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
|
||||||
//
|
|
||||||
// Licensed 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.
|
|
||||||
|
|
||||||
/*Licensed 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. */
|
|
||||||
|
|
||||||
#include "mkldnn.hpp"
|
|
||||||
#include "paddle/fluid/framework/tensor.h"
|
|
||||||
#include "paddle/fluid/operators/math/selected_rows_functor.h"
|
|
||||||
#include "paddle/fluid/operators/sum_op.h"
|
|
||||||
#include "paddle/fluid/platform/device_context.h"
|
|
||||||
#include "paddle/fluid/platform/mkldnn_helper.h"
|
|
||||||
|
|
||||||
namespace paddle {
|
|
||||||
namespace operators {
|
|
||||||
|
|
||||||
using paddle::framework::Tensor;
|
|
||||||
using paddle::platform::MKLDNNDeviceContext;
|
|
||||||
using paddle::platform::CPUDeviceContext;
|
|
||||||
using framework::DataLayout;
|
|
||||||
using mkldnn::memory;
|
|
||||||
using mkldnn::primitive;
|
|
||||||
using mkldnn::stream;
|
|
||||||
using mkldnn::sum;
|
|
||||||
using mkldnn::reorder;
|
|
||||||
using platform::to_void_cast;
|
|
||||||
|
|
||||||
template <typename T>
|
|
||||||
class SumMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
|
|
||||||
public:
|
|
||||||
void Compute(const paddle::framework::ExecutionContext& ctx) const override {
|
|
||||||
PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
|
|
||||||
"It must use CPUPlace.");
|
|
||||||
auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
|
|
||||||
const auto& mkldnn_engine = dev_ctx.GetEngine();
|
|
||||||
auto in_vars = ctx.MultiInputVar("X");
|
|
||||||
|
|
||||||
const int N = in_vars.size();
|
|
||||||
auto out_var = ctx.OutputVar("Out");
|
|
||||||
bool in_place = out_var == in_vars[0];
|
|
||||||
|
|
||||||
if (out_var->IsType<framework::LoDTensor>()) {
|
|
||||||
LoDTensor* output = ctx.Output<LoDTensor>("Out");
|
|
||||||
T* output_data = output->mutable_data<T>(ctx.GetPlace());
|
|
||||||
|
|
||||||
std::vector<int> dst_tz = framework::vectorize2int(output->dims());
|
|
||||||
auto src_tz = dst_tz;
|
|
||||||
memory::format output_format{memory::format::format_undef};
|
|
||||||
std::vector<float> scales;
|
|
||||||
std::vector<memory::primitive_desc> srcs_mpd;
|
|
||||||
std::vector<mkldnn::memory> srcs_mem;
|
|
||||||
|
|
||||||
PADDLE_ENFORCE(in_vars[0]->IsType<LoDTensor>(),
|
|
||||||
"Input[0] must be LoDTensors");
|
|
||||||
auto& input0 = in_vars[0]->Get<LoDTensor>();
|
|
||||||
PADDLE_ENFORCE(input0.layout() == DataLayout::kMKLDNN &&
|
|
||||||
input0.format() != memory::format::format_undef,
|
|
||||||
"Wrong layout/format for inputs[0]");
|
|
||||||
|
|
||||||
memory::format input_format = input0.format();
|
|
||||||
|
|
||||||
if (src_tz.size() == 1 && (input_format == memory::format::nchw ||
|
|
||||||
input_format == memory::format::nhwc)) {
|
|
||||||
input_format = memory::format::x;
|
|
||||||
}
|
|
||||||
if (src_tz.size() == 2 && (input_format == memory::format::nchw ||
|
|
||||||
input_format == memory::format::nhwc)) {
|
|
||||||
input_format = memory::format::nc;
|
|
||||||
}
|
|
||||||
|
|
||||||
for (int i = in_place ? 1 : 0; i < N; i++) {
|
|
||||||
PADDLE_ENFORCE(in_vars[i]->IsType<LoDTensor>(),
|
|
||||||
"all inputs must be all LoDTensors");
|
|
||||||
auto& input = in_vars[i]->Get<LoDTensor>();
|
|
||||||
PADDLE_ENFORCE(input.layout() == DataLayout::kMKLDNN &&
|
|
||||||
input.format() != memory::format::format_undef,
|
|
||||||
"Wrong layout/format for inputs");
|
|
||||||
|
|
||||||
if (input.numel() == 0) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
const T* input_data = input.data<T>();
|
|
||||||
|
|
||||||
auto src_md =
|
|
||||||
memory::desc(src_tz, memory::data_type::f32, input_format);
|
|
||||||
auto src_mpd = memory::primitive_desc(src_md, mkldnn_engine);
|
|
||||||
auto src_mem = memory(src_mpd, to_void_cast(input_data));
|
|
||||||
srcs_mpd.push_back(src_mpd);
|
|
||||||
srcs_mem.push_back(src_mem);
|
|
||||||
scales.push_back(1.0);
|
|
||||||
}
|
|
||||||
|
|
||||||
auto dst_md =
|
|
||||||
memory::desc(dst_tz, memory::data_type::f32, memory::format::any);
|
|
||||||
|
|
||||||
auto sum_pd = sum::primitive_desc(dst_md, scales, srcs_mpd);
|
|
||||||
|
|
||||||
std::shared_ptr<memory> dst_mem;
|
|
||||||
if (in_place) {
|
|
||||||
dst_mem.reset(new memory(sum_pd.dst_primitive_desc()));
|
|
||||||
} else {
|
|
||||||
dst_mem.reset(new memory(sum_pd.dst_primitive_desc(), output_data));
|
|
||||||
}
|
|
||||||
std::vector<mkldnn::primitive::at> inputs;
|
|
||||||
for (size_t i = 0; i < srcs_mem.size(); ++i) {
|
|
||||||
inputs.push_back(srcs_mem[i]);
|
|
||||||
}
|
|
||||||
|
|
||||||
auto sum_prim = mkldnn::sum(sum_pd, inputs, *dst_mem);
|
|
||||||
output_format = (memory::format)platform::GetMKLDNNFormat(sum_pd);
|
|
||||||
|
|
||||||
primitive reorder_prim;
|
|
||||||
std::shared_ptr<memory> target_mem;
|
|
||||||
if (in_place) {
|
|
||||||
output_format = input_format;
|
|
||||||
target_mem.reset(new memory(
|
|
||||||
{{{src_tz}, memory::data_type::f32, output_format}, mkldnn_engine},
|
|
||||||
output_data));
|
|
||||||
reorder_prim = reorder(*dst_mem, *target_mem);
|
|
||||||
}
|
|
||||||
|
|
||||||
std::vector<primitive> pipeline;
|
|
||||||
pipeline.push_back(sum_prim);
|
|
||||||
if (in_place) pipeline.push_back(reorder_prim);
|
|
||||||
stream(stream::kind::eager).submit(pipeline).wait();
|
|
||||||
|
|
||||||
output->set_layout(DataLayout::kMKLDNN);
|
|
||||||
output->set_format(output_format);
|
|
||||||
} else if (out_var->IsType<framework::SelectedRows>()) {
|
|
||||||
// TODO(@mozga-intel) Add MKLDNN SelectedRows support
|
|
||||||
std::unique_ptr<framework::SelectedRows> in0;
|
|
||||||
if (in_place) {
|
|
||||||
// If is in_place, we store the input[0] to in0
|
|
||||||
auto& in_sel0 = in_vars[0]->Get<SelectedRows>();
|
|
||||||
auto& rows = in_sel0.rows();
|
|
||||||
in0.reset(new framework::SelectedRows(rows, in_sel0.height()));
|
|
||||||
in0->mutable_value()->ShareDataWith(in_sel0.value());
|
|
||||||
}
|
|
||||||
|
|
||||||
auto get_selected_row = [&](size_t i) -> const SelectedRows& {
|
|
||||||
if (i == 0 && in0) {
|
|
||||||
return *in0.get();
|
|
||||||
} else {
|
|
||||||
return in_vars[i]->Get<SelectedRows>();
|
|
||||||
}
|
|
||||||
};
|
|
||||||
auto* out = ctx.Output<SelectedRows>("Out");
|
|
||||||
out->mutable_rows()->clear();
|
|
||||||
auto* out_value = out->mutable_value();
|
|
||||||
|
|
||||||
// Runtime InferShape
|
|
||||||
size_t first_dim = 0;
|
|
||||||
for (int i = 0; i < N; i++) {
|
|
||||||
auto& sel_row = get_selected_row(i);
|
|
||||||
first_dim += sel_row.rows().size();
|
|
||||||
}
|
|
||||||
auto in_dim =
|
|
||||||
framework::vectorize(get_selected_row(N - 1).value().dims());
|
|
||||||
in_dim[0] = static_cast<int64_t>(first_dim);
|
|
||||||
|
|
||||||
out_value->Resize(framework::make_ddim(in_dim));
|
|
||||||
|
|
||||||
// if all the input sparse vars are empty, no need to
|
|
||||||
// merge these vars.
|
|
||||||
if (first_dim == 0UL) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
out_value->mutable_data<T>(ctx.GetPlace());
|
|
||||||
math::SelectedRowsAddTo<CPUDeviceContext, T> functor;
|
|
||||||
int64_t offset = 0;
|
|
||||||
for (int i = 0; i < N; i++) {
|
|
||||||
auto& sel_row = get_selected_row(i);
|
|
||||||
if (sel_row.rows().size() == 0) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
PADDLE_ENFORCE_EQ(out->height(), sel_row.height());
|
|
||||||
functor(ctx.template device_context<CPUDeviceContext>(), sel_row,
|
|
||||||
offset, out);
|
|
||||||
offset += sel_row.value().numel();
|
|
||||||
}
|
|
||||||
} else if (out_var->IsType<framework::LoDTensorArray>()) {
|
|
||||||
// TODO(@mozga-intel) Add MKLDNN LoDTensorArray support
|
|
||||||
auto& out_array = *out_var->GetMutable<framework::LoDTensorArray>();
|
|
||||||
for (size_t i = in_place ? 1 : 0; i < in_vars.size(); ++i) {
|
|
||||||
PADDLE_ENFORCE(in_vars[i]->IsType<framework::LoDTensorArray>(),
|
|
||||||
"Only support all inputs are TensorArray");
|
|
||||||
auto& in_array = in_vars[i]->Get<framework::LoDTensorArray>();
|
|
||||||
|
|
||||||
for (size_t i = 0; i < in_array.size(); ++i) {
|
|
||||||
if (in_array[i].numel() != 0) {
|
|
||||||
if (i >= out_array.size()) {
|
|
||||||
out_array.resize(i + 1);
|
|
||||||
}
|
|
||||||
if (out_array[i].numel() == 0) {
|
|
||||||
framework::TensorCopy(in_array[i], in_array[i].place(),
|
|
||||||
ctx.device_context(), &out_array[i]);
|
|
||||||
out_array[i].set_lod(in_array[i].lod());
|
|
||||||
} else {
|
|
||||||
PADDLE_ENFORCE(out_array[i].lod() == in_array[i].lod());
|
|
||||||
auto in = EigenVector<T>::Flatten(in_array[i]);
|
|
||||||
auto result = EigenVector<T>::Flatten(out_array[i]);
|
|
||||||
result.device(*ctx.template device_context<MKLDNNDeviceContext>()
|
|
||||||
.eigen_device()) = result + in;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
PADDLE_THROW("Unexpected branch, output variable type is %s",
|
|
||||||
out_var->Type().name());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
} // namespace operators
|
|
||||||
} // namespace paddle
|
|
||||||
|
|
||||||
REGISTER_OP_KERNEL(sum, MKLDNN, ::paddle::platform::CPUPlace,
|
|
||||||
paddle::operators::SumMKLDNNOpKernel<float>);
|
|
||||||
File diff suppressed because it is too large
Load Diff
@ -1,26 +0,0 @@
|
|||||||
# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
|
||||||
#
|
|
||||||
# Licensed 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.
|
|
||||||
|
|
||||||
import unittest
|
|
||||||
|
|
||||||
from test_sum_op import TestSumOp
|
|
||||||
|
|
||||||
|
|
||||||
class TestMKLDNN(TestSumOp):
|
|
||||||
def init_kernel_type(self):
|
|
||||||
self.use_mkldnn = True
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
unittest.main()
|
|
||||||
Loading…
Reference in new issue