[Phi] Migrate save kernel (#45665)

* add save kernel

* add save_sr_kernel

* remove original save_op

* add save gpu kernel

* remove combine kernel

* add port.h include

* add save selected rows test

* remove useless kernel.h

paddle/fluid/operators/save_load_op_test.cc

@@ -15,8 +15,12 @@ limitations under the License. */
 #include "gtest/gtest.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/platform/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
 
-USE_CPU_ONLY_OP(save);
+USE_OP_ITSELF(save);
+PD_DECLARE_KERNEL(save, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(save_sr, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(cast, CPU, ALL_LAYOUT);
 USE_CPU_ONLY_OP(load);
 
 TEST(SaveLoadOp, CPU) {
@@ -63,6 +67,43 @@ TEST(SaveLoadOp, CPU) {
   }
 }
 
+TEST(SaveLoadOpSelectedRows, CPU) {
+  paddle::framework::Scope scope;
+  paddle::platform::CPUPlace place;
+
+  auto var = scope.Var("test_var_sr");
+  auto selected_rows = var->GetMutable<phi::SelectedRows>();
+  selected_rows->set_height(3);
+  selected_rows->set_rows({0, 1, 2});
+  auto* tensor = selected_rows->mutable_value();
+  tensor->Resize({3, 10});
+  int* expect = tensor->mutable_data<int>(place);
+  for (int64_t i = 0; i < tensor->numel(); ++i) {
+    expect[i] = static_cast<int>(i);
+  }
+  paddle::framework::AttributeMap attrs;
+  attrs.insert({"file_path", std::string("selected_rows.save")});
+
+  auto save_op = paddle::framework::OpRegistry::CreateOp(
+      "save", {{"X", {"test_var_sr"}}}, {}, attrs);
+  save_op->Run(scope, place);
+
+  auto load_var = scope.Var("out_var_sr");
+  auto target = load_var->GetMutable<phi::SelectedRows>();
+  auto load_op = paddle::framework::OpRegistry::CreateOp(
+      "load", {}, {{"Out", {"out_var_sr"}}}, attrs);
+  load_op->Run(scope, place);
+  const int* actual = target->value().data<int>();
+  for (int64_t i = 0; i < tensor->numel(); ++i) {
+    EXPECT_EQ(expect[i], actual[i]);
+  }
+  EXPECT_EQ(target->height(), 3);
+  auto& rows = target->rows();
+  for (size_t i = 0; i < rows.size(); ++i) {
+    EXPECT_EQ(rows[i], static_cast<int64_t>(i));
+  }
+}
+
 TEST(SaveFP16Op, CPU) {
   paddle::framework::Scope scope;
   paddle::platform::CPUPlace place;

paddle/fluid/operators/save_op.cc

@@ -87,15 +87,3 @@ REGISTER_OPERATOR(save,
                   ops::SaveOp,
                   ops::SaveOpProtoMaker,
                   ops::SaveOpVarTypeInference);
-
-REGISTER_OP_CPU_KERNEL(
-    save,
-    ops::SaveOpKernel<phi::CPUContext, float>,
-    ops::SaveOpKernel<phi::CPUContext, double>,
-    ops::SaveOpKernel<phi::CPUContext, paddle::platform::float16>,
-    ops::SaveOpKernel<phi::CPUContext, paddle::platform::bfloat16>,
-    ops::SaveOpKernel<phi::CPUContext, int>,
-    ops::SaveOpKernel<phi::CPUContext, uint8_t>,
-    ops::SaveOpKernel<phi::CPUContext, int8_t>,
-    ops::SaveOpKernel<phi::CPUContext, int16_t>,
-    ops::SaveOpKernel<phi::CPUContext, int64_t>);

paddle/phi/core/kernel_utils.h

@@ -350,8 +350,6 @@ struct KernelImpl<Return (*)(DevCtx, Args...), kernel_fn> {
     static void Compute(KernelContext* ctx, DevCtx dev_ctx, Args&... args) {
       static_assert(dev_ctx_idx > 0,
                     "Kernel should pass DeviceContext as argument.");
-      static_assert(out_idx > 0, "Kernel should have output argument.");
-      // TODO(chenweihang): check dev_ctx, in, attr, out number
       return kernel_fn(dev_ctx, args...);
     }
   };

paddle/phi/infermeta/unary.cc

@@ -337,6 +337,7 @@ void CastInferMeta(const MetaTensor& x, DataType out_dtype, MetaTensor* out) {
   out->set_dims(x.dims());
   out->set_dtype(out_dtype);
   out->set_layout(x.layout());
+  out->share_lod(x);
 }
 
 void CholeskyInferMeta(const MetaTensor& x, bool upper, MetaTensor* out) {

paddle/phi/kernels/save_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/save_kernel.h"
+
+#include <fstream>
+
+#include "paddle/phi/backends/dynload/port.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/serialization.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SaveKernel(const Context& dev_ctx,
+                const DenseTensor& x,
+                const std::string& file_path,
+                bool overwrite,
+                bool save_as_fp16) {
+  PADDLE_ENFORCE_EQ(
+      FileExists(file_path) && !overwrite,
+      false,
+      phi::errors::PreconditionNotMet(
+          "%s exists!, cannot save to it when overwrite is set to false.",
+          file_path,
+          overwrite));
+
+  MkDirRecursively(DirName(file_path).c_str());
+
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ofstream fout(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(
+      static_cast<bool>(fout),
+      true,
+      phi::errors::Unavailable("Cannot open %s to save variables.", file_path));
+
+  auto in_dtype = x.dtype();
+  auto out_dtype = save_as_fp16 ? DataType::FLOAT16 : in_dtype;
+
+  if (in_dtype != out_dtype) {
+    auto out = Cast<T>(dev_ctx, x, out_dtype);
+    SerializeToStream(fout, out, dev_ctx);
+  } else {
+    SerializeToStream(fout, x, dev_ctx);
+  }
+  fout.close();
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(save,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SaveKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(save,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SaveKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+#endif

paddle/fluid/operators/save_op.cu → paddle/phi/kernels/save_kernel.h

-/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+/* 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.
@@ -12,17 +12,19 @@ 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/fluid/operators/save_op.h"
-#include "paddle/fluid/platform/float16.h"
+#pragma once
 
-namespace ops = paddle::operators;
+#include <string>
 
-REGISTER_OP_CUDA_KERNEL(
-    save,
-    ops::SaveOpKernel<phi::GPUContext, float>,
-    ops::SaveOpKernel<phi::GPUContext, double>,
-    ops::SaveOpKernel<phi::GPUContext, int>,
-    ops::SaveOpKernel<phi::GPUContext, uint8_t>,
-    ops::SaveOpKernel<phi::GPUContext, int8_t>,
-    ops::SaveOpKernel<phi::GPUContext, int64_t>,
-    ops::SaveOpKernel<phi::GPUContext, paddle::platform::float16>);
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SaveKernel(const Context& dev_ctx,
+                const DenseTensor& x,
+                const std::string& file_path,
+                bool overwrite,
+                bool save_as_fp16);
+
+}  // namespace phi

paddle/phi/kernels/selected_rows/save_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/selected_rows/save_kernel.h"
+
+#include <fstream>
+
+#include "paddle/phi/backends/dynload/port.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/serialization.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void SaveKernel(const Context& dev_ctx,
+                const SelectedRows& x,
+                const std::string& file_path,
+                bool overwrite,
+                bool save_as_fp16) {
+  PADDLE_ENFORCE_EQ(
+      FileExists(file_path) && !overwrite,
+      false,
+      phi::errors::PreconditionNotMet(
+          "%s exists!, cannot save to it when overwrite is set to false.",
+          file_path,
+          overwrite));
+  PADDLE_ENFORCE_EQ(save_as_fp16,
+                    false,
+                    phi::errors::Unimplemented(
+                        "SelectedRows is not supported to save as float16."));
+
+  MkDirRecursively(DirName(file_path).c_str());
+
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ofstream fout(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(
+      static_cast<bool>(fout),
+      true,
+      phi::errors::Unavailable("Cannot open %s to save variables.", file_path));
+  SerializeToStream(fout, x, dev_ctx);
+  fout.close();
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(save_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::SaveKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(save_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::SaveKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+#endif

paddle/phi/kernels/selected_rows/save_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 <string>
+
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void SaveKernel(const Context& dev_ctx,
+                const SelectedRows& x,
+                const std::string& file_path,
+                bool overwrite,
+                bool save_as_fp16);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/ops/compat/save_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 SaveOpArgumentMapping(const ArgumentMappingContext& ctx) {
+  if (ctx.IsDenseTensorInput("X")) {
+    return KernelSignature(
+        "save", {"X"}, {"file_path", "overwrite", "save_as_fp16"}, {});
+  } else if (ctx.IsSelectedRowsInput("X")) {
+    return KernelSignature(
+        "save_sr", {"X"}, {"file_path", "overwrite", "save_as_fp16"}, {});
+  } else {
+    return KernelSignature("unregistered", {}, {}, {});
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(save, phi::SaveOpArgumentMapping);