Enable ConstantFolding for Product, Reverse; add a couple of dynamic tests (#3187)

* Add a few more toy dynamic examples

* Enable constant folding for Reverse; enable reverse_shape test

* Add constant folding for Product, and enable two more tests

src/ngraph/pass/constant_folding.cpp

@@ -31,9 +31,11 @@
 #include "ngraph/op/multiply.hpp"
 #include "ngraph/op/negative.hpp"
 #include "ngraph/op/pad.hpp"
+#include "ngraph/op/product.hpp"
 #include "ngraph/op/quantize.hpp"
 #include "ngraph/op/relu.hpp"
 #include "ngraph/op/reshape.hpp"
+#include "ngraph/op/reverse.hpp"
 #include "ngraph/op/sqrt.hpp"
 #include "ngraph/op/subtract.hpp"
 #include "ngraph/pattern/matcher.hpp"
@@ -49,9 +51,11 @@
 #include "ngraph/runtime/reference/multiply.hpp"
 #include "ngraph/runtime/reference/negate.hpp"
 #include "ngraph/runtime/reference/pad.hpp"
+#include "ngraph/runtime/reference/product.hpp"
 #include "ngraph/runtime/reference/quantize.hpp"
 #include "ngraph/runtime/reference/relu.hpp"
 #include "ngraph/runtime/reference/reshape.hpp"
+#include "ngraph/runtime/reference/reverse.hpp"
 #include "ngraph/runtime/reference/sqrt.hpp"
 #include "ngraph/runtime/reference/subtract.hpp"
 #include "ngraph/util.hpp"
@@ -215,7 +219,13 @@ void pass::ConstantFolding::construct_constant_reshape()
         if (type == element::i32)
         {
             replace_node(m.get_match_root(),
-                         fold_constant_reshape<int>(constant_match, reshape_match, func));
+                         fold_constant_reshape<int32_t>(constant_match, reshape_match, func));
+            return true;
+        }
+        if (type == element::i64)
+        {
+            replace_node(m.get_match_root(),
+                         fold_constant_reshape<int64_t>(constant_match, reshape_match, func));
             return true;
         }
         else if (type == element::i8)
@@ -956,3 +966,181 @@ void pass::ConstantFolding::construct_constant_shape_of()
         make_shared<pattern::Matcher>(shape_of_op, "ConstantFolding.ConstantShapeOf");
     this->add_matcher(shape_of_matcher, constant_shape_of_callback, all_pass_property_off);
 }
