chore: import upstream snapshot with attribution

test/custom_runtime/CMakeLists.txt

@@ -0,0 +1,40 @@
+if(WITH_CUSTOM_DEVICE AND NOT WITH_GPU)
+  set(PLUGIN_URL https://github.com/PaddlePaddle/PaddleCustomDevice.git)
+  set(PLUGIN_TAG develop)
+
+  file(
+    GLOB TEST_OPS
+    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+    "test_*.py")
+  string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+  list(REMOVE_ITEM TEST_OPS test_custom_cpu_profiler_plugin)
+  list(REMOVE_ITEM TEST_OPS test_custom_cpu_to_static)
+  list(REMOVE_ITEM TEST_OPS test_collective_process_group_xccl)
+
+  foreach(TEST_OP ${TEST_OPS})
+    py_test(
+      ${TEST_OP}
+      SRCS ${TEST_OP}.py ENVS FLAGS_allocator_strategy=naive_best_fit
+           PLUGIN_URL=${PLUGIN_URL} PLUGIN_TAG=${PLUGIN_TAG}
+           FLAGS_enable_pir_with_pt_in_dy2st=False)
+  endforeach()
+
+  bash_test_modules(
+    test_fleet_launch_custom_device
+    START_BASH
+    test_fleet_launch_custom_device.sh
+    ENVS
+    PYTHONPATH=""
+    FLAGS_allocator_strategy=naive_best_fit
+    PADDLE_BINARY_DIR=${PADDLE_BINARY_DIR}
+    PLUGIN_URL=${PLUGIN_URL}
+    PLUGIN_TAG=${PLUGIN_TAG}
+    PYTHON_EXECUTABLE=${PYTHON_EXECUTABLE})
+  set_tests_properties(test_fleet_launch_custom_device PROPERTIES TIMEOUT 120)
+  set_tests_properties(test_custom_cpu_plugin PROPERTIES TIMEOUT 120)
+  set_tests_properties(test_custom_op_setup PROPERTIES TIMEOUT 120)
+  # cpp testing
+  paddle_test(extension_header_test SRCS extension_header_test.cc DEPS phi
+              common)
+endif()

test/custom_runtime/custom_device_multi_process_collective.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2019 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.
+
+import os
+import sys
+
+
+def train(prefix):
+    selected_accelerators = os.getenv("FLAGS_selected_accelerators")
+    selected_custom_devices = os.getenv("FLAGS_selected_custom_cpus")
+    trainer_id = int(os.getenv("PADDLE_TRAINER_ID"))
+    worker_endpoints_env = os.getenv("PADDLE_TRAINER_ENDPOINTS")
+    current_endpoint = os.getenv("PADDLE_CURRENT_ENDPOINT")
+    worker_endpoints = worker_endpoints_env
+    trainers_num = len(worker_endpoints.split(','))
+    device_ids = os.getenv("PADDLE_WORLD_DEVICE_IDS")
+    current_device_id = os.getenv("PADDLE_LOCAL_DEVICE_IDS")
+
+    details = f"selected_accelerators:{selected_accelerators} selected_custom_devices:{selected_custom_devices} worker_endpoints:{worker_endpoints} trainers_num:{trainers_num} current_endpoint:{current_endpoint} trainer_id:{trainer_id} device_ids:{device_ids} device_id:{current_device_id}"
+
+    print(details)
+    with open(f"multi_process_{prefix}.check_{trainer_id}.log", "w") as f:
+        f.write(details)
+
+
+if __name__ == '__main__':
+    prefix = sys.argv[1]
+    train(prefix)

test/custom_runtime/custom_op.cc

@@ -0,0 +1,215 @@
+// 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 <iostream>
+#include <vector>
+
+#include "paddle/extension.h"
+#include "paddle/phi/backends/context_pool.h"
+
+#define CHECK_CPU_INPUT(x) PD_CHECK(x.is_cpu(), #x " must be a CPU Tensor.")
+#define CHECK_CUSTOM_INPUT(x) \
+  PD_CHECK(x.is_custom_device(), #x " must be a custom Tensor.")
+
+template <typename data_t>
+void relu_cpu_forward_kernel(const data_t* x_data,
+                             data_t* out_data,
+                             int64_t x_numel) {
+  PD_CHECK(x_data != nullptr, "x_data is nullptr.");
+  PD_CHECK(out_data != nullptr, "out_data is nullptr.");
+  for (int64_t i = 0; i < x_numel; ++i) {
+    out_data[i] = std::max(static_cast<data_t>(0.), x_data[i]);
+  }
+}
+
+template <typename data_t>
+void relu_cpu_backward_kernel(const data_t* grad_out_data,
+                              const data_t* out_data,
+                              data_t* grad_x_data,
+                              int64_t out_numel) {
+  for (int64_t i = 0; i < out_numel; ++i) {
+    grad_x_data[i] =
+        grad_out_data[i] * (out_data[i] > static_cast<data_t>(0) ? 1. : 0.);
+  }
+}
+
+template <typename data_t>
+void relu_cpu_double_backward_kernel(const data_t* out_data,
+                                     const data_t* ddx_data,
+                                     data_t* ddout_data,
+                                     int64_t ddout_numel) {
+  for (int64_t i = 0; i < ddout_numel; ++i) {
+    ddout_data[i] =
+        ddx_data[i] * (out_data[i] > static_cast<data_t>(0) ? 1. : 0.);
+  }
+}
+
+std::vector<paddle::Tensor> relu_cpu_forward(const paddle::Tensor& x) {
+  CHECK_CPU_INPUT(x);
+  auto out = paddle::empty_like(x);
+
+  PD_DISPATCH_FLOATING_TYPES(
+      x.type(), "relu_cpu_forward", ([&] {
+        relu_cpu_forward_kernel<data_t>(
+            x.data<data_t>(), out.data<data_t>(), x.numel());
+      }));
+
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_cpu_backward(const paddle::Tensor& x,
+                                              const paddle::Tensor& out,
+                                              const paddle::Tensor& grad_out) {
+  auto grad_x = paddle::empty_like(x);
+
+  PD_DISPATCH_FLOATING_TYPES(out.type(), "relu_cpu_backward", ([&] {
+                               relu_cpu_backward_kernel<data_t>(
+                                   grad_out.data<data_t>(),
+                                   out.data<data_t>(),
+                                   grad_x.data<data_t>(),
+                                   out.size());
+                             }));
+
+  return {grad_x};
+}
+
+std::vector<paddle::Tensor> relu_cpu_double_backward(
+    const paddle::Tensor& out, const paddle::Tensor& ddx) {
+  CHECK_CPU_INPUT(out);
+  CHECK_CPU_INPUT(ddx);
+  auto ddout = paddle::empty(out.shape(), out.dtype(), out.place());
+
+  PD_DISPATCH_FLOATING_TYPES(out.type(), "relu_cpu_double_backward", ([&] {
+                               relu_cpu_double_backward_kernel<data_t>(
+                                   out.data<data_t>(),
+                                   ddx.data<data_t>(),
+                                   ddout.mutable_data<data_t>(out.place()),
+                                   ddout.size());
+                             }));
+
+  return {ddout};
+}
+
+std::vector<paddle::Tensor> relu_custom_forward(const paddle::Tensor& x) {
+  CHECK_CUSTOM_INPUT(x);
+  auto out = paddle::relu(x);
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_custom_backward(
+    const paddle::Tensor& x,
+    const paddle::Tensor& out,
+    const paddle::Tensor& grad_out) {
+  CHECK_CUSTOM_INPUT(x);
+  CHECK_CUSTOM_INPUT(out);
+  auto grad_x = paddle::empty_like(x, x.dtype(), x.place());
+  auto ones = paddle::experimental::full_like(x, 1.0, x.dtype(), x.place());
+  auto zeros = paddle::experimental::full_like(x, 0.0, x.dtype(), x.place());
+  auto condition = paddle::experimental::greater_than(x, zeros);
+
+  grad_x = paddle::multiply(grad_out, paddle::where(condition, ones, zeros));
+
+  return {grad_x};
+}
+
+std::vector<paddle::Tensor> relu_custom_double_backward(
+    const paddle::Tensor& out, const paddle::Tensor& ddx) {
+  CHECK_CUSTOM_INPUT(out);
+  auto ddout = paddle::empty(out.shape(), out.dtype(), out.place());
+  auto ones =
+      paddle::experimental::full_like(out, 1.0, out.dtype(), out.place());
+  auto zeros =
+      paddle::experimental::full_like(out, 0.0, out.dtype(), out.place());
+  auto condition = paddle::experimental::greater_than(out, zeros);
+
+  ddout = paddle::multiply(ddx, paddle::where(condition, ones, zeros));
+
+  return {ddout};
+}
+
+std::vector<paddle::Tensor> ReluForward(const paddle::Tensor& x) {
+  if (x.is_cpu()) {
+    return relu_cpu_forward(x);
+  } else if (x.is_custom_device()) {
+    return relu_custom_forward(x);
+  } else {
+    PD_THROW("Not implemented.");
+  }
+}
+
+std::vector<paddle::Tensor> ReluBackward(const paddle::Tensor& x,
+                                         const paddle::Tensor& out,
+                                         const paddle::Tensor& grad_out) {
+  if (x.is_cpu()) {
+    return relu_cpu_backward(x, out, grad_out);
+  } else if (x.is_custom_device()) {
+    return relu_custom_backward(x, out, grad_out);
+  } else {
+    PD_THROW("Not implemented.");
+  }
+}
+
+std::vector<paddle::Tensor> ReluDoubleBackward(const paddle::Tensor& out,
+                                               const paddle::Tensor& ddx) {
+  if (out.is_cpu()) {
+    return relu_cpu_double_backward(out, ddx);
+  } else if (out.is_custom_device()) {
+    return relu_custom_double_backward(out, ddx);
+  } else {
+    PD_THROW("Not implemented.");
+  }
+}
+
+std::vector<std::vector<int64_t>> ReluDoubleBackwardInferShape(
+    const std::vector<int64_t>& out_shape,
+    const std::vector<int64_t>& ddx_shape) {
+  return {out_shape};
+}
+
+PD_BUILD_OP(custom_relu)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(ReluForward));
+
+PD_BUILD_GRAD_OP(custom_relu)
+    .Inputs({"X", "Out", paddle::Grad("Out")})
+    .Outputs({paddle::Grad("X")})
+    .SetKernelFn(PD_KERNEL(ReluBackward));
+
+PD_BUILD_DOUBLE_GRAD_OP(custom_relu)
+    .Inputs({"Out", paddle::Grad(paddle::Grad("X"))})
+    .Outputs({paddle::Grad(paddle::Grad("Out"))})
+    .SetKernelFn(PD_KERNEL(ReluDoubleBackward))
+    .SetInferShapeFn(PD_INFER_SHAPE(ReluDoubleBackwardInferShape));
+
+std::vector<paddle::Tensor> StreamForward(const paddle::Tensor& x) {
+  CHECK_CUSTOM_INPUT(x);
+
+  auto dev_ctx =
+      paddle::experimental::DeviceContextPool::Instance().Get(x.place());
+  auto custom_ctx = static_cast<const phi::CustomContext*>(dev_ctx);
+  std::shared_ptr<phi::stream::Stream> stream = custom_ctx->GetStream();
+
+  PD_CHECK(stream != nullptr);
+  std::cout << "Check stream != nullptr successfully" << std::endl;
+  custom_ctx->Wait();
+  std::cout << "Check Wait successfully" << std::endl;
+
+  return {x};
+}
+
+PD_BUILD_OP(custom_stream)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(StreamForward));

test/custom_runtime/extension_header_test.cc

@@ -0,0 +1,24 @@
+// Copyright (c) 2023 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 "glog/logging.h"
+#include "gtest/gtest.h"
+
+// Note(qili93): ensure compile with one header file 'extension.h' only,
+// !!! do not fix this ut by adding other header files (PR#60842) !!!
+#include "paddle/phi/extension.h"
+
+TEST(CustomDevice, extension_header) {
+  VLOG(1) << "check extension header support compile only";
+}

test/custom_runtime/process_group_xccl.py

@@ -0,0 +1,235 @@
+# 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.
+
+import random
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+
+def init_process_group(strategy=None):
+    nranks = paddle.distributed.ParallelEnv().nranks
+    rank = paddle.distributed.ParallelEnv().local_rank
+    is_master = True if rank == 0 else False
+    store = paddle.base.core.TCPStore("127.0.0.1", 6173, is_master, nranks)
+    pg_group = core.ProcessGroupCustom.create(
+        store,
+        paddle.distributed.ParallelEnv().device_type,
+        rank,
+        nranks,
+    )
+
+    return pg_group
+
+
+class TestProcessGroupFp32(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2022)
+        random.seed(2022)
+        np.random.seed(2022)
+        self.config()
+
+    def config(self):
+        self.dtype = "float32"
+        self.shape = (2, 10, 5)
+
+    def test_create_process_group_xccl(self):
+        device_id = paddle.distributed.ParallelEnv().dev_id
+        paddle.set_device(f'custom_cpu:{device_id}')
+
+        pg = init_process_group()
+
+        x = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_y = paddle.to_tensor(y)
+
+        sum_result = tensor_x + tensor_y
+        if pg.rank() == 0:
+            task = pg.all_reduce(tensor_x, core.ReduceOp.SUM, sync_op=True)
+            task.wait()
+            # assert np.array_equal(tensor_x, sum_result)
+        else:
+            task = pg.all_reduce(tensor_y, core.ReduceOp.SUM, sync_op=True)
+            task.wait()
+            # assert np.array_equal(tensor_y, sum_result)
+
+        print("test allreduce sum api ok", flush=True)
+
+        x = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_y = paddle.to_tensor(y)
+
+        max_result = paddle.maximum(tensor_x, tensor_y)
+
+        if pg.rank() == 0:
+            task = pg.all_reduce(tensor_x, core.ReduceOp.MAX, sync_op=True)
+            task.wait()
+            # assert np.array_equal(tensor_x, max_result)
+        else:
+            task = pg.all_reduce(tensor_y, core.ReduceOp.MAX, sync_op=True)
+            task.wait()
+            # assert np.array_equal(tensor_y, max_result)
+
+        print("test allreduce max api ok", flush=True)
+
+        # test broadcast
+        # rank 0
+        x = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        # rank 1
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_y = paddle.to_tensor(y)
+
+        broadcast_result = paddle.assign(tensor_x)
+        if pg.rank() == 0:
+            task = pg.broadcast(tensor_x, 0, sync_op=True)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+            assert task.is_completed()
+            # assert np.array_equal(broadcast_result, tensor_x)
+        else:
+            task = pg.broadcast(tensor_y, 0, sync_op=True)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+            assert task.is_completed()
+            # assert np.array_equal(broadcast_result, tensor_y)
+
+        print("test broadcast api ok", flush=True)
+
+        # test barrier
+        # rank 0
+        if pg.rank() == 0:
+            task = pg.barrier(device_id)
+            task.wait()
+        # rank 1
+        else:
+            task = pg.barrier(device_id)
+            task.wait()
+
+        print("test barrier api ok\n", flush=True)
+        return
+
+        # test allgather
+        # rank 0
+        x = np.random.random(self.shape).astype(self.dtype)
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        tensor_y = paddle.to_tensor(y)
+        out_shape = list(self.shape)
+        out_shape[0] *= 2
+        out = np.random.random(out_shape).astype(self.dtype)
+        tensor_out = paddle.to_tensor(out)
+        if pg.rank() == 0:
+            task = pg.all_gather(tensor_out, tensor_x, sync_op=True)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        # rank 1
+        else:
+            task = pg.all_gather(tensor_out, tensor_y, sync_op=True)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        out_1 = paddle.slice(tensor_out, [0], [0], [out_shape[0] // 2])
+        out_2 = paddle.slice(
+            tensor_out, [0], [out_shape[0] // 2], [out_shape[0]]
+        )
+        # assert np.array_equal(tensor_x, out_1)
+        # assert np.array_equal(tensor_y, out_2)
+        print("test allgather api ok\n", flush=True)
+
+        # test alltoall
+        # rank 0
+        x = np.random.random(self.shape).astype(self.dtype)
+        y = np.random.random(self.shape).astype(self.dtype)
+        out1 = np.random.random(self.shape).astype(self.dtype)
+        out2 = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        tensor_y = paddle.to_tensor(y)
+        tensor_out1 = paddle.to_tensor(out1)
+        tensor_out2 = paddle.to_tensor(out2)
+        raw_tensor_x_2 = paddle.slice(
+            tensor_x, [0], [self.shape[0] // 2], [self.shape[0]]
+        )
+        raw_tensor_y_1 = paddle.slice(tensor_y, [0], [0], [self.shape[0] // 2])
+        if pg.rank() == 0:
+            task = pg.alltoall(tensor_out1, tensor_x)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        # rank 1
+        else:
+            task = pg.alltoall(tensor_out2, tensor_y)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        out1_2 = paddle.slice(
+            tensor_out1, [0], [self.shape[0] // 2], [self.shape[0]]
+        )
+        out2_1 = paddle.slice(tensor_out2, [0], [0], [self.shape[0] // 2])
+        # if pg.rank() == 0:
+        #     assert np.array_equal(out1_2.numpy(), raw_tensor_y_1.numpy())
+        # else:
+        #     assert np.array_equal(out2_1, raw_tensor_x_2)
+        print("test alltoall api ok\n", flush=True)
+
+        # test Reduce
+        # rank 0
+        x = np.random.random(self.shape).astype(self.dtype)
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        tensor_y = paddle.to_tensor(y)
+        sum_result = tensor_x + tensor_y
+        if pg.rank() == 0:
+            task = pg.reduce(tensor_x, 0)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        # rank 1
+        else:
+            task = pg.reduce(tensor_y, 0)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        # if pg.rank() == 0:
+        #     assert np.array_equal(tensor_x, sum_result)
+        print("test reduce sum api ok\n", flush=True)
+
+        # test Scatter
+        # rank 0
+        in_shape = list(self.shape)
+        in_shape[0] *= 2
+        x = np.random.random(in_shape).astype(self.dtype)
+        y = np.random.random(self.shape).astype(self.dtype)
+        tensor_x = paddle.to_tensor(x)
+        tensor_y = paddle.to_tensor(y)
+        if pg.rank() == 0:
+            task = pg.scatter(tensor_x, tensor_y, 0)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        # rank 1
+        else:
+            task = pg.scatter(tensor_x, tensor_y, 0)
+            task.wait()
+            # paddle.base.core._custom_device_synchronize("custom_cpu", -1)
+        out1 = paddle.slice(tensor_x, [0], [0], [self.shape[0]])
+        out2 = paddle.slice(tensor_x, [0], [self.shape[0]], [self.shape[0] * 2])
+        # if pg.rank() == 0:
+        #     assert np.array_equal(tensor_y, out1)
+        # else:
+        #     assert np.array_equal(tensor_y, out2)
+        print("test scatter api ok\n", flush=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/custom_runtime/test_collective_process_group_xccl.py

@@ -0,0 +1,192 @@
+# 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.
+
+import copy
+import os
+import subprocess
+import sys
+import tempfile
+import time
+import unittest
+
+
+def start_local_trainers(
+    cluster,
+    pod,
+    training_script,
+    training_script_args,
+    eager_mode=True,
+    log_dir=None,
+):
+    from paddle.distributed.utils.launch_utils import (  # noqa: F401
+        TrainerProc,
+        find_free_ports,
+        get_cluster,
+        watch_local_trainers,
+    )
+
+    current_env = copy.copy(os.environ.copy())
+    # paddle broadcast ncclUniqueId use socket, and
+    # proxy maybe make trainers unreachable, so delete them.
+    # if we set them to "", grpc will log error message "bad uri"
+    # so just delete them.
+    current_env.pop("http_proxy", None)
+    current_env.pop("https_proxy", None)
+
+    procs = []
+
+    os.system("rm -rf log && mkdir -p log")
+    for idx, t in enumerate(pod.trainers):
+        proc_env = {
+            "FLAGS_selected_custom_cpus": "{}".format(
+                ",".join([str(g) for g in t.gpus])
+            ),
+            "PADDLE_TRAINER_ID": str(t.rank),
+            "PADDLE_CURRENT_ENDPOINT": str(t.endpoint),
+            "PADDLE_TRAINERS_NUM": str(cluster.trainers_nranks()),
+            "PADDLE_TRAINER_ENDPOINTS": ",".join(cluster.trainers_endpoints()),
+            "PADDLE_DISTRI_CUSTOM_DEVICE_TYPE": "custom_cpu",
+        }
+
+        current_env.update(proc_env)
+
+        print(f"trainer proc env:{current_env}")
+
+        if os.getenv('WITH_COVERAGE', 'OFF') == 'ON':
+            cmd = "python -m coverage run --branch -p " + training_script
+        else:
+            cmd = "python -u " + training_script
+
+        print(f"start trainer proc:{cmd} env:{proc_env}")
+
+        fn = open(f"workerlog.{idx}", "a")
+        proc = subprocess.Popen(
+            cmd.split(" "), env=current_env, stdout=fn, stderr=fn
+        )
+
+        tp = TrainerProc()
+        tp.proc = proc
+        tp.rank = t.rank
+        tp.log_fn = fn
+        tp.cmd = cmd
+
+        procs.append(tp)
+
+    return procs
+
+
+def get_cluster_from_args(selected_gpus):
+    from paddle.distributed.utils.launch_utils import (  # noqa: F401
+        TrainerProc,
+        find_free_ports,
+        get_cluster,
+        watch_local_trainers,
+    )
+
+    cluster_node_ips = '127.0.0.1'
+    node_ip = '127.0.0.1'
+
+    node_ips = [x.strip() for x in cluster_node_ips.split(',')]
+
+    node_ips.index(node_ip)
+
+    free_ports = None
+
+    free_ports = find_free_ports(len(selected_gpus))
+    if free_ports is not None:
+        free_ports = list(free_ports)
+
+    trainer_endpoints = []
+    for ip in node_ips:
+        trainer_endpoints.append([f"{ip}:{port}" for port in free_ports])
+    return get_cluster(node_ips, node_ip, trainer_endpoints, selected_gpus)
+
+
+class TestMultipleCustomCPU(unittest.TestCase):
+    def run_mnist_2custom_cpu(self, target_file_name, eager_mode=True):
+        from paddle.distributed.utils.launch_utils import (  # noqa: F401
+            TrainerProc,
+            find_free_ports,
+            get_cluster,
+            watch_local_trainers,
+        )
+
+        selected_devices = [0, 1]
+        cluster = None
+        pod = None
+
+        cluster, pod = get_cluster_from_args(selected_devices)
+
+        procs = start_local_trainers(
+            cluster,
+            pod,
+            eager_mode=eager_mode,
+            training_script=target_file_name,
+            training_script_args=[],
+        )
+
+        while True:
+            alive = watch_local_trainers(procs, cluster.trainers_endpoints())
+
+            if not alive:
+                print(f"Local procs complete, POD info:{pod}")
+                break
+            time.sleep(3)
+
+
+class TestProcessGroup(TestMultipleCustomCPU):
+    def setUp(self):
+        # compile so and set to current path
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        self.temp_dir = tempfile.TemporaryDirectory()
+        cmd = 'cd {} \
+            && git clone --depth 1 {} \
+            && cd PaddleCustomDevice \
+            && git fetch origin \
+            && git checkout {} -b dev \
+            && cd backends/custom_cpu \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8'.format(
+            self.temp_dir.name,
+            os.getenv('PLUGIN_URL'),
+            os.getenv('PLUGIN_TAG'),
+            sys.executable,
+        )
+        os.system(cmd)
+
+        # set environment for loading and registering compiled custom kernels
+        # only valid in current process
+        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
+            cur_dir,
+            f'{self.temp_dir.name}/PaddleCustomDevice/backends/custom_cpu/build',
+        )
+        os.environ['FLAGS_selected_custom_cpus'] = '0,1'
+        os.environ['CUSTOM_CPU_VISIBLE_DEVICES'] = '0,1'
+        os.environ['PADDLE_XCCL_BACKEND'] = 'custom_cpu'
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_process_group_xccl(self):
+        from paddle.distributed.utils.launch_utils import (  # noqa: F401
+            TrainerProc,
+            find_free_ports,
+            get_cluster,
+            watch_local_trainers,
+        )
+
+        self.run_mnist_2custom_cpu('process_group_xccl.py')
+
+
+if __name__ == "__main__":
+    unittest.main()

test/custom_runtime/test_custom_cpu_plugin.py

@@ -0,0 +1,336 @@
+# 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.
+
+import os
+import sys
+import tempfile
+import unittest
+
+import numpy as np
+
+
+class TestCustomCPUPlugin(unittest.TestCase):
+    def setUp(self):
+        # compile so and set to current path
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        self.temp_dir = tempfile.TemporaryDirectory()
+        cmd = 'cd {} \
+            && git clone --depth 1 {} \
+            && cd PaddleCustomDevice \
+            && git fetch origin \
+            && git checkout {} -b dev \
+            && cd backends/custom_cpu \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8'.format(
+            self.temp_dir.name,
+            os.getenv('PLUGIN_URL'),
+            os.getenv('PLUGIN_TAG'),
+            sys.executable,
+        )
+        os.system(cmd)
+
+        # set environment for loading and registering compiled custom kernels
+        # only valid in current process
+        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
+            cur_dir,
+            f'{self.temp_dir.name}/PaddleCustomDevice/backends/custom_cpu/build',
+        )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+        del os.environ['CUSTOM_DEVICE_ROOT']
+
+    def test_custom_device(self):
+        self._test_custom_device_dataloader()
+        self._test_custom_device_mnist()
+        self._test_eager_backward_api()
+        self._test_eager_copy_to()
+        self._test_fallback_kernel()
+        self._test_scalar()
+        self._test_custom_device_py_api()
+        self._test_custom_device_mix_precision()
+
+    def _test_custom_device_dataloader(self):
+        import paddle
+
+        paddle.set_device('custom_cpu')
+        dataset = paddle.vision.datasets.MNIST(
+            mode='test',
+            transform=paddle.vision.transforms.Compose(
+                [
+                    paddle.vision.transforms.CenterCrop(20),
+                    paddle.vision.transforms.RandomResizedCrop(14),
+                    paddle.vision.transforms.Normalize(),
+                    paddle.vision.transforms.ToTensor(),
+                ]
+            ),
+        )
+        loader = paddle.io.DataLoader(
+            dataset, batch_size=32, num_workers=1, shuffle=True
+        )
+        for image, label in loader:
+            self.assertTrue(image.place.is_custom_place())
+            self.assertTrue(label.place.is_custom_place())
+            break
+
+    def _test_custom_device_mnist(self):
+        import paddle
+
+        class MNIST(paddle.nn.Layer):
+            def __init__(self):
+                super().__init__()
+                self.shape = 1 * 28 * 28
+                self.size = 10
+                self.output_weight = self.create_parameter(
+                    [self.shape, self.size]
+                )
+                self.accuracy = paddle.metric.Accuracy()
+
+            def forward(self, inputs, label=None):
+                x = paddle.reshape(inputs, shape=[-1, self.shape])
+                x = paddle.matmul(x, self.output_weight)
+                x = paddle.nn.functional.softmax(x)
+                if label is not None:
+                    self.accuracy.reset()
+                    correct = self.accuracy.compute(x, label)
+                    self.accuracy.update(correct)
+                    acc = self.accuracy.accumulate()
+                    return x, acc
+                else:
+                    return x
+
+        paddle.set_device('custom_cpu')
+        dataset = paddle.vision.datasets.MNIST(
+            mode='train',
+            transform=paddle.vision.transforms.Compose(
+                [paddle.vision.transforms.ToTensor()]
+            ),
+        )
+        loader = paddle.io.DataLoader(
+            dataset, batch_size=64, num_workers=1, shuffle=True
+        )
+
+        mnist = MNIST()
+        sgd = paddle.optimizer.SGD(
+            learning_rate=0.01, parameters=mnist.parameters()
+        )
+
+        data = next(loader())
+        img = data[0]
+        label = data[1]
+        label_int32 = paddle.cast(label, 'int32')
+
+        pred, acc = mnist(img, label_int32)
+        avg_loss = paddle.nn.functional.cross_entropy(pred, label_int32)
+        avg_loss.backward()
+        sgd.step()
+        sgd.clear_grad()
+
+        self.assertTrue(pred.place.is_custom_place())
+
+    def _test_eager_backward_api(self):
+        x = np.random.random([2, 2]).astype("float32")
+        y = np.random.random([2, 2]).astype("float32")
+        grad = np.ones([2, 2]).astype("float32")
+
+        import paddle
+
+        paddle.set_device('custom_cpu')
+        paddle.device.get_available_device()
+        x_tensor = paddle.to_tensor(x, stop_gradient=False)
+        y_tensor = paddle.to_tensor(y)
+        z1_tensor = paddle.matmul(x_tensor, y_tensor)
+        z2_tensor = paddle.matmul(x_tensor, y_tensor)
+
+        grad_tensor = paddle.to_tensor(grad)
+        paddle.autograd.backward([z1_tensor, z2_tensor], [grad_tensor, None])
+
+        self.assertTrue(x_tensor.grad.place.is_custom_place())
+
+    def _test_eager_copy_to(self):
+        import paddle
+
+        x = np.random.random([2, 2]).astype("float32")
+        # cpu -> custom
+        cpu_tensor = paddle.to_tensor(
+            x, dtype='float32', place=paddle.CPUPlace()
+        )
+        custom_cpu_tensor = cpu_tensor._copy_to(
+            paddle.CustomPlace('custom_cpu', 0), True
+        )
+        np.testing.assert_array_equal(custom_cpu_tensor, x)
+        self.assertTrue(custom_cpu_tensor.place.is_custom_place())
+        # custom -> custom
+        another_custom_cpu_tensor = custom_cpu_tensor._copy_to(
+            paddle.CustomPlace('custom_cpu', 0), True
+        )
+        np.testing.assert_array_equal(another_custom_cpu_tensor, x)
+        self.assertTrue(another_custom_cpu_tensor.place.is_custom_place())
+        # custom -> cpu
+        another_cpu_tensor = custom_cpu_tensor._copy_to(paddle.CPUPlace(), True)
+        np.testing.assert_array_equal(another_cpu_tensor, x)
+        self.assertTrue(another_cpu_tensor.place.is_cpu_place())
+        # custom -> custom self
+        another_custom_cpu_tensor = another_custom_cpu_tensor._copy_to(
+            paddle.CustomPlace('custom_cpu', 0), True
+        )
+        np.testing.assert_array_equal(another_custom_cpu_tensor, x)
+        self.assertTrue(another_custom_cpu_tensor.place.is_custom_place())
+
+    def _test_fallback_kernel(self):
+        # using (custom_cpu, add, int16) which is not registered
+        import paddle
+
+        r = np.array([6, 6, 6], 'int16')
+        x = paddle.to_tensor([5, 4, 3], 'int16')
+        y = paddle.to_tensor([1, 2, 3], 'int16')
+        z = paddle.add(x, y)
+        np.testing.assert_array_equal(z, r)
+
+    def _test_scalar(self):
+        import paddle
+
+        data_1 = paddle.to_tensor(
+            [[[[1.0, 4.0, 5.0, 7.0], [3.0, 4.0, 5.0, 6.0]]]]
+        )
+        k_t = paddle.to_tensor([3], dtype="int32")
+        value_1, indices_1 = paddle.topk(data_1, k=k_t)
+
+    def _test_custom_device_gradient_accumulation(self):
+        import paddle
+
+        class MNIST(paddle.nn.Layer):
+            def __init__(self):
+                super().__init__()
+                self.shape = 1 * 28 * 28
+                self.size = 10
+                self.output_weight = self.create_parameter(
+                    [self.shape, self.size]
+                )
+                self.accuracy = paddle.metric.Accuracy()
+
+            def forward(self, inputs, label=None):
+                x = paddle.reshape(inputs, shape=[-1, self.shape])
+                x = paddle.matmul(x, self.output_weight)
+                x = paddle.nn.functional.softmax(x)
+                if label is not None:
+                    self.accuracy.reset()
+                    correct = self.accuracy.compute(x, label)
+                    self.accuracy.update(correct)
+                    acc = self.accuracy.accumulate()
+                    return x, acc
+                else:
+                    return x
+
+        paddle.set_device('custom_cpu')
+        dataset = paddle.vision.datasets.MNIST(
+            mode='train',
+            transform=paddle.vision.transforms.Compose(
+                [paddle.vision.transforms.ToTensor()]
+            ),
+        )
+        loader = paddle.io.DataLoader(
+            dataset, batch_size=64, num_workers=1, shuffle=True
+        )
+
+        mnist = MNIST()
+        sgd = paddle.optimizer.SGD(
+            learning_rate=0.01, parameters=mnist.parameters()
+        )
+
+        data = next(loader())
+        img = data[0]
+        label = data[1]
+        label_int32 = paddle.cast(label, 'int32')
+
+        pred, acc = mnist(img, label_int32)
+        avg_loss = paddle.nn.functional.cross_entropy(pred, label_int32)
+        avg_loss.backward(retain_graph=True)
+        avg_loss = paddle.nn.functional.cross_entropy(pred, label_int32)
+        avg_loss.backward()
+        sgd.step()
+
+    def _test_custom_device_mix_precision(self):
+        import tempfile
+
+        import paddle
+        from paddle.inference import (
+            PlaceType,
+            PrecisionType,
+            convert_to_mixed_precision,
+        )
+        from paddle.jit import to_static
+        from paddle.static import InputSpec
+        from paddle.vision.models import resnet50
+
+        self.temp_dir = tempfile.TemporaryDirectory()
+        model = resnet50(True)
+        net = to_static(
+            model,
+            input_spec=[InputSpec(shape=[None, 3, 224, 224], name='x')],
+            full_graph=True,
+        )
+        paddle.jit.save(
+            net, os.path.join(self.temp_dir.name, 'resnet50/inference')
+        )
+        if paddle.framework.use_pir_api():
+            return
+
+        convert_to_mixed_precision(
+            os.path.join(self.temp_dir.name, 'resnet50/inference.pdmodel'),
+            os.path.join(self.temp_dir.name, 'resnet50/inference.pdiparams'),
+            os.path.join(
+                self.temp_dir.name, 'mixed_precision/inference.pdmodel'
+            ),
+            os.path.join(
+                self.temp_dir.name, 'mixed_precision/inference.pdiparams'
+            ),
+            backend=PlaceType.CUSTOM,
+            mixed_precision=PrecisionType.Half,
+        )
+        self.temp_dir.cleanup()
+
+    def _test_custom_device_py_api(self):
+        import paddle
+
+        p = paddle.set_device('custom_cpu')
+        paddle.device.synchronize('custom_cpu')
+
+        s1 = paddle.device.Stream()
+        s2 = paddle.device.Stream(p)
+
+        s1 = paddle.device.current_stream()
+        s2 = paddle.device.current_stream(p)
+
+        e1 = paddle.device.Event()
+        e2 = paddle.device.Event(p)
+
+        s = paddle.device.Stream()
+        e = paddle.device.Event()
+        s.query()
+        s.synchronize()
+        s.wait_event(e)
+        s.record_event(e)
+        s.wait_stream(s)
+        paddle.device.set_stream(s)
+
+        e.query()
+        e.synchronize()
+        e.record(s)
+
+
+if __name__ == '__main__':
+    if os.name == 'nt' or sys.platform.startswith('darwin'):
+        # only support Linux now
+        sys.exit()
+    unittest.main()

test/custom_runtime/test_custom_cpu_profiler_plugin.py

@@ -0,0 +1,76 @@
+# 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.
+
+import os
+import sys
+import tempfile
+import unittest
+
+
+class TestCustomCPUProfilerPlugin(unittest.TestCase):
+    def setUp(self):
+        # compile so and set to current path
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        self.temp_dir = tempfile.TemporaryDirectory()
+        cmd = 'cd {} \
+            && git clone --depth 1 {} \
+            && cd PaddleCustomDevice \
+            && git fetch origin \
+            && git checkout {} -b dev \
+            && cd backends/custom_cpu \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8'.format(
+            self.temp_dir.name,
+            os.getenv('PLUGIN_URL'),
+            os.getenv('PLUGIN_TAG'),
+            sys.executable,
+        )
+        os.system(cmd)
+
+        # set environment for loading and registering compiled custom kernels
+        # only valid in current process
+        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
+            cur_dir,
+            f'{self.temp_dir.name}/PaddleCustomDevice/backends/custom_cpu/build',
+        )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+        del os.environ['CUSTOM_DEVICE_ROOT']
+
+    def test_custom_profiler(self):
+        import paddle
+        from paddle import profiler
+
+        paddle.set_device('custom_cpu')
+
+        x = paddle.to_tensor([1, 2, 3])
+        p = profiler.Profiler(
+            targets=[
+                profiler.ProfilerTarget.CPU,
+                profiler.ProfilerTarget.CUSTOM_DEVICE,
+            ]
+        )
+        p.start()
+        for iter in range(10):
+            x = x + 1
+            p.step()
+        p.stop()
+        p.summary()
+
+
+if __name__ == '__main__':
+    if os.name == 'nt' or sys.platform.startswith('darwin'):
+        # only support Linux now
+        sys.exit()
+    unittest.main()

test/custom_runtime/test_custom_cpu_to_static.py

@@ -0,0 +1,281 @@
+#   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.
+
+import os
+import sys
+import tempfile
+import time
+import unittest
+
+import numpy as np
+
+EPOCH_NUM = 1
+BATCH_SIZE = 1024
+
+
+def train_func_base(epoch_id, train_loader, model, cost, optimizer):
+    total_step = len(train_loader)
+    epoch_start = time.time()
+    for batch_id, (images, labels) in enumerate(train_loader()):
+        # forward
+        outputs = model(images)
+        loss = cost(outputs, labels)
+        # backward and optimize
+        loss.backward()
+        optimizer.step()
+        optimizer.clear_grad()
+        print(
+            f"Epoch [{epoch_id + 1}/{EPOCH_NUM}], Step [{batch_id + 1}/{total_step}], Loss: {loss.numpy()}"
+        )
+    epoch_end = time.time()
+    print(
+        f"Epoch ID: {epoch_id + 1}, FP32 train epoch time: {(epoch_end - epoch_start) * 1000} ms"
+    )
+
+
+def train_func_ampo1(epoch_id, train_loader, model, cost, optimizer, scaler):
+    import paddle
+
+    total_step = len(train_loader)
+    epoch_start = time.time()
+    for batch_id, (images, labels) in enumerate(train_loader()):
+        # forward
+        with paddle.amp.auto_cast(
+            custom_black_list={
+                "flatten_contiguous_range",
+                "greater_than",
+                "matmul_v2",
+            },
+            level='O1',
+        ):
+            outputs = model(images)
+            loss = cost(outputs, labels)
+        # backward and optimize
+        scaled = scaler.scale(loss)
+        scaled.backward()
+        scaler.minimize(optimizer, scaled)
+        optimizer.clear_grad()
+        print(
+            f"Epoch [{epoch_id + 1}/{EPOCH_NUM}], Step [{batch_id + 1}/{total_step}], Loss: {loss.numpy()}"
+        )
+    epoch_end = time.time()
+    print(
+        f"Epoch ID: {epoch_id + 1}, AMPO1 train epoch time: {(epoch_end - epoch_start) * 1000} ms"
+    )
+
+
+def test_func(epoch_id, test_loader, model, cost):
+    import paddle
+
+    # evaluation every epoch finish
+    model.eval()
+    avg_acc = [[], []]
+    for batch_id, (images, labels) in enumerate(test_loader()):
+        # forward
+        outputs = model(images)
+        loss = cost(outputs, labels)
+        # accuracy
+        acc_top1 = paddle.metric.accuracy(input=outputs, label=labels, k=1)
+        acc_top5 = paddle.metric.accuracy(input=outputs, label=labels, k=5)
+        avg_acc[0].append(acc_top1.numpy())
+        avg_acc[1].append(acc_top5.numpy())
+    model.train()
+    print(
+        f"Epoch ID: {epoch_id + 1}, Top1 accuracy: {np.array(avg_acc[0]).mean()}, Top5 accuracy: {np.array(avg_acc[1]).mean()}"
+    )
+
+
+class TestCustomCPUPlugin(unittest.TestCase):
+    def setUp(self):
+        # compile so and set to current path
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        self.temp_dir = tempfile.TemporaryDirectory()
+        cmd = 'cd {} \
+            && git clone --depth 1 {} \
+            && cd PaddleCustomDevice \
+            && git fetch origin \
+            && git checkout {} -b dev \
+            && cd backends/custom_cpu \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8'.format(
+            self.temp_dir.name,
+            os.getenv('PLUGIN_URL'),
+            os.getenv('PLUGIN_TAG'),
+            sys.executable,
+        )
+        os.system(cmd)
+
+        # set environment for loading and registering compiled custom kernels
+        # only valid in current process
+        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
+            cur_dir,
+            f'{self.temp_dir.name}/PaddleCustomDevice/backends/custom_cpu/build',
+        )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_custom_cpu_plugin(self):
+        self._test_to_static()
+        self._test_amp_o1()
+
+    def _test_to_static(self):
+        import paddle
+
+        class LeNet5(paddle.nn.Layer):
+            def __init__(self):
+                super().__init__()
+                self.fc = paddle.nn.Linear(in_features=1024, out_features=10)
+                self.relu = paddle.nn.ReLU()
+                self.fc1 = paddle.nn.Linear(in_features=10, out_features=10)
+
+            def forward(self, x):
+                out = paddle.flatten(x, 1)
+                out = self.fc(out)
+                out = self.relu(out)
+                out = self.fc1(out)
+                return out
+
+        # set device
+        paddle.set_device('custom_cpu')
+
+        # model
+        model = LeNet5()
+
+        # cost and optimizer
+        cost = paddle.nn.CrossEntropyLoss()
+        optimizer = paddle.optimizer.Adam(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+
+        # convert to static model
+        build_strategy = paddle.static.BuildStrategy()
+        mnist = paddle.jit.to_static(
+            model, build_strategy=build_strategy, full_graph=True
+        )
+
+        # data loader
+        transform = paddle.vision.transforms.Compose(
+            [
+                paddle.vision.transforms.Resize((32, 32)),
+                paddle.vision.transforms.ToTensor(),
+                paddle.vision.transforms.Normalize(
+                    mean=(0.1307,), std=(0.3081,)
+                ),
+            ]
+        )
+        train_dataset = paddle.vision.datasets.MNIST(
+            mode='train', transform=transform, download=True
+        )
+        test_dataset = paddle.vision.datasets.MNIST(
+            mode='test', transform=transform, download=True
+        )
+        train_loader = paddle.io.DataLoader(
+            train_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            drop_last=True,
+            num_workers=2,
+        )
+        test_loader = paddle.io.DataLoader(
+            test_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            drop_last=True,
+            num_workers=2,
+        )
+
+        # train and eval
+        for epoch_id in range(EPOCH_NUM):
+            train_func_base(epoch_id, train_loader, model, cost, optimizer)
+            test_func(epoch_id, test_loader, model, cost)
+
+    def _test_amp_o1(self):
+        import paddle
+
+        class LeNet5(paddle.nn.Layer):
+            def __init__(self):
+                super().__init__()
+                self.fc = paddle.nn.Linear(in_features=1024, out_features=10)
+                self.relu = paddle.nn.ReLU()
+                self.fc1 = paddle.nn.Linear(in_features=10, out_features=10)
+
+            def forward(self, x):
+                out = paddle.flatten(x, 1)
+                out = self.fc(out)
+                out = self.relu(out)
+                out = self.fc1(out)
+                return out
+
+        # set device
+        paddle.set_device('custom_cpu')
+
+        # model
+        model = LeNet5()
+
+        # cost and optimizer
+        cost = paddle.nn.CrossEntropyLoss()
+        optimizer = paddle.optimizer.Adam(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+
+        # convert to static model
+        scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
+        model, optimizer = paddle.amp.decorate(
+            models=model, optimizers=optimizer, level='O1'
+        )
+
+        # data loader
+        transform = paddle.vision.transforms.Compose(
+            [
+                paddle.vision.transforms.Resize((32, 32)),
+                paddle.vision.transforms.ToTensor(),
+                paddle.vision.transforms.Normalize(
+                    mean=(0.1307,), std=(0.3081,)
+                ),
+            ]
+        )
+        train_dataset = paddle.vision.datasets.MNIST(
+            mode='train', transform=transform, download=True
+        )
+        test_dataset = paddle.vision.datasets.MNIST(
+            mode='test', transform=transform, download=True
+        )
+        train_loader = paddle.io.DataLoader(
+            train_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            drop_last=True,
+            num_workers=2,
+        )
+        test_loader = paddle.io.DataLoader(
+            test_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            drop_last=True,
+            num_workers=2,
+        )
+
+        # train and eval
+        for epoch_id in range(EPOCH_NUM):
+            train_func_ampo1(
+                epoch_id, train_loader, model, cost, optimizer, scaler
+            )
+            test_func(epoch_id, test_loader, model, cost)
+
+
+if __name__ == '__main__':
+    if os.name == 'nt' or sys.platform.startswith('darwin'):
+        # only support Linux now
+        sys.exit()
+    unittest.main()

test/custom_runtime/test_custom_op_setup.py

@@ -0,0 +1,280 @@
+# 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.
+
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from site import getsitepackages
+
+import numpy as np
+
+from paddle.utils.cpp_extension.extension_utils import (
+    _get_all_paddle_includes_from_include_root,
+)
+
+
+def custom_relu_dynamic(func, device, dtype, np_x, use_func=True):
+    import paddle
+
+    paddle.set_device(device)
+
+    t = paddle.to_tensor(np_x, dtype=dtype)
+    t.stop_gradient = False
+    t.retain_grads()
+    sys.stdout.flush()
+
+    out = func(t) if use_func else paddle.nn.functional.relu(t)
+    out.stop_gradient = False
+
+    out.backward()
+
+    if t.grad is None:
+        return out.numpy(), t.grad
+    else:
+        return out.numpy(), t.grad.numpy()
+
+
+def custom_relu_static(func, device, dtype, np_x, use_func=True):
+    import paddle
+    from paddle import static
+
+    paddle.enable_static()
+    paddle.set_device(device)
+
+    with (
+        static.scope_guard(static.Scope()),
+        static.program_guard(static.Program()),
+    ):
+        x = static.data(name="X", shape=[None, 8], dtype=dtype)
+        x.stop_gradient = False
+        out = func(x) if use_func else paddle.nn.functional.relu(x)
+        static.append_backward(out)
+
+        exe = static.Executor()
+        exe.run(static.default_startup_program())
+        # in static mode, x data has been covered by out
+        out_v = exe.run(
+            static.default_main_program(),
+            feed={"X": np_x},
+            fetch_list=[out],
+        )
+
+    paddle.disable_static()
+    return out_v
+
+
+def custom_relu_double_grad_dynamic(func, device, dtype, np_x, use_func=True):
+    import paddle
+
+    paddle.set_device(device)
+
+    t = paddle.to_tensor(np_x, dtype=dtype, stop_gradient=False)
+    t.retain_grads()
+
+    out = func(t) if use_func else paddle.nn.functional.relu(t)
+    out.retain_grads()
+    dx = paddle.grad(
+        outputs=out,
+        inputs=t,
+        grad_outputs=paddle.ones_like(t),
+        create_graph=True,
+        retain_graph=True,
+    )
+
+    ddout = paddle.grad(
+        outputs=dx[0],
+        inputs=out.grad,
+        grad_outputs=paddle.ones_like(t),
+        create_graph=False,
+    )
+
+    assert ddout[0].numpy() is not None
+    return dx[0].numpy(), ddout[0].numpy()
+
+
+class TestNewCustomOpSetUpInstall(unittest.TestCase):
+    def setUp(self):
+        # compile so and set to current path
+        self.cur_dir = os.path.dirname(os.path.abspath(__file__))
+        self.temp_dir = tempfile.TemporaryDirectory()
+        cmd = 'cd {} \
+            && git clone --depth 1 {} \
+            && cd PaddleCustomDevice \
+            && git fetch origin \
+            && git checkout {} -b dev \
+            && cd backends/custom_cpu \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8 \
+            && cd {}'.format(
+            self.temp_dir.name,
+            os.getenv('PLUGIN_URL'),
+            os.getenv('PLUGIN_TAG'),
+            sys.executable,
+            self.cur_dir,
+        )
+        os.system(cmd)
+
+        # set environment for loading and registering compiled custom kernels
+        # only valid in current process
+        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
+            self.cur_dir,
+            f'{self.temp_dir.name}/PaddleCustomDevice/backends/custom_cpu/build',
+        )
+
+        # `import paddle` loads custom_cpu.so, hence we must import paddle after finishing build PaddleCustomDevice
+        import paddle
+
+        # [Why specific paddle_includes directory?]
+        # Add paddle_includes to pass CI, for more details,
+        # please refer to the comments in `paddle/tests/custom_op/utils.py``
+        paddle_includes = []
+        for site_packages_path in getsitepackages():
+            paddle_include_dir = Path(site_packages_path) / "paddle/include"
+            paddle_includes.extend(
+                _get_all_paddle_includes_from_include_root(
+                    str(paddle_include_dir)
+                )
+            )
+
+        custom_module = paddle.utils.cpp_extension.load(
+            name='custom_device',
+            sources=['custom_op.cc'],
+            extra_include_paths=paddle_includes,  # add for Coverage CI
+            extra_cxx_cflags=["-w", "-g"],  # test for cc flags
+            # build_directory=self.cur_dir,
+            verbose=True,
+        )
+        self.custom_op = custom_module.custom_relu
+        self.custom_stream_op = custom_module.custom_stream
+
+        self.dtypes = ["float32", "float64"]
+        self.device = "custom_cpu"
+
+        # config seed
+        SEED = 2021
+        paddle.seed(SEED)
+        paddle.framework.random._manual_program_seed(SEED)
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+        del os.environ['CUSTOM_DEVICE_ROOT']
+
+    def test_custom_device(self):
+        self._test_static()
+        self._test_dynamic()
+        self._test_double_grad_dynamic()
+        self._test_with_dataloader()
+        self._test_stream()
+
+    def _test_static(self):
+        for dtype in self.dtypes:
+            x = np.random.uniform(-1, 1, [4, 8]).astype(dtype)
+            out = custom_relu_static(self.custom_op, self.device, dtype, x)
+            pd_out = custom_relu_static(
+                self.custom_op, self.device, dtype, x, False
+            )
+            np.testing.assert_array_equal(
+                out,
+                pd_out,
+                err_msg=f"custom op out: {out},\n paddle api out: {pd_out}",
+            )
+
+    def _test_dynamic(self):
+        for dtype in self.dtypes:
+            x = np.random.uniform(-1, 1, [4, 8]).astype(dtype)
+            out, x_grad = custom_relu_dynamic(
+                self.custom_op, self.device, dtype, x
+            )
+            pd_out, pd_x_grad = custom_relu_dynamic(
+                self.custom_op, self.device, dtype, x, False
+            )
+            np.testing.assert_array_equal(
+                out,
+                pd_out,
+                err_msg=f"custom op out: {out},\n paddle api out: {pd_out}",
+            )
+            np.testing.assert_array_equal(
+                x_grad,
+                pd_x_grad,
+                err_msg=f"custom op x grad: {x_grad},\n paddle api x grad: {pd_x_grad}",
+            )
+
+    def _test_double_grad_dynamic(self):
+        for dtype in self.dtypes:
+            x = np.random.uniform(-1, 1, [4, 8]).astype(dtype)
+            out, dx_grad = custom_relu_double_grad_dynamic(
+                self.custom_op, self.device, dtype, x
+            )
+            pd_out, pd_dx_grad = custom_relu_double_grad_dynamic(
+                self.custom_op, self.device, dtype, x, False
+            )
+            np.testing.assert_array_equal(
+                out,
+                pd_out,
+                err_msg=f"custom op out: {out},\n paddle api out: {pd_out}",
+            )
+            np.testing.assert_array_equal(
+                dx_grad,
+                pd_dx_grad,
+                err_msg=f"custom op dx grad: {dx_grad},\n paddle api dx grad: {pd_dx_grad}",
+            )
+
+    def _test_with_dataloader(self):
+        import paddle
+        from paddle.vision.transforms import Compose, Normalize
+
+        paddle.set_device(self.device)
+        # data loader
+        transform = Compose(
+            [Normalize(mean=[127.5], std=[127.5], data_format="CHW")]
+        )
+        train_dataset = paddle.vision.datasets.MNIST(
+            mode="train", transform=transform
+        )
+        train_loader = paddle.io.DataLoader(
+            train_dataset,
+            batch_size=64,
+            shuffle=True,
+            drop_last=True,
+            num_workers=0,
+        )
+
+        for batch_id, (image, _) in enumerate(train_loader()):
+            out = self.custom_op(image)
+            pd_out = paddle.nn.functional.relu(image)
+            np.testing.assert_array_equal(
+                out,
+                pd_out,
+                err_msg=f"custom op out: {out},\n paddle api out: {pd_out}",
+            )
+
+            if batch_id == 5:
+                break
+
+    def _test_stream(self):
+        import paddle
+
+        paddle.set_device(self.device)
+        x = paddle.ones([2, 2], dtype='float32')
+        out = self.custom_stream_op(x)
+
+        np.testing.assert_array_equal(x.numpy(), out.numpy())
+
+
+if __name__ == "__main__":
+    if os.name == 'nt' or sys.platform.startswith('darwin'):
+        # only support Linux now
+        sys.exit()
+    unittest.main()

test/custom_runtime/test_fleet_launch_custom_device.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+# 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.
+
+set -e
+
+temp_dir=$(mktemp --directory)
+pushd ${temp_dir} \
+&& git clone --depth 1 ${PLUGIN_URL} \
+&& pushd PaddleCustomDevice/ \
+&& git fetch origin \
+&& git checkout ${PLUGIN_TAG} -b dev \
+&& pushd backends/custom_cpu \
+&& mkdir build && pushd build && cmake .. -DPython_EXECUTABLE=${PYTHON_EXECUTABLE} -DWITH_TESTING=OFF && make -j8 && popd && popd && popd && popd
+
+echo "begin test use custom_cpu"
+
+export FLAGS_selected_custom_cpus=0,1
+export CUSTOM_CPU_VISIBLE_DEVICES=0,1
+export CUSTOM_DEVICE_ROOT=${temp_dir}/PaddleCustomDevice/backends/custom_cpu/build
+
+distributed_args="--devices=0,1"
+python -m paddle.distributed.fleet.launch ${distributed_args} custom_device_multi_process_collective.py fleetlaunch_custom_cpu