SpaceToDepth operator (#979)

Inverse of DepthToSpace op
Co-authored-by: Shucai Xiao <shucai@gmail.com>

src/onnx/parse_spacetodepth.cpp

+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_spacetodepth : op_parser<parse_spacetodepth>
+{
+    std::vector<op_desc> operators() const { return {{"SpaceToDepth"}}; }
+
+    instruction_ref parse(const op_desc& /*opd*/,
+                          const onnx_parser& /*parser*/,
+                          const onnx_parser::node_info& info,
+                          std::vector<instruction_ref> args) const
+    {
+        auto s = args[0]->get_shape();
+        // blocksize attribute of SpaceToDepth
+        int blocksize = 1; // if blockSize of 1 then, this is a no-op
+        if(contains(info.attributes, "blocksize"))
+        {
+            blocksize = info.attributes.at("blocksize").i();
+        }
+        if(blocksize < 1)
+        {
+            // blockSize less than 1 would rather result in DepthToSpace instead of SpaceToDepth
+            MIGRAPHX_THROW("SpaceToDepth: blocksize is less than 1");
+        }
+        // calculate dimensions
+        auto res_lens = s.lens(); // {N, C, H, W}
+        if(((res_lens[2] % blocksize) == 0) and ((res_lens[3] % blocksize) == 0))
+        {
+            // Co = C * (blocksize ^ 2)
+            res_lens[1] = res_lens[1] * blocksize * blocksize;
+            // Ho = (H / blocksize)
+            res_lens[2] = res_lens[2] / blocksize;
+            // Wo = (W / blocksize)
+            res_lens[3] = res_lens[3] / blocksize;
+        } // res_shape = (N, Co, Ho, Wo)
+        else
+            MIGRAPHX_THROW("SpaceToDepth: div by blocksize quotient not int ");
+
+        auto trans_lens = s.lens(); // {N, C, H, W}
+        trans_lens[2]   = res_lens[2];
+        trans_lens[3]   = blocksize;
+        trans_lens.push_back(res_lens[3]);
+        trans_lens.push_back(blocksize); // {N, C, Ho, blocksize, Wo, blocksize}
+        std::vector<int64_t> perm = {0, 3, 5, 1, 2, 4};
+        auto temp1 = info.add_instruction(make_op("reshape", {{"dims", trans_lens}}), args[0]);
+        auto temp2 = info.add_instruction(make_op("transpose", {{"permutation", perm}}), temp1);
+        return info.add_instruction(make_op("reshape", {{"dims", res_lens}}),
+                                    info.make_contiguous(temp2));
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/onnx/gen_onnx.py

@@ -1061,6 +1061,62 @@ def depthtospace_crd_test():
     return ([node], [x], [y])
 
 
+@onnx_test
+def spacetodepth_test():
+
+    x = helper.make_tensor_value_info('x', TensorProto.float, [2, 2, 10, 10])
+    y = helper.make_tensor_value_info('y', TensorProto.float, [2, 8, 5, 5])
+
+    node = onnx.helper.make_node('spacetodepth',
+                                 inputs=['x'],
+                                 outputs=['y'],
+                                 blocksize=2)
+
+    return ([node], [x], [y])
+
+
+@onnx_test
+def spacetodepth_simple_test():
+
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 2, 4, 6])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 8, 2, 3])
+
+    node = onnx.helper.make_node('SpaceToDepth',
+                                 inputs=['x'],
+                                 outputs=['y'],
+                                 blocksize=2)
+
+    return ([node], [x], [y])
+
+
+@onnx_test
+def spacetodepth_invalid_blocksize_test():
+
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 2, 4, 6])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 8, 2, 3])
+
+    node = onnx.helper.make_node('SpaceToDepth',
+                                 inputs=['x'],
+                                 outputs=['y'],
+                                 blocksize=0.3)
+
+    return ([node], [x], [y])
+
+
+@onnx_test
+def spacetodepth_nondivisibility_test():
+
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 2, 5, 5])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 8, 2, 2])
+
+    node = onnx.helper.make_node('SpaceToDepth',
+                                 inputs=['x'],
+                                 outputs=['y'],
+                                 blocksize=2)
+
+    return ([node], [x], [y])
+
+
 @onnx_test
 def dequantizelinear_test():
     arg0 = helper.make_tensor_value_info('0', TensorProto.INT8, [5])

test/onnx/onnx_test.cpp