+
+template <typename T>
+static shared_ptr<op::Constant> fold_constant_reverse_helper(shared_ptr<op::Constant> constant,
+                                                             const AxisSet& reversed_axes)
+{
+    auto out_shape = constant->get_shape();
+    vector<T> out_vec(shape_size(out_shape));
+
+    runtime::reference::reverse<T>(
+        constant->get_vector<T>().data(), out_vec.data(), out_shape, out_shape, reversed_axes);
+
+    return make_shared<op::Constant>(constant->get_output_element_type(0), out_shape, out_vec);
+}
+
+static shared_ptr<op::Constant> fold_constant_reverse(shared_ptr<op::Constant> constant,
+                                                      const AxisSet& reversed_axes)
+{
+    auto& input_element_type = constant->get_output_element_type(0);
+
+#if !(defined(__GNUC__) && (__GNUC__ == 4 && __GNUC_MINOR__ == 8))
+#pragma GCC diagnostic push
+#pragma GCC diagnostic error "-Wswitch"
+#pragma GCC diagnostic error "-Wswitch-enum"
+#endif
+    switch (input_element_type.get_type_enum())
+    {
+    case element::Type_t::undefined:
+        NGRAPH_CHECK(false, "Encountered 'undefined' element type in fold_constant_convert");
+        break;
+    case element::Type_t::dynamic:
+        NGRAPH_CHECK(false, "Encountered 'dynamic' element type in fold_constant_convert");
+        break;
+    case element::Type_t::boolean:
+        return fold_constant_reverse_helper<char>(constant, reversed_axes);
+    case element::Type_t::bf16:
+        return fold_constant_reverse_helper<bfloat16>(constant, reversed_axes);
+    case element::Type_t::f16:
+        return fold_constant_reverse_helper<float16>(constant, reversed_axes);
+    case element::Type_t::f32: return fold_constant_reverse_helper<float>(constant, reversed_axes);
+    case element::Type_t::f64: return fold_constant_reverse_helper<double>(constant, reversed_axes);
+    case element::Type_t::i8: return fold_constant_reverse_helper<int8_t>(constant, reversed_axes);
+    case element::Type_t::i16:
+        return fold_constant_reverse_helper<int16_t>(constant, reversed_axes);
+    case element::Type_t::i32:
+        return fold_constant_reverse_helper<int32_t>(constant, reversed_axes);
+    case element::Type_t::i64:
+        return fold_constant_reverse_helper<int64_t>(constant, reversed_axes);
+    case element::Type_t::u8: return fold_constant_reverse_helper<uint8_t>(constant, reversed_axes);
+    case element::Type_t::u16:
+        return fold_constant_reverse_helper<uint16_t>(constant, reversed_axes);
+    case element::Type_t::u32:
+        return fold_constant_reverse_helper<uint32_t>(constant, reversed_axes);
+    case element::Type_t::u64:
+        return fold_constant_reverse_helper<uint64_t>(constant, reversed_axes);
+    }
+#if !(defined(__GNUC__) && (__GNUC__ == 4 && __GNUC_MINOR__ == 8))
+#pragma GCC diagnostic pop
+#endif
+}
+
+void pass::ConstantFolding::construct_constant_reverse()
+{
+    auto constant_label = make_shared<pattern::op::Label>(
+        element::i32, Shape{2, 3, 4}, pattern::has_class<op::Constant>());
+    auto convert_op = make_shared<op::Reverse>(constant_label, AxisSet{0, 1, 2});
+
+    auto constant_reverse_callback = [constant_label](pattern::Matcher& m) {
+        NGRAPH_DEBUG << "In callback for constant_reverse_callback against node = "
+                     << m.get_match_root()->get_name();
+
+        auto pattern_map = m.get_pattern_map();
+
+        auto constant_match = static_pointer_cast<op::Constant>(pattern_map[constant_label]);
+        auto reverse_match = static_pointer_cast<op::Reverse>(m.get_match_root());
+
+        replace_node(m.get_match_root(),
+                     fold_constant_reverse(constant_match, reverse_match->get_reversed_axes()));
+        return true;
+    };
+
+    auto convert_matcher =
+        make_shared<pattern::Matcher>(convert_op, "ConstantFolding.ConstantReverse");
+    this->add_matcher(convert_matcher, constant_reverse_callback, all_pass_property_off);
+}
+
+template <typename T>
+static shared_ptr<op::Constant> fold_constant_product_helper(shared_ptr<op::Constant> constant,
+                                                             const AxisSet& reduction_axes,
+                                                             const Shape& result_shape)
+{
+    vector<T> out_vec(shape_size(result_shape));
+
+    runtime::reference::product<T>(constant->get_vector<T>().data(),
+                                   out_vec.data(),
+                                   constant->get_output_shape(0),
+                                   result_shape,
+                                   reduction_axes);
+
+    return make_shared<op::Constant>(constant->get_output_element_type(0), result_shape, out_vec);
+}
+
+static shared_ptr<op::Constant> fold_constant_product(shared_ptr<op::Constant> constant,
+                                                      const AxisSet& reduction_axes,
+                                                      const Shape& result_shape)
+{
+    auto& input_element_type = constant->get_output_element_type(0);
+
+#if !(defined(__GNUC__) && (__GNUC__ == 4 && __GNUC_MINOR__ == 8))
+#pragma GCC diagnostic push
+#pragma GCC diagnostic error "-Wswitch"
+#pragma GCC diagnostic error "-Wswitch-enum"
+#endif
+    switch (input_element_type.get_type_enum())
+    {
+    case element::Type_t::undefined:
+        NGRAPH_CHECK(false, "Encountered 'undefined' element type in fold_constant_product");
+        break;
+    case element::Type_t::dynamic:
+        NGRAPH_CHECK(false, "Encountered 'dynamic' element type in fold_constant_product");
+        break;
+    case element::Type_t::boolean:
+        return fold_constant_product_helper<char>(constant, reduction_axes, result_shape);
+    case element::Type_t::bf16:
+        return fold_constant_product_helper<bfloat16>(constant, reduction_axes, result_shape);
+    case element::Type_t::f16:
+        return fold_constant_product_helper<float16>(constant, reduction_axes, result_shape);
+    case element::Type_t::f32:
+        return fold_constant_product_helper<float>(constant, reduction_axes, result_shape);
+    case element::Type_t::f64:
+        return fold_constant_product_helper<double>(constant, reduction_axes, result_shape);
+    case element::Type_t::i8:
+        return fold_constant_product_helper<int8_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::i16:
+        return fold_constant_product_helper<int16_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::i32:
+        return fold_constant_product_helper<int32_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::i64:
+        return fold_constant_product_helper<int64_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::u8:
+        return fold_constant_product_helper<uint8_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::u16:
+        return fold_constant_product_helper<uint16_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::u32:
+        return fold_constant_product_helper<uint32_t>(constant, reduction_axes, result_shape);
+    case element::Type_t::u64:
+        return fold_constant_product_helper<uint64_t>(constant, reduction_axes, result_shape);
+    }
+#if !(defined(__GNUC__) && (__GNUC__ == 4 && __GNUC_MINOR__ == 8))
+#pragma GCC diagnostic pop
+#endif
+}
+
+void pass::ConstantFolding::construct_constant_product()
+{
+    auto constant_label = make_shared<pattern::op::Label>(
+        element::i32, Shape{2, 3, 4}, pattern::has_class<op::Constant>());
+    auto convert_op = make_shared<op::Product>(constant_label, AxisSet{0, 1, 2});
+
+    auto constant_product_callback = [constant_label](pattern::Matcher& m) {
+        NGRAPH_DEBUG << "In callback for constant_product_callback against node = "
+                     << m.get_match_root()->get_name();
+
+        auto pattern_map = m.get_pattern_map();
+
+        auto constant_match = static_pointer_cast<op::Constant>(pattern_map[constant_label]);
+        auto product_match = static_pointer_cast<op::Product>(m.get_match_root());
+
+        replace_node(m.get_match_root(),
+                     fold_constant_product(constant_match,
+                                           product_match->get_reduction_axes(),
+                                           product_match->get_output_shape(0)));
+        return true;
+    };
+
+    auto convert_matcher =
+        make_shared<pattern::Matcher>(convert_op, "ConstantFolding.ConstantProduct");
+    this->add_matcher(convert_matcher, constant_product_callback, all_pass_property_off);
+}

