Run cpu constand folding with cpu constant folders (#4525)

Keep cpu and non-cpu constant folders isolated (sqrt test was shared)
Co-authored-by: Sang Ik Lee <sang.ik.lee@intel.com>

src/ngraph/pass/constant_folding_unary.cpp

@@ -52,16 +52,6 @@ shared_ptr<op::Constant> fold_constant_unary(shared_ptr<op::Constant> constant,
                                              shared_ptr<Node> unary,
                                              NodeExecutorTy func)
 {
-    // check sqrt arg
-    if (is_type<op::Sqrt>(unary))
-    {
-        std::vector<T> values{constant->get_vector<T>()};
-        if (std::any_of(values.begin(), values.end(), [](T i) { return i < T(0); }))
-        {
-            throw ngraph_error("Square root of negative value");
-        }
-    }
-
     const Shape& out_shape = unary->get_shape();
     runtime::AlignedBuffer buffer(shape_size(out_shape) * sizeof(T));
 
@@ -123,6 +113,11 @@ shared_ptr<op::Constant> fold_constant_unary(shared_ptr<op::Constant> constant,
         }
         else if (is_type<op::Sqrt>(unary))
         {
+            std::vector<T> values{constant->get_vector<T>()};
+            if (std::any_of(values.begin(), values.end(), [](T i) { return i < T(0); }))
+            {
+                throw ngraph_error("Square root of negative value");
+            }
             runtime::reference::sqrt<T>(
                 constant->get_data_ptr<T>(), buffer.get_ptr<T>(), shape_size(out_shape));
         }

src/ngraph/runtime/cpu/cpu_builder.cpp

@@ -564,7 +564,7 @@ namespace ngraph
             template <>
             NodeExecutorTy Builder::BUILDER_CF_DECL(ngraph::op::Sqrt)
             {
-                BUILD_UNARY_ELEMWISE_CF_FUNCTOR(runtime::cpu::kernel::sqrt);
+                BUILD_UNARY_ELEMWISE_CF_FUNCTOR(runtime::cpu::kernel::checked_sqrt);
             }
 
             template <>

src/ngraph/runtime/cpu/kernel/sqrt.hpp

@@ -44,6 +44,29 @@ namespace ngraph
                     out.device(ngraph::runtime::cpu::executor::GetCPUExecutor().get_device(arena)) =
                         in.sqrt();
                 }
+
+                template <typename ElementType>
+                void checked_sqrt(void* input, void* output, size_t count, int arena)
+                {
+                    ElementType* elts = static_cast<ElementType*>(input);
+                    if (std::any_of(
+                            elts, elts + count, [](ElementType i) { return i < ElementType(0); }))
+                    {
+                        throw ngraph_error("Square root of negative value");
+                    }
+
+                    Eigen::array<Eigen::Index, 1> out_dims, in_dims;
+
+                    out_dims[0] = in_dims[0] = count;
+
+                    Eigen::TensorMap<Eigen::Tensor<ElementType, 1, Eigen::RowMajor>> out(
+                        static_cast<ElementType*>(output), out_dims);
+                    Eigen::TensorMap<Eigen::Tensor<ElementType, 1, Eigen::RowMajor>> in(
+                        static_cast<ElementType*>(input), in_dims);
+
+                    out.device(ngraph::runtime::cpu::executor::GetCPUExecutor().get_device(arena)) =
+                        in.sqrt();
+                }
             }
         }
     }

test/cpu_test.cpp

@@ -1059,6 +1059,8 @@ TEST(cpu_test, thread_safe_calls_convolution_2d_2items)
 // Thus this test is disabled with MLIR enabled.
 TEST(cpu_test, MLIR_DISABLE_TEST(constant_convertlayout))
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape data_shape{1, 64, 56, 56};
     auto data = make_shared<op::Parameter>(element::f32, data_shape);
     Shape weights_shape{64, 64, 3, 3};
@@ -1073,15 +1075,15 @@ TEST(cpu_test, MLIR_DISABLE_TEST(constant_convertlayout))
     auto convbias = make_shared<op::ConvolutionBias>(conv, bias);
 
     auto f = make_shared<Function>(convbias, ParameterVector{data, bias});
-    auto backend = runtime::Backend::create("CPU");
     auto handle = backend->compile(f);
-
     size_t convert_layout = count_ops_of_type<runtime::cpu::op::ConvertLayout>(f);
     ASSERT_EQ(convert_layout, 1);
 }
 
 TEST(cpu_test, constant_reshape)
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape shape_in{2, 4};
     Shape shape_out{2, 4, 1};
 
@@ -1107,6 +1109,8 @@ TEST(cpu_test, constant_reshape)
 
 TEST(cpu_test, constant_reshape_permute)
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape shape_in{2, 4};
     Shape shape_out{4, 2};
 
@@ -1133,6 +1137,8 @@ TEST(cpu_test, constant_reshape_permute)
 
 TEST(cpu_test, constant_broadcast)
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape shape_in{2};
     Shape shape_out{2, 4};
 
@@ -1159,6 +1165,8 @@ TEST(cpu_test, constant_broadcast)
 
 TEST(cpu_test, constant_pad_exterior)
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape shape_in{2};
 
     vector<int> values_in{777, 888};
@@ -1196,6 +1204,8 @@ static std::vector<T> get_result_constant(std::shared_ptr<Function> f, size_t po
 
 TEST(cpu_test, constant_unary_binary)
 {
+    // Initialize CPU constant folders
+    auto backend = runtime::Backend::create("CPU");
     Shape shape_in{2, 2};
     vector<int> values_a{1, 2, 3, 4};
     vector<int> values_b{1, 2, 3, 4};