diff --git a/paddle/phi/kernels/cpu/randint_kernel.cc b/paddle/phi/kernels/cpu/randint_kernel.cc index 5fe56b57452d5a1afd718d85a038310e27f0ff50..feb418949ba40d3bf553c2df0b4300cc686a0ef7 100644 --- a/paddle/phi/kernels/cpu/randint_kernel.cc +++ b/paddle/phi/kernels/cpu/randint_kernel.cc @@ -22,42 +22,43 @@ namespace phi { template -void RandintRawKernel(const Context& ctx, +void RandintRawKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, DataType dtype, int seed, DenseTensor* out) { - out->ResizeAndAllocate(phi::make_ddim(shape.GetData())); - auto size = out->numel(); + out->Resize(phi::make_ddim(shape.GetData())); + T* data = dev_ctx.template Alloc(out); + auto numel = out->numel(); std::shared_ptr engine; if (seed) { engine = std::make_shared(); engine->seed(seed); } else { - engine = ctx.GetGenerator()->GetCPUEngine(); + engine = dev_ctx.GetGenerator()->GetCPUEngine(); } std::uniform_int_distribution dist(low, high - 1); - auto data = out->data(); - for (int64_t i = 0; i < size; ++i) { + for (int64_t i = 0; i < numel; ++i) { data[i] = dist(*engine); } } template -void RandintKernel(const Context& ctx, +void RandintKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, DataType dtype, DenseTensor* out) { - RandintRawKernel(ctx, low, high, shape, dtype, 0, out); + RandintRawKernel(dev_ctx, low, high, shape, dtype, 0, out); } } // namespace phi PD_REGISTER_KERNEL( randint_raw, CPU, ALL_LAYOUT, phi::RandintRawKernel, int, int64_t) {} + PD_REGISTER_KERNEL(randint, CPU, ALL_LAYOUT, phi::RandintKernel, int, int64_t) { } diff --git a/paddle/phi/kernels/cpu/randperm_kernel.cc b/paddle/phi/kernels/cpu/randperm_kernel.cc index 28092c8df6d153c6b5e787027f0c2239bd257cc1..6cb435f53b85bd22afba1a0d31b16ecd4c27204b 100644 --- a/paddle/phi/kernels/cpu/randperm_kernel.cc +++ b/paddle/phi/kernels/cpu/randperm_kernel.cc @@ -13,20 +13,23 @@ // limitations under the License. #include "paddle/phi/kernels/randperm_kernel.h" -#include "paddle/phi/core/dense_tensor.h" -#include "paddle/phi/core/device_context.h" + #include "paddle/phi/core/kernel_registry.h" namespace phi { template -void RandpermKernel(const Context& ctx, - int n, - DataType dtype, - DenseTensor* out) { - T* out_data = ctx.template Alloc(out); - auto gen_ptr = ctx.GetHostGenerator(); - auto engine = gen_ptr->GetCPUEngine(); +void RandpermRawKernel( + const Context& dev_ctx, int n, DataType dtype, int seed, DenseTensor* out) { + T* out_data = dev_ctx.template Alloc(out); + + std::shared_ptr engine; + if (seed) { + engine = std::make_shared(); + engine->seed(seed); + } else { + engine = dev_ctx.GetGenerator()->GetCPUEngine(); + } for (int i = 0; i < n; ++i) { out_data[i] = static_cast(i); @@ -34,8 +37,25 @@ void RandpermKernel(const Context& ctx, std::shuffle(out_data, out_data + n, *engine); } +template +void RandpermKernel(const Context& dev_ctx, + int n, + DataType dtype, + DenseTensor* out) { + RandpermRawKernel(dev_ctx, n, dtype, 0, out); +} + } // namespace phi +PD_REGISTER_KERNEL(randperm_raw, + CPU, + ALL_LAYOUT, + phi::RandpermRawKernel, + float, + double, + int, + int64_t) {} + PD_REGISTER_KERNEL(randperm, CPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/cpu/unbind_kernel.cc b/paddle/phi/kernels/cpu/unbind_kernel.cc index 655f8c8aafbf201dc07db0fa1af79605c2a76763..39cc2f8fc4662a0893fb8b73b138a52b810f59b8 100644 --- a/paddle/phi/kernels/cpu/unbind_kernel.cc +++ b/paddle/phi/kernels/cpu/unbind_kernel.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "paddle/phi/kernels/unbind_kernel.h" + #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/impl/unbind_kernel_impl.h" diff --git a/paddle/phi/kernels/funcs/concat_and_split_functor.cc b/paddle/phi/kernels/funcs/concat_and_split_functor.cc index c8405703a5c16ae9eae583638d1c89c22a736531..aa73ba5f689906e73f3f0e3a845aa397ad0a33c1 100644 --- a/paddle/phi/kernels/funcs/concat_and_split_functor.cc +++ b/paddle/phi/kernels/funcs/concat_and_split_functor.cc @@ -12,21 +12,6 @@ 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 -#include - -#include "paddle/fluid/framework/convert_utils.h" -#include "paddle/fluid/framework/eigen.h" -#include "paddle/fluid/framework/operator.h" -#include "paddle/fluid/framework/tensor.h" -#include "paddle/fluid/framework/tensor_util.h" -#include "paddle/fluid/platform/device_context.h" -#include "paddle/fluid/platform/enforce.h" -#include "paddle/phi/core/utils/data_type.h" - -#include "paddle/phi/backends/cpu/cpu_context.h" -#include "paddle/phi/core/dense_tensor.h" #include "paddle/phi/kernels/funcs/concat_and_split_functor.h" namespace phi { diff --git a/paddle/phi/kernels/funcs/concat_and_split_functor.cu b/paddle/phi/kernels/funcs/concat_and_split_functor.cu index 2abfdb606e7e6c410f6f9deb45aed536bea88207..840c8872f50f83c2859f07be2e0e7242a74004a7 100644 --- a/paddle/phi/kernels/funcs/concat_and_split_functor.cu +++ b/paddle/phi/kernels/funcs/concat_and_split_functor.cu @@ -12,23 +12,11 @@ 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 -#include - -#include "paddle/fluid/framework/convert_utils.h" -#include "paddle/fluid/framework/eigen.h" -#include "paddle/fluid/framework/operator.h" -#include "paddle/fluid/framework/tensor.h" -#include "paddle/fluid/framework/tensor_util.h" -#include "paddle/fluid/platform/cuda_graph_with_memory_pool.h" -#include "paddle/fluid/platform/device_context.h" -#include "paddle/fluid/platform/enforce.h" - -#include "paddle/phi/backends/cpu/cpu_context.h" -#include "paddle/phi/core/utils/data_type.h" #include "paddle/phi/kernels/funcs/concat_and_split_functor.h" +#include "paddle/fluid/memory/malloc.h" +#include "paddle/fluid/platform/cuda_graph_with_memory_pool.h" + namespace phi { namespace funcs { diff --git a/paddle/phi/kernels/funcs/concat_and_split_functor.h b/paddle/phi/kernels/funcs/concat_and_split_functor.h index 3af4d878d3cab03eb80a6ba878cc4fa5a62103c9..4cb15fe539b66b8a6fddccf18d92b95976db2a65 100644 --- a/paddle/phi/kernels/funcs/concat_and_split_functor.h +++ b/paddle/phi/kernels/funcs/concat_and_split_functor.h @@ -13,20 +13,18 @@ See the License for the specific language governing permissions and limitations under the License. */ #pragma once -#include -#include + #include -#include "paddle/fluid/framework/convert_utils.h" -#include "paddle/fluid/framework/eigen.h" -#include "paddle/fluid/framework/operator.h" -#include "paddle/fluid/framework/tensor.h" -#include "paddle/fluid/framework/tensor_util.h" -#include "paddle/fluid/platform/device_context.h" -#include "paddle/fluid/platform/enforce.h" +#include "paddle/phi/backends/cpu/cpu_context.h" +#include "paddle/phi/backends/gpu/gpu_context.h" #include "paddle/phi/core/dense_tensor.h" +#include "paddle/phi/core/device_context.h" #include "paddle/phi/core/utils/data_type.h" +// See Note [ Why still include the fluid headers? ] +#include "paddle/fluid/memory/memcpy.h" + namespace phi { namespace funcs { diff --git a/paddle/phi/kernels/gpu/randint_kernel.cu b/paddle/phi/kernels/gpu/randint_kernel.cu index b89b714c73d92c896cde3beef182703018b6aa12..66dc5f72a5c7067a08127bce65740851b123efd3 100644 --- a/paddle/phi/kernels/gpu/randint_kernel.cu +++ b/paddle/phi/kernels/gpu/randint_kernel.cu @@ -25,7 +25,7 @@ namespace phi { template -void RandintRawKernel(const Context& ctx, +void RandintRawKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, @@ -34,21 +34,22 @@ void RandintRawKernel(const Context& ctx, DenseTensor* out) { DenseTensor tmp; tmp.Resize(phi::make_ddim(shape.GetData())); - T* tmp_data = ctx.template HostAlloc(&tmp); + T* tmp_data = dev_ctx.template HostAlloc(&tmp); - out->ResizeAndAllocate(tmp.dims()); - auto size = out->numel(); + out->Resize(tmp.dims()); + T* data = dev_ctx.template Alloc(out); std::shared_ptr engine; if (seed) { engine = std::make_shared(); engine->seed(seed); } else { - engine = ctx.GetHostGenerator()->GetCPUEngine(); + engine = dev_ctx.GetHostGenerator()->GetCPUEngine(); } + std::uniform_int_distribution dist(low, high - 1); - auto data = out->data(); - for (int64_t i = 0; i < size; ++i) { + auto numel = out->numel(); + for (int64_t i = 0; i < numel; ++i) { tmp_data[i] = dist(*engine); } @@ -57,18 +58,18 @@ void RandintRawKernel(const Context& ctx, data, tmp.place(), tmp_data, - size * paddle::experimental::SizeOf(out->dtype()), + numel * paddle::experimental::SizeOf(out->dtype()), 0); } template -void RandintKernel(const Context& ctx, +void RandintKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, DataType dtype, DenseTensor* out) { - RandintRawKernel(ctx, low, high, shape, dtype, 0, out); + RandintRawKernel(dev_ctx, low, high, shape, dtype, 0, out); } } // namespace phi diff --git a/paddle/phi/kernels/gpu/randperm_kernel.cu b/paddle/phi/kernels/gpu/randperm_kernel.cu index f75f768b633a31a9d3d6eadcf036640f50309a8b..d4d90cac917a2c35e26eca0d57d1c5349b878599 100644 --- a/paddle/phi/kernels/gpu/randperm_kernel.cu +++ b/paddle/phi/kernels/gpu/randperm_kernel.cu @@ -12,41 +12,60 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "paddle/phi/core/dense_tensor.h" -#include "paddle/phi/core/device_context.h" #include "paddle/phi/kernels/randperm_kernel.h" +#include "paddle/phi/core/kernel_registry.h" + // See Note [ Why still include the fluid headers? ] #include "paddle/fluid/memory/memcpy.h" -#include "paddle/phi/core/kernel_registry.h" namespace phi { template -void RandpermKernel(const Context& ctx, - int n, - DataType dtype, - DenseTensor* out) { +void RandpermRawKernel( + const Context& dev_ctx, int n, DataType dtype, int seed, DenseTensor* out) { DenseTensor tmp; tmp.Resize(phi::make_ddim({n})); - T* tmp_data = ctx.template HostAlloc(&tmp); + T* tmp_data = dev_ctx.template HostAlloc(&tmp); - auto gen_ptr = ctx.GetHostGenerator(); - auto engine = gen_ptr->GetCPUEngine(); + std::shared_ptr engine; + if (seed) { + engine = std::make_shared(); + engine->seed(seed); + } else { + engine = dev_ctx.GetHostGenerator()->GetCPUEngine(); + } for (int i = 0; i < n; ++i) { tmp_data[i] = static_cast(i); } std::shuffle(tmp_data, tmp_data + n, *engine); - T* out_data = ctx.template Alloc(out); + T* out_data = dev_ctx.template Alloc(out); auto size = out->numel() * paddle::experimental::SizeOf(out->dtype()); paddle::memory::Copy( out->place(), out_data, tmp.place(), tmp_data, size, 0); } +template +void RandpermKernel(const Context& dev_ctx, + int n, + DataType dtype, + DenseTensor* out) { + RandpermRawKernel(dev_ctx, n, dtype, 0, out); +} + } // namespace phi +PD_REGISTER_KERNEL(randperm_raw, + GPU, + ALL_LAYOUT, + phi::RandpermRawKernel, + float, + double, + int, + int64_t) {} + PD_REGISTER_KERNEL(randperm, GPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/gpu/unbind_kernel.cu b/paddle/phi/kernels/gpu/unbind_kernel.cu index 1efc3a1094da253c27fec5108b536837d868425e..8a7aa8f6033ab9b86f87e792bc37f912562578a7 100644 --- a/paddle/phi/kernels/gpu/unbind_kernel.cu +++ b/paddle/phi/kernels/gpu/unbind_kernel.cu @@ -12,9 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include "paddle/phi/kernels/unbind_kernel.h" + #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/impl/unbind_kernel_impl.h" -#include "paddle/phi/kernels/unbind_kernel.h" PD_REGISTER_KERNEL(unbind, GPU, diff --git a/paddle/phi/kernels/impl/unbind_kernel_impl.h b/paddle/phi/kernels/impl/unbind_kernel_impl.h index 8a1342559bd908bf197e4949ce66f9b3e504b499..3e233a2038e48098d8c78bf81d922a812a87187a 100644 --- a/paddle/phi/kernels/impl/unbind_kernel_impl.h +++ b/paddle/phi/kernels/impl/unbind_kernel_impl.h @@ -20,7 +20,7 @@ namespace phi { template -void UnbindKernel(const Context& ctx, +void UnbindKernel(const Context& dev_ctx, const DenseTensor& x, int axis, std::vector outs) { @@ -29,12 +29,12 @@ void UnbindKernel(const Context& ctx, std::vector shape_refer; for (size_t j = 0; j < outs.size(); ++j) { - ctx.template Alloc(outs[j]); + dev_ctx.template Alloc(outs[j]); shape_refer.emplace_back(outs[j]); } phi::funcs::SplitFunctor functor; - functor(ctx, x, shape_refer, axis, &outs); + functor(dev_ctx, x, shape_refer, axis, &outs); } } // namespace phi diff --git a/paddle/phi/kernels/randint_kernel.h b/paddle/phi/kernels/randint_kernel.h index 1a78e73d863e33619d10b72bf9e368aba4c856c5..bfefc628614fbc03c484a43f31a7194da15a2bf9 100644 --- a/paddle/phi/kernels/randint_kernel.h +++ b/paddle/phi/kernels/randint_kernel.h @@ -20,7 +20,7 @@ namespace phi { template -void RandintKernel(const Context& ctx, +void RandintKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, @@ -28,7 +28,7 @@ void RandintKernel(const Context& ctx, DenseTensor* out); template -void RandintRawKernel(const Context& ctx, +void RandintRawKernel(const Context& dev_ctx, int low, int high, const ScalarArray& shape, diff --git a/paddle/phi/kernels/randperm_kernel.h b/paddle/phi/kernels/randperm_kernel.h index 63bdac6da6fdc12955e3743f23941840696a0ce6..70b95db98bef95f364802ad14a54966dc47d13fe 100644 --- a/paddle/phi/kernels/randperm_kernel.h +++ b/paddle/phi/kernels/randperm_kernel.h @@ -20,7 +20,11 @@ namespace phi { template -void RandpermKernel(const Context& ctx, +void RandpermRawKernel( + const Context& dev_ctx, int n, DataType dtype, int seed, DenseTensor* out); + +template +void RandpermKernel(const Context& dev_ctx, int n, DataType dtype, DenseTensor* out); diff --git a/paddle/phi/ops/compat/randperm_sig.cc b/paddle/phi/ops/compat/randperm_sig.cc index 14b28512e402a377b4f2f8f7d8f1e90f7ef37b71..89548beff6762a59edbe317383714502b4382efe 100644 --- a/paddle/phi/ops/compat/randperm_sig.cc +++ b/paddle/phi/ops/compat/randperm_sig.cc @@ -17,7 +17,12 @@ namespace phi { KernelSignature RandpermOpArgumentMapping(const ArgumentMappingContext& ctx) { - return KernelSignature("randperm", {}, {"n", "dtype"}, {"Out"}); + int seed = paddle::any_cast(ctx.Attr("seed")); + if (seed) { + return KernelSignature("randperm", {}, {"n", "dtype", "seed"}, {"Out"}); + } else { + return KernelSignature("randperm", {}, {"n", "dtype"}, {"Out"}); + } } } // namespace phi