src/ngraph/pass/constant_folding.hpp

@@ -40,7 +40,9 @@ public:
         BINARY,
         QUANTIZE,
         CONVERT,
-        SHAPE_OF
+        SHAPE_OF,
+        REVERSE,
+        PRODUCT
     };
 
     ConstantFolding(const ngraph::BuildNodeExecutorMap& cfmap = ngraph::BuildNodeExecutorMap())
@@ -56,6 +58,8 @@ public:
         construct_constant_dequantize();
         construct_constant_convert();
         construct_constant_shape_of();
+        construct_constant_reverse();
+        construct_constant_product();
     }
 
     //this allows to specify the order in which matchers will be run
@@ -78,6 +82,8 @@ public:
             case CFTransformations::QUANTIZE: construct_constant_quantize(); break;
             case CFTransformations::CONVERT: construct_constant_convert(); break;
             case CFTransformations::SHAPE_OF: construct_constant_shape_of(); break;
+            case CFTransformations::REVERSE: construct_constant_reverse(); break;
+            case CFTransformations::PRODUCT: construct_constant_product(); break;
             }
         }
     }
@@ -92,6 +98,8 @@ private:
     void construct_constant_dequantize();
     void construct_constant_convert();
     void construct_constant_shape_of();
+    void construct_constant_reverse();
+    void construct_constant_product();
 
     ngraph::BuildNodeExecutorMap m_cfmap;
 };

src/ngraph/runtime/reference/product.hpp

@@ -47,8 +47,9 @@ namespace ngraph
                 {
                     Coordinate output_coord = reduce(input_coord, reduction_axes);
 
-                    out[output_transform.index(output_coord)] *=
-                        arg[input_transform.index(input_coord)];
+                    size_t output_index = output_transform.index(output_coord);
+
+                    out[output_index] = out[output_index] * arg[input_transform.index(input_coord)];
                 }
             }
         }

test/constant_folding.cpp

@@ -358,6 +358,56 @@ TEST(constant_folding, shape_of_rank_dynamic)
         PartialShape{Dimension::dynamic()}));
 }
 
