Dynamic ref transpose (#1438)

Extends ref transpose operator for dynamic shapes
Make dynamic tests more consistent naming

src/include/migraphx/op/transpose.hpp

@@ -29,6 +29,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -45,17 +46,15 @@ struct transpose
     }
 
     std::string name() const { return "transpose"; }
+
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
-        auto input         = inputs.at(0);
-        auto input_lens    = input.lens();
-        auto input_strides = input.strides();
-        auto t             = input.type();
+        check_shapes{inputs, *this, true}.has(1);
+        auto input = inputs.at(0);
 
-        if(dims.size() != input_lens.size())
+        if(dims.size() != input.ndim())
         {
-            MIGRAPHX_THROW("Permutation has wrong number of axes");
+            MIGRAPHX_THROW("TRANSPOSE: Permutation has wrong number of axes");
         }
         std::vector<int64_t> axes(dims.size());
         std::iota(axes.begin(), axes.end(), 0);
@@ -63,19 +62,36 @@ struct transpose
         {
             MIGRAPHX_THROW("TRANSPOSE: Invalid permutation");
         }
-        std::vector<size_t> output_lens(input_lens.size());
-        std::vector<size_t> output_strides(input_lens.size());
-        for(std::size_t i = 0; i < output_lens.size(); i++)
+
+        if(input.dynamic())
         {
-            output_lens[i]    = input_lens[dims[i]];
-            output_strides[i] = input_strides[dims[i]];
+            std::vector<shape::dynamic_dimension> output_dyn_dims(input.ndim());
+            std::transform(dims.cbegin(), dims.cend(), output_dyn_dims.begin(), [&](auto dim) {
+                return input.dyn_dims()[dim];
+            });
+            return {input.type(), output_dyn_dims};
+        }
+        else
+        {
+            auto input_lens    = input.lens();
+            auto input_strides = input.strides();
+
+            std::vector<size_t> output_lens(input.ndim());
+            std::vector<size_t> output_strides(input.ndim());
+            for(std::size_t i = 0; i < input.ndim(); i++)
+            {
+                output_lens[i]    = input_lens[dims[i]];
+                output_strides[i] = input_strides[dims[i]];
+            }
+            return {input.type(), output_lens, output_strides};
         }
-        return {t, output_lens, output_strides};
     }
-    argument compute(shape output_shape, std::vector<argument> args) const
+
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(output_shape);
+        return args[0].reshape(dyn_out.computed_shape);
     }
+
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 

src/onnx/parse_transpose.cpp

@@ -47,7 +47,7 @@ struct parse_transpose : op_parser<parse_transpose>
         }
 
         // if perm is empty, use the default value
-        auto n_dim = args.front()->get_shape().lens().size();
+        auto n_dim = args.front()->get_shape().ndim();
         if(perm.empty())
         {
             perm.resize(n_dim);

test/onnx/gen_onnx.py

@@ -6277,6 +6277,21 @@ def transpose_test():
     return ([node], [x], [y])
 
 
+@onnx_test
+def transpose_dyn_test():
+    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [None, 2, 2, 3])
+    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [None, 3, 2, 2])
+
+    node = onnx.helper.make_node(
+        'Transpose',
+        perm=[0, 3, 1, 2],
+        inputs=['0'],
+        outputs=['1'],
+    )
+
+    return ([node], [x], [y])
+
+
 @onnx_test
 def transpose_gather_test():
     x = helper.make_tensor_value_info('data', TensorProto.FLOAT, [3, 5, 4, 6])

test/onnx/onnx_test.cpp

@@ -5973,6 +5973,24 @@ TEST_CASE(transpose_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(transpose_dyn_test)
+{
+    migraphx::program p;
+    auto* mm   = p.get_main_module();
+    auto input = mm->add_parameter(
+        "0",
+        migraphx::shape{migraphx::shape::float_type, {{1, 4, 0}, {2, 2, 0}, {2, 2, 0}, {3, 3, 0}}});
+    std::vector<int64_t> perm{0, 3, 1, 2};
+    auto t0 = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", perm}}), input);
+    mm->add_return({t0});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {1, 4, 0};
+    auto prog                     = migraphx::parse_onnx("transpose_dyn_test.onnx", options);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(topk_attrk_test)
 {
     migraphx::program p;

test/op_shape_test.cpp

@@ -2226,6 +2226,28 @@ TEST_CASE(transpose_shape)
     throws_shape(migraphx::make_op("transpose", {{"permutation", {1, 2}}}), input);
 }
 
+TEST_CASE(transpose_dyn_shape0)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {2, 2, 0}}};
+    migraphx::shape output{migraphx::shape::float_type, {{2, 2, 0}, {1, 4, 0}}};
+    expect_shape(input, migraphx::make_op("transpose", {{"permutation", {0, 1}}}), input);
+    expect_shape(output, migraphx::make_op("transpose", {{"permutation", {1, 0}}}), input);
+}
+
+TEST_CASE(transpose_dyn_shape1)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4, 0}, {4, 4, 0}, {4, 4, 0}}};
+    migraphx::shape output{migraphx::shape::float_type, {{4, 4, 0}, {4, 4, 0}, {1, 4, 0}}};
+    expect_shape(input, migraphx::make_op("transpose", {{"permutation", {0, 1, 2}}}), input);
+    expect_shape(output, migraphx::make_op("transpose", {{"permutation", {2, 1, 0}}}), input);
+}
+
+TEST_CASE(transpose_axes_error)
+{
+    migraphx::shape input{migraphx::shape::float_type, {2, 2}};
+    throws_shape(migraphx::make_op("transpose", {{"permutation", {1}}}), input);
+}
+
 TEST_CASE(step_test)
 {
     migraphx::shape s1{migraphx::shape::float_type, {1, 2, 4}};