diff --git a/paddle/fluid/operators/slice_op.h b/paddle/fluid/operators/slice_op.h index 9c30c4e07fa774f4da65f1e78bc74252fa3b2e86..22f6fa9e3e6f206b33c46369086d1637fdc83457 100644 --- a/paddle/fluid/operators/slice_op.h +++ b/paddle/fluid/operators/slice_op.h @@ -259,7 +259,20 @@ class SliceKernel : public framework::OpKernel { auto out_t = framework::EigenTensor::From( *out, new_out_dims); - out_t.device(place) = in_t.slice(offsets, extents); + + if (in->numel() <= Eigen::NumTraits::highest()) { + // similar to tf.slice: + // if element number less than INT_MAX, change the type of index to int + Eigen::DSizes offsets_32bit, extents_32bit; + for (size_t i = 0; i < D; i++) { + offsets_32bit[i] = offsets[i]; + extents_32bit[i] = extents[i]; + } + framework::To32BitIndex(out_t).device(place) = + framework::To32BitIndex(in_t).slice(offsets_32bit, extents_32bit); + } else { + out_t.device(place) = in_t.slice(offsets, extents); + } out->Resize(out_dims); } @@ -300,8 +313,6 @@ class SliceGradKernel : public framework::OpKernel { private: template void SliceCompute(const framework::ExecutionContext& context) const { - auto& place = - *context.template device_context().eigen_device(); auto axes = context.Attr>("axes"); auto starts_int = context.Attr>("starts"); @@ -435,13 +446,189 @@ class SliceGradKernel : public framework::OpKernel { paddings[i].first = offsets[i]; paddings[i].second = (in_dims[i] - out_dims[i]) - offsets[i]; } + EigenPaddingCompute(context, d_input, in_dims, d_out, out_dims, paddings); + } + + template + void EigenPaddingCompute( + const framework::ExecutionContext& context, framework::Tensor* d_input, + const framework::DDim& in_dims, const framework::Tensor* d_out, + const framework::DDim& out_dims, + const Eigen::array, D>& paddings) const { + if (D <= 3) { + // if dimension less than 3, cannot reduce dimension + LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims, paddings); + } else { // else we can reduce dimension + // count not-zero padding number, and record the dimension + int need_pad_num = 0, pad_dim = -1; + for (size_t i = 0; i < D; i++) { + if (paddings[i].first != 0 || paddings[i].second != 0) { + need_pad_num++; + pad_dim = i; + } + } + + if (need_pad_num == 0) { + // do not need padding, pass if data address same, else copy + if (d_input->mutable_data(context.GetPlace()) == d_out->data()) { + // inplace, do not any operator, pass + } else { + framework::TensorCopy( + *d_out, context.GetPlace(), + context.template device_context(), + d_input); + } + } else if (need_pad_num == 1) { + // only need padding one dimension, we can reduce dimension. + // only the padding dimension is available for us. + // How to reduce dimension(5 to 3 for example): + // before(D=5): + // in_dims: [x1, x2, x3, x4, x5] + // padding.first: [0, 0, a, 0, 0] + // padding.second: [0, 0, b, 0, 0] + // | | + // V V + // after(D=3): + // reshaped_in_dims: [x1*x2, x3, x4*x5] + // reshaped_padding.first: [0, a, 0] + // reshaped_padding.second: [0, b, 0] + + if (pad_dim == D - 1) { + // only last dimension need padding, + // reshape the dimension of tensor in 2: [preceding, padding] + std::vector in_tore_shape(2, 1), out_tore_shape(2, 1); + Eigen::array, 2> reshaped_padding; + + // first dimension is the accumulate of preceding dimension + for (int i = 0; i < pad_dim; i++) { + in_tore_shape[0] *= in_dims[i]; + out_tore_shape[0] *= out_dims[i]; + } + // second dimension is the padding dimension + in_tore_shape[1] = in_dims[pad_dim]; + out_tore_shape[1] = out_dims[pad_dim]; + + // convert array from std::vector to DDim + framework::DDim reshaped_in_dims = + framework::make_ddim(in_tore_shape); + framework::DDim reshaped_out_dims = + framework::make_ddim(out_tore_shape); + + // after reshape: the first dimension do not need padding, + // set padding[0] zero + reshaped_padding[0].first = reshaped_padding[0].second = 0; + // the second dimension is the previous padding dimension + reshaped_padding[1].first = paddings[pad_dim].first; + reshaped_padding[1].second = paddings[pad_dim].second; + + LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out, + reshaped_out_dims, reshaped_padding); + } else if (pad_dim == 0) { + // only first dimension need padding, + // reshape the dimension of tensor in 2: [padding, succeeding] + // similar to (D - 1) + std::vector in_tore_shape(2, 1), out_tore_shape(2, 1); + Eigen::array, 2> reshaped_padding; + + // first dimension is the padding dimension + in_tore_shape[0] = in_dims[pad_dim]; + out_tore_shape[0] = out_dims[pad_dim]; + // sencond dimension is the accumulate of succeeding dimension + for (size_t i = pad_dim + 1; i < D; i++) { + in_tore_shape[1] *= in_dims[i]; + out_tore_shape[1] *= out_dims[i]; + } + + // convert array from std::vector to DDim + framework::DDim reshaped_in_dims = + framework::make_ddim(in_tore_shape); + framework::DDim reshaped_out_dims = + framework::make_ddim(out_tore_shape); + + // after reshape: + // the first dimension is the previous padding dimension + reshaped_padding[0].first = paddings[pad_dim].first; + reshaped_padding[0].second = paddings[pad_dim].second; + // the second dimension do not need padding, set padding[1] zero + reshaped_padding[1].first = reshaped_padding[1].second = 0; + + LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out, + reshaped_out_dims, reshaped_padding); + } else { + // other dimension need padding + // reshape the dimension of tensor in 3: + // [preceding, padding, succeeding] + std::vector in_tore_shape(3, 1), out_tore_shape(3, 1); + Eigen::array, 3> reshaped_padding; + + // first dimension is the accumulate of preceding dimension + for (int i = 0; i < pad_dim; i++) { + in_tore_shape[0] *= in_dims[i]; + out_tore_shape[0] *= out_dims[i]; + } + // second dimension is the padding dimension + in_tore_shape[1] = in_dims[pad_dim]; + out_tore_shape[1] = out_dims[pad_dim]; + // third dimension is the accumulate of succeeding dimension + for (size_t i = pad_dim + 1; i < D; i++) { + in_tore_shape[2] *= in_dims[i]; + out_tore_shape[2] *= out_dims[i]; + } + + // convert array from std::vector to DDim + framework::DDim reshaped_in_dims = + framework::make_ddim(in_tore_shape); + framework::DDim reshaped_out_dims = + framework::make_ddim(out_tore_shape); + + // after reshape: + // the first dimension do not need padding, set padding[0] zero + reshaped_padding[0].first = reshaped_padding[2].second = 0; + // the second dimension is the previous padding dimension + reshaped_padding[1].first = paddings[pad_dim].first; + reshaped_padding[1].second = paddings[pad_dim].second; + // the third dimension do not need padding, set padding[2] zero + reshaped_padding[2].first = reshaped_padding[2].second = 0; + + LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out, + reshaped_out_dims, reshaped_padding); + } + } else { + // need padding at many dimension, cannot reduce dimension + LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims, + paddings); + } + } + } + + template + void LaunchEigenPadding( + const framework::ExecutionContext& context, framework::Tensor* d_input, + const framework::DDim& in_dims, const framework::Tensor* d_out, + const framework::DDim& out_dims, + const Eigen::array, D>& paddings) const { + auto& place = + *context.template device_context().eigen_device(); auto d_in_t = framework::EigenTensor::From( - *d_input); + *d_input, in_dims); auto d_out_t = framework::EigenTensor::From( *d_out, out_dims); - d_in_t.device(place) = d_out_t.pad(paddings, T(0)); + + if (d_input->numel() <= Eigen::NumTraits::highest()) { + // similar to tf.pad: + // if element number less than INT_MAX, change the type of index to int + Eigen::array, D> paddings_32bit; + for (size_t i = 0; i < D; i++) { + paddings_32bit[i] = + std::make_pair(paddings[i].first, paddings[i].second); + } + framework::To32BitIndex(d_in_t).device(place) = + framework::To32BitIndex(d_out_t).pad(paddings_32bit, T(0)); + } else { + d_in_t.device(place) = d_out_t.pad(paddings, T(0)); + } } }; } // namespace operators