+TEST(constant_folding, const_reverse)
+{
+    Shape input_shape{3, 3};
+
+    vector<int32_t> values_in{1, 2, 3, 4, 5, 6, 7, 8, 9};
+    auto constant = op::Constant::create(element::i32, input_shape, values_in);
+    auto convert = make_shared<op::Reverse>(constant, AxisSet{1});
+    auto f = make_shared<Function>(convert, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::Reverse>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    auto new_const =
+        std::dynamic_pointer_cast<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    auto values_out = new_const->get_vector<int32_t>();
+
+    vector<int32_t> values_expected{3, 2, 1, 6, 5, 4, 9, 8, 7};
+    ASSERT_EQ(values_expected, values_out);
+}
+
+TEST(constant_folding, const_product)
+{
+    Shape input_shape{3, 3};
+
+    vector<int32_t> values_in{1, 2, 3, 4, 5, 6, 7, 8, 9};
+    auto constant = op::Constant::create(element::i32, input_shape, values_in);
+    auto convert = make_shared<op::Product>(constant, AxisSet{1});
+    auto f = make_shared<Function>(convert, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::Product>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    auto new_const =
+        std::dynamic_pointer_cast<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    auto values_out = new_const->get_vector<int32_t>();
+
+    vector<int32_t> values_expected{6, 120, 504};
+    ASSERT_EQ(values_expected, values_out);
+}
+
 TEST(constant_folding, pass_property)
 {
     auto pass = std::make_shared<ngraph::pass::ConstantFolding>();

test/dynamic.in.cpp

@@ -456,3 +456,190 @@ NGRAPH_TEST(dynamic_${BACKEND_NAME}, reshape)
         ASSERT_EQ(results, data);
     }
 }
+
+static void axpy_test(const PartialShape& input_pshape, const std::vector<Shape>& input_shapes)
+{
+    auto a = make_shared<op::Parameter>(element::f32, input_pshape);
+    auto x = make_shared<op::Parameter>(element::f32, input_pshape);
+    auto y = make_shared<op::Parameter>(element::f32, input_pshape);
+
+    auto axpy = a * x + y;
+
+    auto f = make_shared<Function>(NodeVector{axpy}, ParameterVector{a, x, y});
+    auto backend = runtime::Backend::create("${BACKEND_NAME}", true);
+    auto ex = backend->compile(f);
+
+    auto t_r = backend->create_dynamic_tensor(element::f32, input_pshape);
+
+    for (auto& shape : input_shapes)
+    {
+        vector<float> inputs(shape_size(shape));
+        for (size_t i = 0; i < shape_size(shape); i++)
+        {
+            inputs[i] = i;
+        }
+
+        auto t_a = backend->create_tensor(element::f32, shape);
+        auto t_x = backend->create_tensor(element::f32, shape);
+        auto t_y = backend->create_tensor(element::f32, shape);
+
+        copy_data(t_a, inputs);
+        copy_data(t_x, inputs);
+        copy_data(t_y, inputs);
+
+        ex->call_with_validate({t_r}, {t_a, t_x, t_y});
+
+        ASSERT_EQ(t_r->get_shape(), shape);
+
+        auto results = read_vector<float>(t_r);
+
+        vector<float> expected_values(shape_size(shape));
+        for (size_t i = 0; i < shape_size(shape); i++)
+        {
+            expected_values[i] = (i * i) + i;
+        }
+
+        EXPECT_TRUE(test::all_close_f(results, expected_values));
+    }
+}
+
+NGRAPH_TEST(dynamic_${BACKEND_NAME}, axpy)
+{
+    // Test with shape {?, 3, 3}.
+    axpy_test(PartialShape{Dimension::dynamic(), 3, 3}, {Shape{2, 3, 3}, Shape{5, 3, 3}});
+
+    // Test with shape {?, ?, ?}.
+    axpy_test(PartialShape::dynamic(3),
+              {Shape{2, 3, 3}, Shape{5, 3, 3}, Shape{2, 5, 2}, Shape{8, 1, 8}});
+
+    // Test with shape ?. (Rank unknown.)
+    axpy_test(PartialShape::dynamic(),
+              {Shape{2, 3, 3},
+               Shape{5, 3, 3},
+               Shape{2, 5, 2},
+               Shape{8, 1, 8},
+               Shape{5},
+               Shape{8, 2},
+               Shape{8, 2, 8, 2},
+               Shape{2, 3, 4, 5, 2}});
+}
+
+static void to_vector_test(const PartialShape& input_pshape, const std::vector<Shape>& input_shapes)
+{
+    auto x = make_shared<op::Parameter>(element::f32, input_pshape);
+
+    shared_ptr<Node> x_new_shape = make_shared<op::ShapeOf>(x);
+    x_new_shape = make_shared<op::Product>(x_new_shape, AxisSet{0});
+    x_new_shape = make_shared<op::Reshape>(x_new_shape, AxisVector{}, Shape{1});
+
+    auto x_reshaped = make_shared<op::DynReshape>(x, x_new_shape);
+
+    auto f = make_shared<Function>(NodeVector{x_reshaped}, ParameterVector{x});
+    auto backend = runtime::Backend::create("${BACKEND_NAME}", true);
+    auto ex = backend->compile(f);
+
+    auto t_r = backend->create_dynamic_tensor(element::f32, PartialShape::dynamic(1));
+
+    for (auto& shape : input_shapes)
+    {
+        vector<float> inputs(shape_size(shape));
+        for (size_t i = 0; i < shape_size(shape); i++)
+        {
+            inputs[i] = i;
+        }
+
+        auto t_x = backend->create_tensor(element::f32, shape);
+
+        copy_data(t_x, inputs);
+
+        ex->call_with_validate({t_r}, {t_x});
+
+        ASSERT_EQ(t_r->get_shape(), (Shape{shape_size(shape)}));
+
+        auto results = read_vector<float>(t_r);
+
+        EXPECT_TRUE(test::all_close_f(results, inputs));
+    }
+}
+
+NGRAPH_TEST(dynamic_${BACKEND_NAME}, to_vector)
+{
+    // Test with shape {?, 3, 3}.
+    to_vector_test(PartialShape{Dimension::dynamic(), 3, 3}, {Shape{2, 3, 3}, Shape{5, 3, 3}});
+
+    // Test with shape {?, ?, ?}.
+    to_vector_test(PartialShape::dynamic(3),
+                   {Shape{2, 3, 3}, Shape{5, 3, 3}, Shape{2, 5, 2}, Shape{8, 1, 8}});
+
+    // Test with shape ?. (Rank unknown.)
+    to_vector_test(PartialShape::dynamic(),
+                   {Shape{2, 3, 3},
+                    Shape{5, 3, 3},
+                    Shape{2, 5, 2},
+                    Shape{8, 1, 8},
+                    Shape{5},
+                    Shape{8, 2},
+                    Shape{8, 2, 8, 2},
+                    Shape{2, 3, 4, 5, 2}});
+}
+
+static void reverse_shape_test(const PartialShape& input_pshape,
+                               const std::vector<Shape>& input_shapes)
+{
+    auto x = make_shared<op::Parameter>(element::f32, input_pshape);
+
+    shared_ptr<Node> x_new_shape = make_shared<op::ShapeOf>(x);
+    x_new_shape = make_shared<op::Reverse>(x_new_shape, AxisSet{0});
+
+    auto x_reshaped = make_shared<op::DynReshape>(x, x_new_shape);
+
+    auto f = make_shared<Function>(NodeVector{x_reshaped}, ParameterVector{x});
+    auto backend = runtime::Backend::create("${BACKEND_NAME}", true);
+    auto ex = backend->compile(f);
+
+    auto t_r = backend->create_dynamic_tensor(element::f32, PartialShape::dynamic());
+
+    for (auto& shape : input_shapes)
+    {
+        vector<float> inputs(shape_size(shape));
+        for (size_t i = 0; i < shape_size(shape); i++)
+        {
+            inputs[i] = i;
+        }
+
+        auto t_x = backend->create_tensor(element::f32, shape);
+
+        copy_data(t_x, inputs);
+
+        ex->call_with_validate({t_r}, {t_x});
+
+        Shape expected_shape = shape;
+        std::reverse(expected_shape.begin(), expected_shape.end());
+        ASSERT_EQ(t_r->get_shape(), expected_shape);
+
+        auto results = read_vector<float>(t_r);
+
+        EXPECT_TRUE(test::all_close_f(results, inputs));
+    }
+}
+
+NGRAPH_TEST(dynamic_${BACKEND_NAME}, reverse_shape)
+{
+    // Test with shape {?, 3, 3}.
+    reverse_shape_test(PartialShape{Dimension::dynamic(), 3, 3}, {Shape{2, 3, 3}, Shape{5, 3, 3}});
+
+    // Test with shape {?, ?, ?}.
+    reverse_shape_test(PartialShape::dynamic(3),
+                       {Shape{2, 3, 3}, Shape{5, 3, 3}, Shape{2, 5, 2}, Shape{8, 1, 8}});
+
+    // Test with shape ?. (Rank unknown.)
+    reverse_shape_test(PartialShape::dynamic(),
+                       {Shape{2, 3, 3},
+                        Shape{5, 3, 3},
+                        Shape{2, 5, 2},
+                        Shape{8, 1, 8},
+                        Shape{5},
+                        Shape{8, 2},
+                        Shape{8, 2, 8, 2},
+                        Shape{2, 3, 4, 5, 2}});
+}