[phi] move randint to phi (#39872)

* move randint to phi

* use host generator

paddle/fluid/operators/randint_op.cc

@@ -24,37 +24,6 @@
 namespace paddle {
 namespace operators {
 
-template <typename T>
-class CPURandintKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    std::vector<int64_t> new_shape;
-    auto list_new_shape_tensor =
-        ctx.MultiInput<framework::Tensor>("ShapeTensorList");
-    if (list_new_shape_tensor.size() > 0 || ctx.HasInput("ShapeTensor")) {
-      if (ctx.HasInput("ShapeTensor")) {
-        auto* shape_tensor = ctx.Input<framework::Tensor>("ShapeTensor");
-        new_shape = GetNewDataFromShapeTensor(shape_tensor);
-      } else if (list_new_shape_tensor.size() > 0) {
-        new_shape = GetNewDataFromShapeTensorList(list_new_shape_tensor);
-      }
-    }
-    auto* out = ctx.Output<framework::LoDTensor>("Out");
-    if (!new_shape.empty()) out->Resize(phi::make_ddim(new_shape));
-    T* data = out->mutable_data<T>(ctx.GetPlace());
-    int64_t size = out->numel();
-
-    std::uniform_int_distribution<T> dist(ctx.Attr<int>("low"),
-                                          ctx.Attr<int>("high") - 1);
-    unsigned int seed = static_cast<unsigned int>(ctx.Attr<int>("seed"));
-    auto engine = framework::GetCPURandomEngine(seed);
-
-    for (int64_t i = 0; i < size; ++i) {
-      data[i] = dist(*engine);
-    }
-  }
-};
-
 class RandintOp : public framework::OperatorWithKernel {
  public:
   using framework::OperatorWithKernel::OperatorWithKernel;
@@ -176,6 +145,3 @@ REGISTER_OPERATOR(
     randint, ops::RandintOp, ops::RandintOpMaker,
     paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
     paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>)
-
-REGISTER_OP_CPU_KERNEL(randint, ops::CPURandintKernel<int>,
-                       ops::CPURandintKernel<int64_t>)

paddle/fluid/operators/randint_op.cu

-// Copyright (c) 2020 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.
-#include <thrust/random.h>
-#include <thrust/transform.h>
-#include "paddle/fluid/framework/convert_utils.h"
-#include "paddle/fluid/framework/generator.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/operators/uniform_random_op.h"
-
-namespace paddle {
-namespace operators {
-
-template <typename T>
-class GPURandintKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    std::vector<int64_t> new_shape;
-    auto list_new_shape_tensor =
-        context.MultiInput<framework::Tensor>("ShapeTensorList");
-    if (list_new_shape_tensor.size() > 0 || context.HasInput("ShapeTensor")) {
-      if (context.HasInput("ShapeTensor")) {
-        auto* shape_tensor = context.Input<framework::Tensor>("ShapeTensor");
-        new_shape = GetNewDataFromShapeTensor(shape_tensor);
-      } else if (list_new_shape_tensor.size() > 0) {
-        new_shape = GetNewDataFromShapeTensorList(list_new_shape_tensor);
-      }
-    }
-
-    platform::CPUPlace cpu;
-    auto dtype = static_cast<framework::proto::VarType::Type>(
-        context.Attr<int>("dtype"));
-    auto* out = context.Output<framework::LoDTensor>("Out");
-    if (!new_shape.empty()) out->Resize(phi::make_ddim(new_shape));
-    T low = static_cast<T>(context.Attr<int>("low"));
-    T high = static_cast<T>(context.Attr<int>("high")) - 1;
-    framework::LoDTensor tensor;
-    tensor.Resize(out->dims());
-    tensor.mutable_data(cpu, framework::TransToPtenDataType(dtype));
-    T* data = tensor.mutable_data<T>(cpu);
-
-    int64_t size = out->numel();
-    unsigned int seed = static_cast<unsigned int>(context.Attr<int>("seed"));
-
-    /*
-    std::minstd_rand engine;
-    if (seed == 0) {
-      std::random_device rd;
-      seed = rd();
-    }
-    engine.seed(seed);
-    */
-
-    std::uniform_int_distribution<> dist(context.Attr<int>("low"),
-                                         context.Attr<int>("high") - 1);
-    auto engine = framework::GetCPURandomEngine(seed);
-
-    for (int64_t i = 0; i < size; ++i) {
-      data[i] = dist(*engine);
-    }
-
-    if (platform::is_gpu_place(context.GetPlace())) {
-      // Copy tensor to out
-      framework::TensorCopy(tensor, context.GetPlace(), out);
-    }
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-REGISTER_OP_CUDA_KERNEL(randint, ops::GPURandintKernel<int>,
-                        ops::GPURandintKernel<int64_t>)

paddle/phi/kernels/cpu/randint_kernel.cc

+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/randint_kernel.h"
+
+#include <random>
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RandintRawKernel(const Context& 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();
+  std::shared_ptr<std::mt19937_64> engine;
+  if (seed) {
+    engine = std::make_shared<std::mt19937_64>();
+    engine->seed(seed);
+  } else {
+    engine = ctx.GetGenerator()->GetCPUEngine();
+  }
+  std::uniform_int_distribution<T> dist(low, high - 1);
+  auto data = out->data<T>();
+  for (int64_t i = 0; i < size; ++i) {
+    data[i] = dist(*engine);
+  }
+}
+
+template <typename T, typename Context>
+void RandintKernel(const Context& ctx,
+                   int low,
+                   int high,
+                   const ScalarArray& shape,
+                   DataType dtype,
+                   DenseTensor* out) {
+  RandintRawKernel<T>(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) {
+}

paddle/phi/kernels/gpu/randint_kernel.cu

+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/randint_kernel.h"
+
+#include <random>
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+// See Note [ Why still include the fluid headers? ]
+#include "paddle/fluid/memory/memcpy.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RandintRawKernel(const Context& ctx,
+                      int low,
+                      int high,
+                      const ScalarArray& shape,
+                      DataType dtype,
+                      int seed,
+                      DenseTensor* out) {
+  DenseTensor tmp;
+  tmp.Resize(phi::make_ddim(shape.GetData()));
+  T* tmp_data = ctx.template HostAlloc<T>(&tmp);
+
+  out->ResizeAndAllocate(tmp.dims());
+  auto size = out->numel();
+
+  std::shared_ptr<std::mt19937_64> engine;
+  if (seed) {
+    engine = std::make_shared<std::mt19937_64>();
+    engine->seed(seed);
+  } else {
+    engine = ctx.GetHostGenerator()->GetCPUEngine();
+  }
+  std::uniform_int_distribution<T> dist(low, high - 1);
+  auto data = out->data<T>();
+  for (int64_t i = 0; i < size; ++i) {
+    tmp_data[i] = dist(*engine);
+  }
+
+  paddle::memory::Copy<phi::GPUPlace, phi::Place>(
+      out->place(),
+      data,
+      tmp.place(),
+      tmp_data,
+      size * paddle::experimental::SizeOf(out->dtype()),
+      0);
+}
+
+template <typename T, typename Context>
+void RandintKernel(const Context& ctx,
+                   int low,
+                   int high,
+                   const ScalarArray& shape,
+                   DataType dtype,
+                   DenseTensor* out) {
+  RandintRawKernel<T>(ctx, low, high, shape, dtype, 0, out);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    randint_raw, GPU, ALL_LAYOUT, phi::RandintRawKernel, int, int64_t) {}
+
+PD_REGISTER_KERNEL(randint, GPU, ALL_LAYOUT, phi::RandintKernel, int, int64_t) {
+}

paddle/phi/kernels/randint_kernel.h

+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar_array.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RandintKernel(const Context& ctx,
+                   int low,
+                   int high,
+                   const ScalarArray& shape,
+                   DataType dtype,
+                   DenseTensor* out);
+
+template <typename T, typename Context>
+void RandintRawKernel(const Context& ctx,
+                      int low,
+                      int high,
+                      const ScalarArray& shape,
+                      DataType dtype,
+                      int seed,
+                      DenseTensor* out);
+
+}  // namespace phi

paddle/phi/ops/compat/randint_sig.cc

+/* Copyright (c) 2022 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. */
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature RandintOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  int seed = paddle::any_cast<int>(ctx.Attr("seed"));
+  if (seed) {
+    if (ctx.InputSize("ShapeTensorList") > 0) {
+      return KernelSignature(
+          "randint_raw",
+          {},
+          {"low", "high", "ShapeTensorList", "seed", "dtype"},
+          {"Out"});
+    } else {
+      const auto& shape =
+          paddle::any_cast<std::vector<int64_t>>(ctx.Attr("shape"));
+      if (ctx.HasInput("ShapeTensor") && shape.empty()) {
+        return KernelSignature("randint_raw",
+                               {},
+                               {"low", "high", "ShapeTensor", "seed", "dtype"},
+                               {"Out"});
+      } else {
+        return KernelSignature("randint_raw",
+                               {},
+                               {"low", "high", "shape", "seed", "dtype"},
+                               {"Out"});
+      }
+    }
+  } else {
+    if (ctx.InputSize("ShapeTensorList") > 0) {
+      return KernelSignature(
+          "randint", {}, {"low", "high", "ShapeTensorList", "dtype"}, {"Out"});
+    } else {
+      const auto& shape =
+          paddle::any_cast<std::vector<int64_t>>(ctx.Attr("shape"));
+      if (ctx.HasInput("ShapeTensor") && shape.empty()) {
+        return KernelSignature(
+            "randint", {}, {"low", "high", "ShapeTensor", "dtype"}, {"Out"});
+      } else {
+        return KernelSignature(
+            "randint", {}, {"low", "high", "shape", "dtype"}, {"Out"});
+      }
+    }
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(randint, phi::RandintOpArgumentMapping);