[CustomDevice] fix SplitDenseTensor (#55615)

paddle/fluid/memory/allocation/naive_best_fit_allocator.cc

@@ -531,7 +531,9 @@ void Free<platform::CustomPlace>(const platform::CustomPlace &place,
                                  size_t size) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   VLOG(10) << "Free pointer=" << p << " on " << platform::Place(place);
+  if (phi::DeviceManager::HasDeviceType(place.GetDeviceType())) {
     GetBuddyAllocator(place)->Free(p);
+  }
 #else
   PADDLE_THROW(platform::errors::PermissionDenied(
       "'CustomPlace' is not supported in CPU only device."));

paddle/fluid/pybind/process_group_utils.h

@@ -96,12 +96,22 @@ struct SplitDenseTensor<platform::CustomDeviceContext, T> {
                                       const phi::Scalar &,
                                       std::vector<phi::DenseTensor *>);
     auto *kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+
+    auto in_dims = phi::vectorize(in.dims());
+    auto origin_out_dims = phi::vectorize(out->at(0)->dims());
+    for (auto *tensor : *out) {
+      if (origin_out_dims.size() != in_dims.size()) {
+        std::vector<int> new_dims({1});
+        new_dims.insert(
+            new_dims.end(), origin_out_dims.begin(), origin_out_dims.end());
+        tensor->Resize(phi::make_ddim(new_dims));
+      }
+    }
     (*kernel_fn)(context, in, out->size(), phi::Scalar(0), *out);
     for (auto *tensor : *out) {
-      auto dim_vec = phi::vectorize<int>(tensor->dims());
-      if (dim_vec.size() > 0 && dim_vec[0] == 1) {
-        tensor->Resize(phi::make_ddim(
-            std::vector<int>(dim_vec.begin() + 1, dim_vec.end())));
+      auto tensor_dims = phi::vectorize(tensor->dims());
+      if (tensor_dims.size() != origin_out_dims.size()) {
+        tensor->Resize(phi::make_ddim(origin_out_dims));
       }
     }
   }

test/custom_runtime/test_collective_process_group_xccl.py

@@ -150,12 +150,12 @@ class TestProcessGroup(TestMultipleCustomCPU):
         cur_dir = os.path.dirname(os.path.abspath(__file__))
         self.temp_dir = tempfile.TemporaryDirectory()
         cmd = 'cd {} \
-            && git clone {} \
+            && git clone --depth 1 {} \
             && cd PaddleCustomDevice \
             && git fetch origin \
             && git checkout {} -b dev \
             && cd backends/custom_cpu \
-            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} && make -j8'.format(
+            && 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'),

test/custom_runtime/test_custom_cpu_plugin.py

@@ -26,12 +26,12 @@ class TestCustomCPUPlugin(unittest.TestCase):
         cur_dir = os.path.dirname(os.path.abspath(__file__))
         self.temp_dir = tempfile.TemporaryDirectory()
         cmd = 'cd {} \
-            && git clone {} \
+            && git clone --depth 1 {} \
             && cd PaddleCustomDevice \
             && git fetch origin \
             && git checkout {} -b dev \
             && cd backends/custom_cpu \
-            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} && make -j8'.format(
+            && 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'),

test/custom_runtime/test_custom_cpu_profiler_plugin.py

@@ -24,12 +24,12 @@ class TestCustomCPUProfilerPlugin(unittest.TestCase):
         cur_dir = os.path.dirname(os.path.abspath(__file__))
         self.temp_dir = tempfile.TemporaryDirectory()
         cmd = 'cd {} \
-            && git clone {} \
+            && git clone --depth 1 {} \
             && cd PaddleCustomDevice \
             && git fetch origin \
             && git checkout {} -b dev \
             && cd backends/custom_cpu \
-            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} && make -j8'.format(
+            && 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'),

test/custom_runtime/test_custom_cpu_to_static.py

@@ -106,12 +106,12 @@ class TestCustomCPUPlugin(unittest.TestCase):
         cur_dir = os.path.dirname(os.path.abspath(__file__))
         self.temp_dir = tempfile.TemporaryDirectory()
         cmd = 'cd {} \
-            && git clone {} \
+            && git clone --depth 1 {} \
             && cd PaddleCustomDevice \
             && git fetch origin \
             && git checkout {} -b dev \
             && cd backends/custom_cpu \
-            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} && make -j8'.format(
+            && 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'),

test/custom_runtime/test_custom_op_setup.py

@@ -104,12 +104,12 @@ class TestNewCustomOpSetUpInstall(unittest.TestCase):
         self.cur_dir = os.path.dirname(os.path.abspath(__file__))
         self.temp_dir = tempfile.TemporaryDirectory()
         cmd = 'cd {} \
-            && git clone {} \
+            && git clone --depth 1 {} \
             && cd PaddleCustomDevice \
             && git fetch origin \
             && git checkout {} -b dev \
             && cd backends/custom_cpu \
-            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} && make -j8 \
+            && mkdir build && cd build && cmake .. -DPython_EXECUTABLE={} -DWITH_TESTING=OFF && make -j8 \
             && cd {}'.format(
             self.temp_dir.name,
             os.getenv('PLUGIN_URL'),

test/custom_runtime/test_fleet_launch_custom_device.sh

@@ -18,12 +18,12 @@ set -e
 
 temp_dir=$(mktemp --directory)
 pushd ${temp_dir} \
-&& git clone ${PLUGIN_URL} \
+&& 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} && make -j8 && popd && popd && popd && popd
+&& mkdir build && pushd build && cmake .. -DPython_EXECUTABLE=${PYTHON_EXECUTABLE} -DWITH_TESTING=OFF && make -j8 && popd && popd && popd && popd
 
 echo "begin test use custom_cpu"