@@ -965,6 +965,46 @@ TEST_CASE(depthtospace_simple_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(spacetodepth_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0  = mm->add_parameter("x", {migraphx::shape::float_type, {2, 2, 10, 10}});
+    auto tmp1 =
+        mm->add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2, 5, 2, 5, 2}}}), l0);
+    auto tmp2 = mm->add_instruction(
+        migraphx::make_op("transpose", {{"permutation", {0, 3, 5, 1, 2, 4}}}), tmp1);
+    auto tmp3 = mm->add_instruction(migraphx::make_op("contiguous"), tmp2);
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", {2, 8, 5, 5}}}), tmp3);
+    auto prog = optimize_onnx("spacetodepth_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(spacetodepth_simple_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    auto l0  = mm->add_parameter("x", {migraphx::shape::float_type, {1, 2, 4, 6}});
+    auto tmp1 =
+        mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 2, 2, 3, 2}}}), l0);
+    auto tmp2 = mm->add_instruction(
+        migraphx::make_op("transpose", {{"permutation", {0, 3, 5, 1, 2, 4}}}), tmp1);
+    auto tmp3 = mm->add_instruction(migraphx::make_op("contiguous"), tmp2);
+    mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 8, 2, 3}}}), tmp3);
+    auto prog = optimize_onnx("spacetodepth_simple_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(spacetodepth_invalid_blocksize)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("spacetodepth_invalid_blocksize_test.onnx"); }));
+}
+
+TEST_CASE(spacetodepth_nondivisibility_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("spacetodepth_nondivisibility_test.onnx"); }));
+}
+
 TEST_CASE(dequantizelinear_test)
 {
     migraphx::program p;

test/onnx/spacetodepth_invalid_blocksize_test.onnx

+#spacetodepth_invalid_blocksize_test:

+)
+
x
y"SpaceToDepth*
+	blocksize>

#spacetodepth_invalid_blocksize_testZ
+
x
+

+

+
+
+b
+
y
+

+

+
+
+B
\ No newline at end of file

test/onnx/spacetodepth_nondivisibility_test.onnx

+!spacetodepth_nondivisibility_test:

+&
+
x
y"SpaceToDepth*
+	blocksize
!spacetodepth_nondivisibility_testZ
+
x
+

+

+
+
+b
+
y
+

+

+
+
+B
\ No newline at end of file

test/onnx/spacetodepth_simple_test.onnx

+spacetodepth_simple_test:|
+&
+
x
y"SpaceToDepth*
+	blocksize
spacetodepth_simple_testZ
+
x
+

+

+
+
+b
+
y
+

+

+
+
+B
\ No newline at end of file

test/onnx/spacetodepth_test.onnx

+spacetodepth_test:u
+&
+
x
y"SpaceToDepth*
+	blocksize
spacetodepth_testZ
+
x
+

+
+
+
+
+
+b
+
y
+

+
+
+
+B
\ No newline at end of file

test/onnx/verify_onnx.cpp

@@ -63,6 +63,46 @@ TEST_CASE(depthtospace_simple_test)
     EXPECT(migraphx::verify_range(result_vector, gold));
 }
 
+TEST_CASE(spacetodepth_simple_test)
+{
+    auto p = migraphx::parse_onnx("spacetodepth_simple_test.onnx");
+    p.compile(migraphx::ref::target{});
+    std::vector<float> data_in(48);
+    std::iota(std::begin(data_in), std::end(data_in), 0);
+    migraphx::shape s_x{migraphx::shape::float_type, {1, 2, 4, 6}};
+    migraphx::parameter_map pp;
+    pp["x"]     = migraphx::argument(s_x, data_in.data());
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold = {0,  2,  4,  12, 14, 16, 24, 26, 28, 36, 38, 40, 1,  3,  5,  13,
+                               15, 17, 25, 27, 29, 37, 39, 41, 6,  8,  10, 18, 20, 22, 30, 32,
+                               34, 42, 44, 46, 7,  9,  11, 19, 21, 23, 31, 33, 35, 43, 45, 47};
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(spacetodepth_depthtospace_test)
+{
+    // space to depth
+    auto p1 = migraphx::parse_onnx("spacetodepth_simple_test.onnx");
+    p1.compile(migraphx::ref::target{});
+    std::vector<float> data_in(48);
+    std::iota(std::begin(data_in), std::end(data_in), 0);
+    migraphx::shape s_x_1{migraphx::shape::float_type, {1, 2, 4, 6}};
+    migraphx::parameter_map pp1;
+    pp1["x"]     = migraphx::argument(s_x_1, data_in.data());
+    auto result1 = p1.eval(pp1).back();
+    // depth to space
+    auto p2 = migraphx::parse_onnx("depthtospace_simple_test.onnx");
+    p2.compile(migraphx::ref::target{});
+    migraphx::parameter_map pp2;
+    pp2["x"]     = result1;
+    auto result2 = p2.eval(pp2).back();
+    std::vector<float> result_vector2;
+    result2.visit([&](auto output) { result_vector2.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(result_vector2, data_in));
+}
+
 TEST_CASE(gather_elements)
 {
     migraphx::program p = migraphx::parse_onnx("gather_elements_axis0_test.onnx");