Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/AMDMIGraphX into pass_manager

src/include/migraphx/op/common.hpp

@@ -31,6 +31,8 @@ enum class rnn_direction
     bidirectional,
 };
 
+std::ostream& operator<<(std::ostream& os, rnn_direction v);
+
 } // namespace op
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/op/gather.hpp

@@ -19,6 +19,13 @@ namespace op {
 struct gather
 {
     int axis = 0;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.axis, "axis"));
+    }
+
     std::string name() const { return "gather"; }
 
     shape compute_shape(std::vector<shape> inputs) const

src/include/migraphx/op/gru.hpp

@@ -27,6 +27,16 @@ struct gru
     float clip              = 0.0f;
     int linear_before_reset = 0;
 
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.hidden_size, "hidden_size"),
+                    f(self.actv_funcs, "actv_func"),
+                    f(self.direction, "direction"),
+                    f(self.clip, "clip"),
+                    f(self.linear_before_reset, "linear_before_reset"));
+    }
+
     std::string name() const { return "gru"; }
     shape compute_shape(std::vector<shape> inputs) const
     {

src/include/migraphx/op/logsoftmax.hpp

@@ -19,6 +19,13 @@ namespace op {
 struct logsoftmax
 {
     int axis = 1;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.axis, "axis"));
+    }
+
     std::string name() const { return "logsoftmax"; }
     shape compute_shape(std::vector<shape> inputs) const
     {

src/include/migraphx/op/lstm.hpp

@@ -25,6 +25,15 @@ struct lstm
     float clip              = 0.0f;
     int input_forget        = 0;
 
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.hidden_size, "hidden_size"),
+                    f(self.actv_funcs, "actv_func"),
+                    f(self.direction, "direction"),
+                    f(self.input_forget, "input_forget"));
+    }
+
     std::string name() const { return "lstm"; }
     shape compute_shape(std::vector<shape> inputs) const
     {

src/include/migraphx/op/rnn.hpp

@@ -25,6 +25,15 @@ struct rnn
     rnn_direction direction = rnn_direction::forward;
     float clip              = 0.0f;
 
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.hidden_size, "hidden_size"),
+                    f(self.actv_funcs, "actv_func"),
+                    f(self.direction, "direction"),
+                    f(self.clip, "clip"));
+    }
+
     std::string name() const { return "rnn"; }
     shape compute_shape(std::vector<shape> inputs) const
     {

src/rewrite_rnn.cpp

@@ -4,6 +4,7 @@
 #include <migraphx/operators.hpp>
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/dfor.hpp>
+#include <migraphx/op/common.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -1166,5 +1167,14 @@ std::vector<operation> rewrite_rnn::lstm_actv_funcs(instruction_ref ins) const
     }
 }
 
+namespace op {
+std::ostream& operator<<(std::ostream& os, rnn_direction v)
+{
+    std::vector<std::string> rnn_direction_str = {"forward", "reverse", "bidirectional"};
+    os << rnn_direction_str[static_cast<std::underlying_type<rnn_direction>::type>(v)];
+    return os;
+}
+} // namespace op
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/tf/tf.cpp

@@ -363,6 +363,16 @@ struct tf_parser
         int64_t axis = 0;
         if(contains(attributes, "axis"))
             axis = attributes.at("axis").i();
+        size_t input_size = args.front()->get_shape().lens().size();
+        if(axis > input_size)
+        {
+            MIGRAPHX_THROW("TF_PARSER: axis value of " + to_string(axis) +
+                           " must be smaller than input size " + to_string(input_size));
+        }
+        // check if input arg needs axis to be converted to NCHW
+        if(input_size >= 4)
+            axis = parse_axis(axis);
+
         std::transform(
             args.begin(),
             args.end(),

test/tf/pack_test_nhwc.pb

+
+:
+
0Placeholder*
+dtype0
*
+shape:



+:
+
1Placeholder*
+dtype0
*
+shape:



+:
+
2Placeholder*
+dtype0
*
+shape:



+4
+pack1Pack
0
1
2*
+
T0
*
+
+axis*
+
N"
\ No newline at end of file

test/tf/tf_test.cpp

@@ -151,6 +151,28 @@ TEST_CASE(pack_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(pack_test_nhwc)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+    auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+    auto l2 = p.add_parameter("2", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+    std::vector<migraphx::instruction_ref> args{l0, l1, l2};
+    std::vector<migraphx::instruction_ref> unsqueezed_args;
+    int64_t nchw_axis = 1;
+
+    std::transform(args.begin(),
+                   args.end(),
+                   std::back_inserter(unsqueezed_args),
+                   [&](migraphx::instruction_ref arg) {
+                       return p.add_instruction(migraphx::op::unsqueeze{{nchw_axis}}, arg);
+                   });
+    p.add_instruction(migraphx::op::concat{static_cast<size_t>(nchw_axis)}, unsqueezed_args);
+    auto prog = migraphx::parse_tf("pack_test_nhwc.pb", true);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(pooling_test)
 {
     migraphx::program p;