From 32b62c25afbc5a2a87a7c0bd80c463ba36081f10 Mon Sep 17 00:00:00 2001 From: zhaoyuchen2018 <45989343+zhaoyuchen2018@users.noreply.github.com> Date: Tue, 7 May 2019 11:07:07 +0800 Subject: [PATCH] optimize sum op (#16820) * optimize sum op fuse multi eigen kernel calls into one cuda kernel. refine code test=develop. Signed-off-by: zhaoyuchen * Refine code. test=develop Signed-off-by: zhaoyuchen * Refine code according to comments. test=develop * refine code delete sum_op_gpu.h test=develop * Fix test error. test=develop Signed-off-by: zhaoyuchen * refine code in format. test=develop. * refine code test=develop Signed-off-by: zhaoyuchen * refine code test=develop Signed-off-by: zhaoyuchen --- paddle/fluid/operators/sum_op.cu | 237 +++++++++++++++++++++++++++++++ paddle/fluid/operators/sum_op.h | 169 ++++++++++++---------- 2 files changed, 329 insertions(+), 77 deletions(-) diff --git a/paddle/fluid/operators/sum_op.cu b/paddle/fluid/operators/sum_op.cu index 6125ed07b6..43427a4af2 100644 --- a/paddle/fluid/operators/sum_op.cu +++ b/paddle/fluid/operators/sum_op.cu @@ -8,9 +8,246 @@ 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 +#include "paddle/fluid/framework/op_registry.h" #include "paddle/fluid/operators/sum_op.h" #include "paddle/fluid/platform/float16.h" +namespace plat = paddle::platform; + +namespace paddle { +namespace operators { + +#define CEIL_DIV(x, y) (((x) + (y)-1) / (y)) + +using LoDTensor = framework::LoDTensor; + +template +__global__ void Sum2CUDAKernel(const T *in_0, const T *in_1, T *out, + int64_t N) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + while (id < N) { + out[id] = in_0[id] + in_1[id]; + id += blockDim.x * gridDim.x; + } +} + +template +__global__ void SumArrayCUDAKernel(T **in, T *out, int64_t N, size_t in_size, + bool read_dst) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + while (id < N) { + T total(0); + for (int i = 0; i < in_size; ++i) { + const T *tmp = in[i]; + if (tmp) { + total += tmp[id]; + } + } + if (read_dst) { + out[id] += total; + } else { + out[id] = total; + } + id += blockDim.x * gridDim.x; + } +} + +template +__global__ void SumSelectedRowsCUDAKernel(T **sr_in_out, int64_t N, + size_t rows) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + while (id < N) { + for (int i = 0; i < 2 * rows; i += 2) { + const T *tmp = sr_in_out[i]; + T *tmp_out = sr_in_out[i + 1]; + if (tmp && tmp_out) { + tmp_out[id] += tmp[id]; + } + } + id += blockDim.x * gridDim.x; + } +} + +template +__global__ void SumAlign4CUDAKernel(const T *in_0, const T *in_1, T *out, + int64_t N) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = id; i < N / 4; i += blockDim.x * gridDim.x) { + const float4 *in0_4 = reinterpret_cast(in_0); + const float4 *in1_4 = reinterpret_cast(in_1); + float4 tmp; + tmp.x = in0_4[i].x + in1_4[i].x; + tmp.y = in0_4[i].y + in1_4[i].y; + tmp.z = in0_4[i].z + in1_4[i].z; + tmp.w = in0_4[i].w + in1_4[i].w; + reinterpret_cast(out)[i] = tmp; + } +} + +template +void FuseLodTensorSumCompute(const framework::ExecutionContext &context) { + auto in_vars = context.MultiInputVar("X"); + const size_t in_num = in_vars.size(); + + constexpr size_t theory_sm_threads = 1024; + auto &dev_ctx = + context.template device_context(); + auto stream = dev_ctx.stream(); + + auto max_threads = dev_ctx.GetMaxPhysicalThreadCount(); + auto sm_count = max_threads / theory_sm_threads; + size_t tile_size = 0; + dim3 grids; + dim3 blocks; + + auto ComputeKernelParameter = [&](size_t length) { + if (length >= max_threads) + tile_size = 1024; + else if (length < max_threads && length > sm_count * 128) + tile_size = 512; + else if (length <= sm_count * 128) + tile_size = 256; + grids = dim3(CEIL_DIV(length, tile_size), 1, 1); + blocks = dim3(tile_size, 1, 1); + }; + + auto *out = context.Output("Out"); + + auto out_var = context.OutputVar("Out"); + bool in_place = in_vars[0] == out_var; + + if (!in_place) { + out->mutable_data(context.GetPlace()); + } + int start = in_place ? 1 : 0; + if (!in_place) { + // seperate path for a+b,maybe not fast than eigen + if (in_num == 2 && in_vars[0]->IsType() && + in_vars[1]->IsType()) { + auto &in_0 = in_vars[0]->Get(); + auto &in_1 = in_vars[1]->Get(); + + auto length = in_0.numel(); + if (length) { + ComputeKernelParameter(length); + Sum2CUDAKernel<<>>( + in_0.data(), in_1.data(), out->data(), length); + } + return; + } + } + if (!in_place) { + math::SetConstant constant_functor; + constant_functor( + context.template device_context(), out, + static_cast(0)); + } + + std::vector in_data; + std::vector selectrow_index; + int64_t lod_length = 0; + bool dst_write = false; + for (int i = start; i < in_num; ++i) { + if (in_vars[i]->IsType()) { + auto &in_i = in_vars[i]->Get(); + in_data.emplace_back(in_i.data()); + lod_length = in_i.numel(); + } else if (in_vars[i]->IsType()) { + selectrow_index.push_back(i); + } + } + + // compute select rows seperately. + if (!selectrow_index.empty()) { + std::vector sr_in_out_data; + size_t rows = 0; + int64_t length = 0; + for (auto index : selectrow_index) { + auto &sr = in_vars[index]->Get(); + auto &sr_value = sr.value(); + auto &sr_rows = sr.rows(); + + auto row_numel = sr_value.numel() / sr_rows.size(); + auto out_dims = out->dims(); + + PADDLE_ENFORCE_EQ(sr.height(), out_dims[0]); + PADDLE_ENFORCE_EQ(row_numel, out->numel() / sr.height()); + + auto *sr_data = sr_value.data(); + auto *sr_out_data = out->data(); + rows += sr_rows.size(); + length = row_numel; + + for (size_t i = 0; i < sr_rows.size(); ++i) { + sr_in_out_data.emplace_back(&sr_data[i * row_numel]); + sr_in_out_data.emplace_back(&sr_out_data[sr_rows[i] * row_numel]); + } + } + if (!sr_in_out_data.empty()) { + auto tmp_sr_in_out_array = + platform::DeviceTemporaryAllocator::Instance().Get(dev_ctx).Allocate( + sr_in_out_data.size() * sizeof(T *)); + + memory::Copy(boost::get(dev_ctx.GetPlace()), + tmp_sr_in_out_array->ptr(), platform::CPUPlace(), + reinterpret_cast(sr_in_out_data.data()), + sr_in_out_data.size() * sizeof(T *), dev_ctx.stream()); + + T **sr_in_out_array_data = + reinterpret_cast(tmp_sr_in_out_array->ptr()); + + ComputeKernelParameter(length); + SumSelectedRowsCUDAKernel<<>>( + sr_in_out_array_data, length, rows); + dst_write = true; + } + } + // if indata not null, merge into one kernel call. + if (!in_data.empty()) { + auto tmp_in_array = + platform::DeviceTemporaryAllocator::Instance().Get(dev_ctx).Allocate( + in_data.size() * sizeof(T *)); + + memory::Copy(boost::get(dev_ctx.GetPlace()), + tmp_in_array->ptr(), platform::CPUPlace(), + reinterpret_cast(in_data.data()), + in_data.size() * sizeof(T *), dev_ctx.stream()); + + T **in_array_data = reinterpret_cast(tmp_in_array->ptr()); + ComputeKernelParameter(lod_length); + SumArrayCUDAKernel<<>>( + in_array_data, out->data(), lod_length, in_data.size(), + dst_write | in_place); + } +} + +template +class SumKernel + : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext &context) const override { + auto in_vars = context.MultiInputVar("X"); + const size_t in_num = in_vars.size(); + auto out_var = context.OutputVar("Out"); + bool in_place = out_var == in_vars[0]; + + if (out_var->IsType()) { + FuseLodTensorSumCompute(context); + } else if (out_var->IsType()) { + SelectedRowsCompute(context); + } else if (out_var->IsType()) { + LodTensorArrayCompute(context); + } else { + PADDLE_THROW("Unexpected branch, output variable type is %s", + framework::ToTypeName(out_var->Type())); + } + } +}; +} // namespace operators +} // namespace paddle + namespace ops = paddle::operators; namespace plat = paddle::platform; REGISTER_OP_CUDA_KERNEL( diff --git a/paddle/fluid/operators/sum_op.h b/paddle/fluid/operators/sum_op.h index a8b2df186d..0d60947971 100644 --- a/paddle/fluid/operators/sum_op.h +++ b/paddle/fluid/operators/sum_op.h @@ -27,6 +27,96 @@ template using EigenVector = framework::EigenVector; +template +void SelectedRowsCompute(const framework::ExecutionContext &context) { + auto in_vars = context.MultiInputVar("X"); + auto out_var = context.OutputVar("Out"); + bool in_place = out_var == in_vars[0]; + + if (in_place && in_vars.size() < 2) { + return; + } + + std::vector inputs; + SelectedRows temp_in0; + + if (in_place) { + auto &in0 = in_vars[0]->Get(); + temp_in0.set_height(in0.height()); + temp_in0.set_rows(in0.rows()); + framework::TensorCopy(in0.value(), in0.place(), context.device_context(), + temp_in0.mutable_value()); + inputs.push_back(&temp_in0); + for (size_t i = 1; i < in_vars.size(); ++i) { + auto &in = in_vars[i]->Get(); + if (in.rows().size() > 0) { + inputs.push_back(&in); + } + } + } else { + for (auto &in_var : in_vars) { + auto &in = in_var->Get(); + if (in.rows().size() > 0) { + inputs.push_back(&in_var->Get()); + } + } + } + + auto *out = context.Output("Out"); + out->mutable_rows()->clear(); + + bool has_data = false; + for (auto &in : inputs) { + if (in->rows().size() > 0) { + has_data = true; + break; + } + } + if (has_data) { + math::scatter::MergeAdd merge_add; + merge_add(context.template device_context(), inputs, out); + + out->SyncIndex(); + + } else { + // no data, just set a empty out tensor. + out->mutable_value()->mutable_data(framework::make_ddim({0}), + context.GetPlace()); + } +} + +template +void LodTensorArrayCompute(const framework::ExecutionContext &context) { + auto in_vars = context.MultiInputVar("X"); + auto out_var = context.OutputVar("Out"); + bool in_place = out_var == in_vars[0]; + auto &out_array = *out_var->GetMutable(); + for (size_t i = in_place ? 1 : 0; i < in_vars.size(); ++i) { + PADDLE_ENFORCE(in_vars[i]->IsType(), + "Only support all inputs are TensorArray"); + auto &in_array = in_vars[i]->Get(); + + 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(), + context.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::Flatten(in_array[i]); + auto result = EigenVector::Flatten(out_array[i]); + result.device(*context.template device_context() + .eigen_device()) = result + in; + } + } + } + } +} + template class SumKernel : public framework::OpKernel { public: @@ -83,84 +173,9 @@ class SumKernel : public framework::OpKernel { } } } else if (out_var->IsType()) { - if (in_place && in_vars.size() < 2) { - return; - } - - std::vector inputs; - SelectedRows temp_in0; - - if (in_place) { - auto &in0 = in_vars[0]->Get(); - temp_in0.set_height(in0.height()); - temp_in0.set_rows(in0.rows()); - framework::TensorCopy(in0.value(), in0.place(), - context.device_context(), - temp_in0.mutable_value()); - inputs.push_back(&temp_in0); - for (size_t i = 1; i < in_vars.size(); ++i) { - auto &in = in_vars[i]->Get(); - if (in.rows().size() > 0) { - inputs.push_back(&in); - } - } - } else { - for (auto &in_var : in_vars) { - auto &in = in_var->Get(); - if (in.rows().size() > 0) { - inputs.push_back(&in_var->Get()); - } - } - } - - auto *out = context.Output("Out"); - out->mutable_rows()->clear(); - - bool has_data = false; - for (auto &in : inputs) { - if (in->rows().size() > 0) { - has_data = true; - break; - } - } - if (has_data) { - math::scatter::MergeAdd merge_add; - merge_add(context.template device_context(), inputs, - out); - - out->SyncIndex(); - - } else { - // no data, just set a empty out tensor. - out->mutable_value()->mutable_data(framework::make_ddim({0}), - context.GetPlace()); - } + SelectedRowsCompute(context); } else if (out_var->IsType()) { - auto &out_array = *out_var->GetMutable(); - for (size_t i = in_place ? 1 : 0; i < in_vars.size(); ++i) { - PADDLE_ENFORCE(in_vars[i]->IsType(), - "Only support all inputs are TensorArray"); - auto &in_array = in_vars[i]->Get(); - - 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(), - context.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::Flatten(in_array[i]); - auto result = EigenVector::Flatten(out_array[i]); - result.device(*context.template device_context() - .eigen_device()) = result + in; - } - } - } - } + LodTensorArrayCompute(context); } else { PADDLE_THROW("Unexpected branch, output variable type is %s", framework::ToTypeName(out_var->Type())); -- GitLab