From e6ff5edc3dcbc6821550e73eb24585fab96ec12a Mon Sep 17 00:00:00 2001
From: Nikolay Shchegolev <nikolay.shchegolev@intel.com>
Date: Wed, 1 Nov 2023 15:42:25 +0400
Subject: [PATCH] [CPU] Empty tensor used with a custom op leads to CPU plugin
 exception. (#19733)

---
 src/plugins/intel_cpu/src/node.h              |  1 +
 src/plugins/intel_cpu/src/nodes/reference.cpp | 20 +++++++++++++++----
 src/plugins/intel_cpu/src/nodes/reference.h   |  1 +
 .../subgraph_tests/src/custom_op_scalar.cpp   |  1 +
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/plugins/intel_cpu/src/node.h b/src/plugins/intel_cpu/src/node.h
index 4b6fa3a87f7..dba25591386 100644
--- a/src/plugins/intel_cpu/src/node.h
+++ b/src/plugins/intel_cpu/src/node.h
@@ -42,6 +42,7 @@
 #include "nodes/executors/executor.hpp"
 
 #define THROW_CPU_NODE_ERR(...) OPENVINO_THROW(getTypeStr(), " node with name '", getName(), "' ", __VA_ARGS__)
+#define CPU_NODE_ASSERT(condition, ...) OPENVINO_ASSERT(condition, getTypeStr(), " node with name '", getName(), "' ", __VA_ARGS__)
 
 namespace ov {
 namespace intel_cpu {
diff --git a/src/plugins/intel_cpu/src/nodes/reference.cpp b/src/plugins/intel_cpu/src/nodes/reference.cpp
index 091e3181312..0ba82bacec5 100644
--- a/src/plugins/intel_cpu/src/nodes/reference.cpp
+++ b/src/plugins/intel_cpu/src/nodes/reference.cpp
@@ -108,22 +108,34 @@ bool Reference::needShapeInfer() const {
 
 ov::TensorVector Reference::prepareInputs() const {
     ov::TensorVector inputs;
-    for (size_t i = 0; i < inputShapes.size(); i++) {
+    for (size_t i = 0lu; i < inputShapes.size(); i++) {
         void *srcDataPtr = getParentEdgesAtPort(i)[0]->getMemory().getData();
         ov::Shape shape = ovCoreNode->get_input_partial_shape(i).rank().get_length() == 0 ?
                 ov::Shape{} : getParentEdgesAtPort(i)[0]->getMemory().getStaticDims();
-        inputs.push_back(ov::Tensor(ovCoreNode->get_input_element_type(i), shape, srcDataPtr));
+
+        if (std::any_of(shape.begin(), shape.end(), [](const size_t dim) { return dim == 0lu; } )) {
+            inputs.push_back(ov::Tensor(ovCoreNode->get_input_element_type(i), shape));
+        } else {
+            CPU_NODE_ASSERT(srcDataPtr, "has empty input data on port ", i);
+            inputs.push_back(ov::Tensor(ovCoreNode->get_input_element_type(i), shape, srcDataPtr));
+        }
     }
     return inputs;
 }
 
 ov::TensorVector Reference::prepareOutputs() const {
     ov::TensorVector outputs;
-    for (size_t i = 0; i < outputShapes.size(); i++) {
+    for (size_t i = 0lu; i < outputShapes.size(); i++) {
         void *dstDataPtr = getChildEdgesAtPort(i)[0]->getMemory().getData();
         ov::Shape shape = ovCoreNode->get_output_partial_shape(i).rank().get_length() == 0 ?
                 ov::Shape{} : getChildEdgesAtPort(i)[0]->getMemory().getStaticDims();
-        outputs.push_back(ov::Tensor(ovCoreNode->get_output_element_type(i), shape, dstDataPtr));
+
+        if (std::any_of(shape.begin(), shape.end(), [](const size_t dim) { return dim == 0lu; } )) {
+            outputs.push_back(ov::Tensor(ovCoreNode->get_output_element_type(i), shape));
+        } else {
+            CPU_NODE_ASSERT(dstDataPtr, "has empty output data on port ", i);
+            outputs.push_back(ov::Tensor(ovCoreNode->get_output_element_type(i), shape, dstDataPtr));
+        }
     }
     return outputs;
 }
diff --git a/src/plugins/intel_cpu/src/nodes/reference.h b/src/plugins/intel_cpu/src/nodes/reference.h
index c2453835229..9f35f0398e5 100644
--- a/src/plugins/intel_cpu/src/nodes/reference.h
+++ b/src/plugins/intel_cpu/src/nodes/reference.h
@@ -22,6 +22,7 @@ public:
 
     bool needShapeInfer() const override;
     bool needPrepareParams() const override { return false; }
+    bool isExecutable() const override { return true; }
     void executeDynamicImpl(dnnl::stream strm) override;
 
 private:
diff --git a/src/plugins/intel_cpu/tests/functional/subgraph_tests/src/custom_op_scalar.cpp b/src/plugins/intel_cpu/tests/functional/subgraph_tests/src/custom_op_scalar.cpp
index 15eef344b6f..c108d067a47 100644
--- a/src/plugins/intel_cpu/tests/functional/subgraph_tests/src/custom_op_scalar.cpp
+++ b/src/plugins/intel_cpu/tests/functional/subgraph_tests/src/custom_op_scalar.cpp
@@ -140,6 +140,7 @@ TEST_P(CustomOpScalarCPUTest, CompareWithRefs) {
 
 const std::vector<InputShape> inputShapes = {
     {{}, {{2, 3, 16}}},
+    {{2, 3, -1}, {{2, 3, 0}}},
     {{}, {{}}}
 };