提交 68b5e5bf 编写于 作者: W wanghaoshuang

Use stridecpy instead of CUDA kernel

上级 ce709b75
......@@ -128,59 +128,11 @@ class CropOpGrad : public framework::OperatorWithKernel {
}
};
int64_t transIndex(std::vector<int64_t> out_shape, std::vector<int64_t> x_shape,
std::vector<std::pair<int, int>> crop_rules, size_t index) {
int64_t dim_size = out_shape.size();
std::vector<int64_t> pos(dim_size);
for (int64_t i = out_shape.size() - 1; i >= 0; --i) {
pos[i] = (index % out_shape[i]) + crop_rules[i].first;
index = index / out_shape[i];
}
size_t result = pos[0];
for (size_t i = 1; i < x_shape.size(); ++i) {
result = result * x_shape[i] + pos[i];
}
return result;
}
template <typename T>
class CropCPUKernel : public framework::OpKernel {
public:
void Compute(const framework::ExecutionContext &context) const override {
auto *x = context.Input<Tensor>("X");
auto *out = context.Output<Tensor>("Out");
auto x_data = x->data<T>();
T *out_data = out->mutable_data<T>(context.GetPlace());
auto x_dims = x->dims();
auto out_dims = out->dims();
int64_t out_count = out->numel();
std::vector<int64_t> x_shape = framework::vectorize(x_dims);
std::vector<int64_t> out_shape = framework::vectorize(out_dims);
auto offsets = context.Attr<std::vector<int>>("offsets");
PADDLE_ENFORCE_EQ(
x_dims.size(), offsets.size(),
"Offsets size should be equal to dimension size of input tensor.");
std::vector<std::pair<int, int>> crop_rules(x_dims.size());
for (size_t i = 0; i < crop_rules.size(); ++i) {
crop_rules[i].first = offsets[i];
crop_rules[i].second = x_dims[i] - out_dims[i] - offsets[i];
}
for (int64_t i = 0; i < out_count; ++i) {
out_data[i] = x_data[transIndex(out_shape, x_shape, crop_rules, i)];
}
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
REGISTER_OP(crop, ops::CropOp, ops::CropOpMaker, crop_grad, ops::CropOpGrad);
REGISTER_OP_CPU_KERNEL(crop, ops::CropCPUKernel<float>);
REGISTER_OP_CPU_KERNEL(crop, ops::CropKernel<float>);
REGISTER_OP_CPU_KERNEL(crop_grad,
ops::CropGradKernel<paddle::platform::CPUPlace, float>);
......@@ -13,128 +13,9 @@
limitations under the License. */
#define EIGEN_USE_GPU
#include <stdio.h>
#include "paddle/operators/crop_op.h"
namespace paddle {
namespace operators {
using framework::LoDTensor;
using framework::Tensor;
template <typename T, int D>
__global__ void CropKernel(const int N, const int64_t* out_shape,
const int64_t* x_shape, const int* crop_rules,
const T* x_data, T* out_data) {
int64_t pos[D];
int tmp;
int64_t x_index;
for (int out_index = blockIdx.x * blockDim.x + threadIdx.x; out_index < N;
out_index += blockDim.x * gridDim.x) {
tmp = out_index;
for (int64_t i = D - 1; i >= 0; --i) {
pos[i] = (tmp % out_shape[i]) + crop_rules[i * 2];
tmp = tmp / out_shape[i];
}
x_index = pos[0];
for (size_t i = 1; i < D; ++i) {
x_index = x_index * x_shape[i] + pos[i];
}
out_data[out_index] = x_data[x_index];
}
}
template <typename T, int D>
void CropCUDAFunctoin(const framework::ExecutionContext& context) {
PADDLE_ENFORCE(platform::is_gpu_place(context.GetPlace()),
"It must use GPUPlace.");
auto* x = context.Input<LoDTensor>("X");
auto* out = context.Output<LoDTensor>("Out");
auto x_data = x->data<T>();
T* out_data = out->mutable_data<T>(paddle::platform::GPUPlace());
auto x_dims = x->dims();
auto out_dims = out->dims();
int64_t out_count = out->numel();
Tensor x_shape;
Tensor out_shape;
int64_t* x_shape_data =
x_shape.mutable_data<int64_t>({D}, paddle::platform::CPUPlace());
int64_t* out_shape_data =
out_shape.mutable_data<int64_t>({D}, paddle::platform::CPUPlace());
for (int i = 0; i < D; ++i) {
x_shape_data[i] = x_dims[i];
out_shape_data[i] = out_dims[i];
}
Tensor x_shape_gpu;
Tensor out_shape_gpu;
x_shape_gpu.CopyFrom<int64_t>(x_shape, paddle::platform::GPUPlace());
out_shape_gpu.CopyFrom<int64_t>(out_shape, paddle::platform::GPUPlace());
auto offsets = context.op().Attr<std::vector<int>>("offsets");
PADDLE_ENFORCE_EQ(
D, offsets.size(),
"Offsets size should be equal to dimension size of input tensor.");
Tensor crop_rules;
int* crop_rules_data =
crop_rules.mutable_data<int>({D * 2}, paddle::platform::CPUPlace());
for (size_t i = 0; i < D; ++i) {
crop_rules_data[i * 2] = offsets[i];
crop_rules_data[i * 2 + 1] = x_dims[i] - out_dims[i] - offsets[i];
}
Tensor crop_rules_gpu;
crop_rules_gpu.CopyFrom<int>(crop_rules, paddle::platform::GPUPlace());
int n = out_dims[0];
int d = out_dims[1];
int block = 512;
int grid = (n * d + block - 1) / block;
CropKernel<
T,
D><<<grid, block, 0, reinterpret_cast<const platform::CUDADeviceContext&>(
context.device_context())
.stream()>>>(
out_count, out_shape_gpu.data<int64_t>(), x_shape_gpu.data<int64_t>(),
crop_rules_gpu.data<int>(), x_data, out_data);
}
template <typename T>
class CropOpCUDAKernel : public framework::OpKernel {
public:
void Compute(const framework::ExecutionContext& context) const override {
size_t rank = context.Input<LoDTensor>("X")->dims().size();
switch (rank) {
case 1:
CropCUDAFunctoin<T, 1>(context);
break;
case 2:
CropCUDAFunctoin<T, 2>(context);
break;
case 3:
CropCUDAFunctoin<T, 3>(context);
break;
case 4:
CropCUDAFunctoin<T, 4>(context);
break;
case 5:
CropCUDAFunctoin<T, 5>(context);
break;
case 6:
CropCUDAFunctoin<T, 6>(context);
break;
default:
PADDLE_THROW(
"CropOp only support tensors with no more than 6 dimensions.");
}
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
REGISTER_OP_GPU_KERNEL(crop, ops::CropOpCUDAKernel<float>);
REGISTER_OP_GPU_KERNEL(crop, ops::CropKernel<float>);
REGISTER_OP_GPU_KERNEL(crop_grad,
ops::CropGradKernel<paddle::platform::GPUPlace, float>);
......@@ -16,6 +16,7 @@
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"
#include "paddle/operators/strided_memcpy.h"
namespace paddle {
namespace operators { // Internal
......@@ -24,26 +25,58 @@ template <typename T, size_t D, int MajorType = Eigen::RowMajor,
typename IndexType = Eigen::DenseIndex>
using EigenTensor = framework::EigenTensor<T, D, MajorType, IndexType>;
using framework::LoDTensor;
using framework::Tensor;
using framework::DDim;
// TODO(wanghaoshuang): move this function to other place
DDim stride(const DDim& ddim) {
std::vector<int64_t> strides(ddim.size());
strides[ddim.size() - 1] = 1;
for (int i = ddim.size() - 2; i >= 0; --i) {
strides[i] = strides[i + 1] * ddim[i + 1];
}
return make_ddim(strides);
}
template <typename T>
class CropKernel : public framework::OpKernel {
public:
void Compute(const framework::ExecutionContext& context) const override {
auto* x = context.Input<Tensor>("X");
auto* out = context.Output<Tensor>("Out");
T* x_data = x->data<T>();
T* out_data = out->mutable_data<T>(context.GetPlace());
auto x_stride = stride(x->dims());
auto out_stride = stride(out->dims());
auto offsets = context.Attr<std::vector<int>>("offsets");
PADDLE_ENFORCE_EQ(
x_dims.size(), offsets.size(),
"Offsets size should be equal to dimension size of input tensor.");
int64_t offset = 0;
for (int i = 0; i < offsets.size(); ++i) {
offset += (x_stride[i] * offsets[i]);
}
StridedMemcpy<T>(context.device_context(), x_data + offset, x_stride,
out->dims(), out_stride, out_data);
}
};
template <typename Place, typename T, size_t D>
void CropGradFunction(const framework::ExecutionContext& context) {
auto* d_out = context.Input<LoDTensor>(framework::GradVarName("Out"));
auto* d_x = context.Output<LoDTensor>(framework::GradVarName("X"));
auto* d_x = context.Output<Tensor>(framework::GradVarName("X"));
if (d_x != nullptr) {
auto* d_out = context.Input<Tensor>(framework::GradVarName("Out"));
d_x->mutable_data<T>(context.GetPlace());
auto d_x_dims = d_x->dims();
auto d_out_dims = d_out->dims();
auto offsets = context.op().Attr<std::vector<int>>("offsets");
auto offsets = context.Attr<std::vector<int>>("offsets");
Eigen::array<std::pair<int, int>, D> paddings;
for (int i = 0; i < d_out_dims.size(); ++i) {
for (int i = 0; i < D; ++i) {
paddings[i].first = offsets[i];
paddings[i].second = d_x_dims[i] - d_out_dims[i] - offsets[i];
}
auto d_x_tensor = EigenTensor<T, D>::From(*d_x);
auto d_out_tensor = EigenTensor<T, D>::From(*d_out);
auto place = context.GetEigenDevice<Place>();
d_x_tensor.device(place) = d_out_tensor.pad(paddings, 0);
d_x_tensor.device(context.GetEigenDevice<Place>()) =
d_out_tensor.pad(paddings, 0);
}
}
......@@ -52,7 +85,7 @@ class CropGradKernel : public framework::OpKernel {
public:
void Compute(const framework::ExecutionContext& context) const override {
size_t rank =
context.Input<LoDTensor>(framework::GradVarName("Out"))->dims().size();
context.Input<Tensor>(framework::GradVarName("Out"))->dims().size();
switch (rank) {
case 1:
CropGradFunction<Place, T, 1>(context);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册