add NPU support for zero_copy_tensor. (#34629)

* add NPU support for zero_copy_tensor. * revert unnesessary codes. * revert unnesessary codes.
PaddlePaddle · Aug 5, 2021 · a68709d · a68709d
1 parent 7a38b76
commit a68709d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/paddle/fluid/inference/api/details/zero_copy_tensor.cc b/paddle/fluid/inference/api/details/zero_copy_tensor.cc
@@ -65,10 +65,13 @@ T *Tensor::mutable_data(PlaceType place) {
     case static_cast<int>(PlaceType::kXPU): {
       return tensor->mutable_data<T>(paddle::platform::XPUPlace(device_));
     }
+    case static_cast<int>(PlaceType::kNPU): {
+      return tensor->mutable_data<T>(paddle::platform::NPUPlace(device_));
+    }
     default:
       PADDLE_THROW(paddle::platform::errors::Unavailable(
-          "Only CPU / CUDA / XPU places is supported. The place `%d` is not "
-          "supported.",
+          "Only CPU / CUDA / XPU / NPU places is supported. The place `%d` is "
+          "not supported.",
           static_cast<int>(place)));
       break;
   }
@@ -86,6 +89,8 @@ T *Tensor::data(PlaceType *place, int *size) const {
     *place = PlaceType::kGPU;
   } else if (paddle::platform::is_xpu_place(tensor->place())) {
     *place = PlaceType::kXPU;
+  } else if (paddle::platform::is_npu_place(tensor->place())) {
+    *place = PlaceType::kNPU;
   } else {
     *place = PlaceType::kUNK;
   }

diff --git a/paddle/fluid/inference/api/details/zero_copy_tensor_test.cc b/paddle/fluid/inference/api/details/zero_copy_tensor_test.cc
@@ -133,6 +133,10 @@ TEST(Tensor, FillRandomDataAndCheck) {
   ASSERT_TRUE(FillRandomDataAndCheck(PlaceType::kGPU));
   ASSERT_TRUE(SetPlaceAndCheck(PlaceType::kGPU));
 #endif
+#ifdef PADDLE_WITH_ASCEND_CL
+  ASSERT_TRUE(FillRandomDataAndCheck(PlaceType::kNPU));
+  ASSERT_TRUE(SetPlaceAndCheck(PlaceType::kNPU));
+#endif
 }
 
 }  // namespace paddle_infer