fix get tensor backend set bug (PaddlePaddle#41478)

douch · Apr 10, 2022 · a98fbe8 · a98fbe8
1 parent 51347e1
commit a98fbe8
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 3 deletions.
diff --git a/paddle/phi/api/lib/kernel_dispatch.cc b/paddle/phi/api/lib/kernel_dispatch.cc
@@ -14,18 +14,46 @@ limitations under the License. */
 
 #include "paddle/phi/api/lib/kernel_dispatch.h"
 
-#include "paddle/phi/api/include/context_pool.h"
-#include "paddle/phi/core/compat/convert_utils.h"
 #ifdef _MSC_VER
 #include <intrin.h>
 #endif
 
+#include "paddle/phi/api/include/context_pool.h"
+#include "paddle/phi/core/compat/convert_utils.h"
+#include "paddle/phi/core/string_tensor_utils.h"
+#include "paddle/phi/core/tensor_utils.h"
+
 namespace paddle {
 namespace experimental {
 namespace detail {
 
+// We need judge whether the allocation is nullptr,
+// whether the allocation is initialized, wo we need GetHolder method
+bool HasAllocation(const phi::TensorBase& t) {
+  if (phi::DenseTensor::classof(&t)) {
+    return phi::DenseTensorUtils::GetHolder(
+               static_cast<const phi::DenseTensor&>(t)) != nullptr;
+  } else if (phi::SelectedRows::classof(&t)) {
+    return phi::DenseTensorUtils::GetHolder(
+               static_cast<const phi::SelectedRows&>(t).value()) != nullptr;
+  } else if (phi::SparseCsrTensor::classof(&t)) {
+    return phi::DenseTensorUtils::GetHolder(
+               static_cast<const phi::SparseCsrTensor&>(t)
+                   .non_zero_elements()) != nullptr;
+  } else if (phi::SparseCooTensor::classof(&t)) {
+    return phi::DenseTensorUtils::GetHolder(
+               static_cast<const phi::SparseCooTensor&>(t)
+                   .non_zero_elements()) != nullptr;
+  } else if (phi::StringTensor::classof(&t)) {
+    return phi::StringTensorUtils::GetHolder(
+               static_cast<const phi::StringTensor&>(t)) != nullptr;
+  } else {
+    return false;
+  }
+}
+
 BackendSet GetTensorBackendSet(const phi::TensorBase& t) {
-  if (t.initialized()) {
+  if (HasAllocation(t)) {
     BackendSet backend_set(phi::TransToPhiBackend(t.place()));
     switch (t.layout()) {
       case DataLayout::MKLDNN:

diff --git a/paddle/phi/core/string_tensor_utils.h b/paddle/phi/core/string_tensor_utils.h
@@ -23,6 +23,11 @@ class StringTensorUtils {
   static StringTensorMeta* GetMutableMeta(StringTensor* tensor) {
     return &(tensor->meta_);
   }
+
+  static const std::shared_ptr<phi::Allocation>& GetHolder(
+      const StringTensor& tensor) {
+    return tensor.holder_;
+  }
 };
 
 }  // namespace phi
diff --git a/paddle/phi/core/tensor_utils.h b/paddle/phi/core/tensor_utils.h
@@ -25,6 +25,11 @@ class DenseTensorUtils {
     return &(tensor->meta_);
   }
 
+  static const std::shared_ptr<phi::Allocation>& GetHolder(
+      const DenseTensor& tensor) {
+    return tensor.holder_;
+  }
+
   static DenseTensor Slice(const DenseTensor& tensor,
                            int64_t begin_idx,
                            int64_t end_idx) {