[Relay/Topi][Op] Added native DepthToSpace and SpaceToDepth Operators (#4566)

* Added tvm function stencil for subpixel operations to topi.

* Topi subpixel operators added and tested.

* Added subpixel attrs.

* Added depth_to_space relay attributes.

* depth_to_space fully working.

* Fixed NHWC shape bug.

* SpaceToDepth in and all tests passing.

* lint fixes.

* Added string include

* Fixed topi formatting.

* Added DCR/CDR mode to depthtospace operator.

include/tvm/relay/attrs/nn.h

@@ -821,6 +821,26 @@ struct DeformableConv2DAttrs : public tvm::AttrsNode<DeformableConv2DAttrs> {
   }
 };
 
+/*! \brief Attributes used in subpixel operators */
+struct SubPixelAttrs : public tvm::AttrsNode<SubPixelAttrs> {
+  int block_size;
+  std::string layout;
+  std::string mode;
+
+  TVM_DECLARE_ATTRS(SubPixelAttrs, "relay.attrs.SubPixelAttrs") {
+    TVM_ATTR_FIELD(block_size)
+        .describe("The size of subpixel blocks to compose or decompose.")
+        .set_default(1);
+    TVM_ATTR_FIELD(layout).set_default("NCHW").describe(
+        "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
+        "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+        "dimensions respectively.");
+    TVM_ATTR_FIELD(mode).set_default("DCR").describe(
+        "Indicates order in which channels are accessed. Must be one of"
+        "DCR or CDR.");
+  }
+};  // struct SubPixelAttrs
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_NN_H_

python/tvm/relay/frontend/onnx.py

@@ -540,34 +540,7 @@ class DepthToSpace(OnnxOpConverter):
 
         block_size = int(attr['blocksize'])
         mode = attr.get("mode", "DCR")
-
-        # handle NCHW layout
-        indata = infer_value_simulated(inputs[0], params)
-        in_n, in_c, in_h, in_w = indata.shape
-
-        # reshape to proper output
-        new_c = int(in_c / (block_size * block_size))
-        new_h = in_h * block_size
-        new_w = in_w * block_size
-        newshape = (in_n, new_c, new_h, new_w)
-
-        if mode == "DCR":
-            # expand input to larger dimension.
-            expanded = _op.reshape(inputs[0],
-                                   newshape=(in_n, block_size, block_size, new_c, in_h, in_w))
-            # reorder to expand spatial blocks.
-            transposed = _op.transpose(expanded, axes=(0, 3, 4, 1, 5, 2))
-
-        else:  # CRD mode
-            # expand input to larger dimension.
-            expanded = _op.reshape(inputs[0],
-                                   newshape=(in_n, new_c, block_size, block_size, in_h, in_w))
-            # reorder to expand spatial blocks.
-            transposed = _op.transpose(expanded, axes=(0, 1, 4, 2, 5, 3))
-
-        return AttrCvt(op_name="reshape",
-                       extras={'newshape': newshape},
-                       ignores=['mode', 'blocksize'])([transposed], attr)
+        return _op.nn.depth_to_space(inputs[0], block_size, mode=mode)
 
 
 class SpaceToDepth(OnnxOpConverter):
@@ -578,26 +551,7 @@ class SpaceToDepth(OnnxOpConverter):
     def _impl_v1(cls, inputs, attr, params):
 
         block_size = int(attr['blocksize'])
-
-        # handle NCHW layout
-        indata = infer_value_simulated(inputs[0], params)
-        in_n, in_c, in_h, in_w = indata.shape
-
-        # reshape to proper output
-        new_c = in_c * (block_size * block_size)
-        new_h = int(in_h / block_size)
-        new_w = int(in_w / block_size)
-        newshape = (in_n, new_c, new_h, new_w)
-
-        # expand input to larger dimension.
-        expanded = _op.reshape(inputs[0],
-                               newshape=(in_n, in_c, new_h, block_size, new_w, block_size))
-        # reorder to expand spatial blocks.
-        transposed = _op.transpose(expanded, axes=(0, 3, 5, 1, 2, 4))
-
-        return AttrCvt(op_name="reshape",
-                       extras={'newshape': newshape},
-                       ignores=['blocksize'])([transposed], attr)
+        return _op.nn.space_to_depth(inputs[0], block_size)
 
 
 class Concat(OnnxOpConverter):

python/tvm/relay/frontend/tensorflow.py

@@ -728,76 +728,18 @@ def _reshape():
 
 def _depth_to_space():
     def _impl(inputs, attr, params):
-        # Need to handle data layouts differently.
-        input_shape = attr['_input_shapes'][inputs[0]]
         block_size = int(attr['block_size'])
-        if attr['data_format'].decode("utf-8") == 'NHWC':
-            in_n, in_h, in_w, in_c = input_shape
-            new_c = int(in_c / (block_size * block_size))
-
-            # First expand input to larger dimension.
-            expanded = _op.reshape(
-                inputs[0], newshape=(in_n, in_h, in_w, block_size, block_size, new_c))
-            # Now reorder to expand spatial blocks.
-            transposed = _op.transpose(expanded, axes=(0, 1, 3, 2, 4, 5))
-            # Finally reshape to proper output.
-            new_h = in_h * block_size
-            new_w = in_w * block_size
-            newshape = (in_n, new_h, new_w, new_c)
-
-        else: # Handle NCHW layout
-            in_n, in_c, in_h, in_w = input_shape
-            new_c = int(in_c / (block_size * block_size))
-
-            expanded = _op.reshape(
-                inputs[0], newshape=(in_n, block_size, block_size, new_c, in_h, in_w))
-            transposed = _op.transpose(expanded, axes=(0, 3, 4, 1, 5, 2))
-            new_h = in_h * block_size
-            new_w = in_w * block_size
-            newshape = (in_n, new_c, new_h, new_w)
-
-        return AttrCvt(
-            op_name="reshape",
-            extras={'newshape': newshape},
-            ignores=['data_format', 'block_size'])([transposed], attr)
+        layout = attr['data_format'].decode("utf-8")
+        return _op.nn.depth_to_space(inputs[0], block_size, layout)
 
     return _impl
 
 
 def _space_to_depth():
     def _impl(inputs, attr, params):
-        # Need to handle data layouts differently.
-        input_shape = attr['_input_shapes'][inputs[0]]
         block_size = int(attr['block_size'])
-        if attr['data_format'].decode("utf-8") == 'NHWC':
-            in_n, in_h, in_w, in_c = input_shape
-            new_h = int(in_h / block_size)
-            new_w = int(in_w / block_size)
-
-            # First expand input to larger dimension.
-            expanded = _op.reshape(
-                inputs[0], newshape=(in_n, new_h, block_size, new_w, block_size, in_c))
-            # Now reorder to expand spatial blocks.
-            transposed = _op.transpose(expanded, axes=(0, 1, 3, 2, 4, 5))
-            # Finally reshape to proper output.
-            new_c = in_c * block_size * block_size
-            newshape = (in_n, new_h, new_w, new_c)
-
-        else:  # Handle NCHW layout
-            in_n, in_c, in_h, in_w = input_shape
-            new_h = int(in_h / block_size)
-            new_w = int(in_w / block_size)
-
-            expanded = _op.reshape(
-                inputs[0], newshape=(in_n, in_c, new_h, block_size, new_w, block_size))
-            transposed = _op.transpose(expanded, axes=(0, 3, 5, 1, 2, 4))
-            new_c = int(in_c * block_size * block_size)
-            newshape = (in_n, new_c, new_h, new_w)
-
-        return AttrCvt(
-            op_name="reshape",
-            extras={'newshape': newshape},
-            ignores=['data_format', 'block_size'])([transposed], attr)
+        layout = attr['data_format'].decode("utf-8")
+        return _op.nn.space_to_depth(inputs[0], block_size, layout)
 
     return _impl
 

python/tvm/relay/op/nn/_nn.py

@@ -904,6 +904,28 @@ def compute_cross_entropy_with_logits(attrs, inputs, out_dtype, target):
     x, y = inputs
     return [-topi.sum(x * y) / x.shape[0]]
 
+
+@reg.register_compute("nn.depth_to_space")
+def compute_depth_to_space(attrs, inputs, out_dtype, target):
+    block_size = attrs.block_size
+    layout = attrs.layout
+    mode = attrs.mode
+    return [topi.nn.depth_to_space(inputs[0], block_size, layout=layout, mode=mode)]
+
+reg.register_schedule("nn.depth_to_space", schedule_injective)
+reg.register_pattern("nn.depth_to_space", OpPattern.INJECTIVE)
+
+
+@reg.register_compute("nn.space_to_depth")
+def compute_space_to_depth(attrs, inputs, out_dtype, target):
+    block_size = attrs.block_size
+    layout = attrs.layout
+    return [topi.nn.space_to_depth(inputs[0], block_size, layout=layout)]
+
+reg.register_schedule("nn.space_to_depth", schedule_injective)
+reg.register_pattern("nn.space_to_depth", OpPattern.INJECTIVE)
+
+
 # shape func
 @script
 def _conv2d_NCHWc_shape_func(dshape, kshape, strides, padding, dilation, oc_bn):

python/tvm/relay/op/nn/nn.py

@@ -2079,3 +2079,53 @@ def cross_entropy_with_logits(predictions, targets):
       The computed result.
     """
     return _make.cross_entropy_with_logits(predictions, targets)
+
+
+def depth_to_space(data, block_size, layout='NCHW', mode='DCR'):
+    """Convert channels into spatial blocks.
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        Input data with channels divisible by block_size**2
+
+    block_size : int
+        Size of blocks to convert channels into.
+
+    layout : string
+        One of NCHW or NHWC, indicates channel axis.
+
+    mode : string
+        One of DCR or CDR, indicates which order channels
+        are accessed in.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        Tensor with shape [in_batch, in_channel / block_size * block_size,
+                           in_height * block_size, in_width * block_size]
+    """
+    return _make.depth_to_space(data, block_size, layout, mode)
+
+
+def space_to_depth(data, block_size, layout='NCHW'):
+    """Convert spatial blocks into channels.
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        Input data with spatial dimensions divisible by block_size
+
+    block_size : int
+        Size of blocks to decompose into channels.
+
+    layout : string
+        One of NCHW or NHWC, indicates channel axis.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        Tensor with shape [in_batch, in_channel * block_size * block_size,
+                           in_height / block_size, in_width / block_size]
+    """
+    return _make.space_to_depth(data, block_size, layout)

python/tvm/relay/op/op_attrs.py

@@ -299,3 +299,8 @@ class BinaryDenseAttrs(Attrs):
 @register_relay_attr_node
 class Conv2DTransposeAttrs(Attrs):
     """Attributes used in Transposed Conv2D operators"""
+
+
+@register_relay_attr_node
+class SubPixelAttrs(Attrs):
+    """Attributes used in depth to space and space to depth operators"""

src/relay/op/nn/nn.cc

@@ -31,6 +31,7 @@
 #include <topi/nn/softmax.h>
 #include <topi/nn/flatten.h>
 #include <vector>
+#include <string>
 #include "../type_relations.h"
 #include "../../pass/alter_op_layout.h"
 #include "../op_common.h"
@@ -959,6 +960,123 @@ Accept logits.
 .set_support_level(10)
 .add_type_rel("CrossEntropy", CrossEntropyRel);
 
+// Depth to space and space to depth
+TVM_REGISTER_NODE_TYPE(SubPixelAttrs);
+
+bool DepthToSpaceRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                     const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  static const Layout kNCHW("NCHW");
+
+  const SubPixelAttrs* param = attrs.as<SubPixelAttrs>();
+  CHECK(param != nullptr);
+  const int block_size = param->block_size;
+  const Layout in_layout(param->layout);
+  auto layout_converter = BijectiveLayoutNode::make(in_layout, kNCHW);
+  CHECK(layout_converter.defined())
+      << "DepthToSpace only support input layouts that are convertible from NCHW."
+      << " But got " << in_layout;
+
+  auto oshape = layout_converter.ForwardShape(data->shape);
+  oshape.Set(1, indexdiv(oshape[1], (block_size * block_size)));
+  oshape.Set(2, oshape[2] * block_size);
+  oshape.Set(3, oshape[3] * block_size);
+
+  // Assign output type
+  reporter->Assign(types[1],
+                   TensorTypeNode::make(layout_converter.BackwardShape(oshape), data->dtype));
+
+  return true;
+}
+
+// Positional relay function to create DepthToSpace operator
+// used by frontend FFI
+Expr MakeDepthToSpace(Expr data, int block_size, std::string layout, std::string mode) {
+  auto attrs = make_node<SubPixelAttrs>();
+  attrs->block_size = block_size;
+  attrs->layout = std::move(layout);
+  attrs->mode = std::move(mode);
+  static const Op& op = Op::Get("nn.depth_to_space");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.nn._make.depth_to_space").set_body_typed(MakeDepthToSpace);
+
+RELAY_REGISTER_OP("nn.depth_to_space")
+    .describe(R"code(Rearrange input channels into spatial pixels.
+
+- **data**: data is a 4D array of shape
+            (batch, in_channels, in_height, in_width) for NCHW
+
+- **out**: Output is a 4D array of shape
+           (batch, in_channels / block_size * block_size, in_height * block_size, in_width * block_size) for NCHW.
+
+)code" TVM_ADD_FILELINE)
+    .set_attrs_type<SubPixelAttrs>()
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "The input tensor")
+    .set_support_level(5)
+    .add_type_rel("DepthToSpace", DepthToSpaceRel);
+
+bool SpaceToDepthRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                     const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  static const Layout kNCHW("NCHW");
+
+  const SubPixelAttrs* param = attrs.as<SubPixelAttrs>();
+  CHECK(param != nullptr);
+  const int block_size = param->block_size;
+  const Layout in_layout(param->layout);
+  auto layout_converter = BijectiveLayoutNode::make(in_layout, kNCHW);
+  CHECK(layout_converter.defined())
+      << "SpaceToDepth only support input layouts that are convertible from NCHW."
+      << " But got " << in_layout;
+
+  auto oshape = layout_converter.ForwardShape(data->shape);
+  oshape.Set(1, oshape[1] * (block_size * block_size));
+  oshape.Set(2, indexdiv(oshape[2], block_size));
+  oshape.Set(3, indexdiv(oshape[3], block_size));
+
+  // Assign output type
+  reporter->Assign(types[1],
+                   TensorTypeNode::make(layout_converter.BackwardShape(oshape), data->dtype));
+
+  return true;
+}
+
+// Positional relay function to create SpaceToDepth operator
+// used by frontend FFI
+Expr MakeSpaceToDepth(Expr data, int block_size, std::string layout) {
+  auto attrs = make_node<SubPixelAttrs>();
+  attrs->block_size = block_size;
+  attrs->layout = std::move(layout);
+  static const Op& op = Op::Get("nn.space_to_depth");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.nn._make.space_to_depth").set_body_typed(MakeSpaceToDepth);
+
+RELAY_REGISTER_OP("nn.space_to_depth")
+    .describe(R"code(Rearrange spatial pixels into new output channels.
+
+- **data**: data is a 4D array of shape
+            (batch, in_channels, in_height, in_width) for NCHW
+
+- **out**: Output is a 4D array of shape
+           (batch, in_channels * block_size * block_size, in_height / block_size, in_width / block_size) for NCHW.
+
+)code" TVM_ADD_FILELINE)
+    .set_attrs_type<SubPixelAttrs>()
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "The input tensor")
+    .set_support_level(5)
+    .add_type_rel("SpaceToDepth", SpaceToDepthRel);
 
 }  // namespace relay
 }  // namespace tvm

tests/python/relay/test_op_level5.py

@@ -573,6 +573,69 @@ def test_deformable_conv2d():
     test_run(2, 4, 16, 4, 4, 1)
 
 
+def test_depth_to_space():
+    def verify_depth_to_space(dshape, block_size, layout, mode):
+        if layout == "NHWC":
+            out_shape = [dshape[0], dshape[1] * block_size, dshape[2] * block_size, dshape[3] / (block_size * block_size)]
+        else:
+            out_shape = [dshape[0], dshape[1] / (block_size * block_size), dshape[2] * block_size, dshape[3] * block_size]
+
+        x_data = np.random.uniform(size=dshape).astype("float32")
+        if layout == "NHWC":
+            x_data = np.transpose(x_data, axes=[0, 3, 1, 2])
+        ref_res = topi.testing.depth_to_space_python(x_data, block_size, mode=mode)
+        if layout == "NHWC":
+            x_data = np.transpose(x_data, axes=[0, 2, 3, 1])
+            ref_res = np.transpose(ref_res, axes=[0, 2, 3, 1])
+
+        x = relay.var("x", relay.TensorType(dshape, "float32"))
+        z = relay.nn.depth_to_space(x, block_size, layout, mode)
+        assert "block_size=" in z.astext()
+        zz = run_infer_type(z)
+        assert zz.checked_type == relay.TensorType(ref_res.shape, "float32")
+        func = relay.Function([x], z)
+
+        for target, ctx in ctx_list():
+            for kind in ["graph", "debug"]:
+                intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                op_res = intrp.evaluate(func)(x_data)
+                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-4)
+    for layout in ["NHWC", "NCHW"]:
+        for mode in ["DCR", "CDR"]:
+            verify_depth_to_space((1, 4, 4, 4), 2, layout, mode)
+
+
+def test_space_to_depth():
+    def verify_space_to_depth(dshape, block_size, layout):
+        if layout == "NHWC":
+            out_shape = [dshape[0], dshape[1] / block_size, dshape[2] / block_size, dshape[3] * (block_size * block_size)]
+        else:
+            out_shape = [dshape[0], dshape[1] * (block_size * block_size), dshape[2] / block_size, dshape[3] / block_size]
+
+        x_data = np.random.uniform(size=dshape).astype("float32")
+        if layout == "NHWC":
+            x_data = np.transpose(x_data, axes=[0, 3, 1, 2])
+        ref_res = topi.testing.space_to_depth_python(x_data, block_size)
+        if layout == "NHWC":
+            x_data = np.transpose(x_data, axes=[0, 2, 3, 1])
+            ref_res = np.transpose(ref_res, axes=[0, 2, 3, 1])
+
+        x = relay.var("x", relay.TensorType(dshape, "float32"))
+        z = relay.nn.space_to_depth(x, block_size, layout)
+        assert "block_size=" in z.astext()
+        zz = run_infer_type(z)
+        assert zz.checked_type == relay.TensorType(ref_res.shape, "float32")
+        func = relay.Function([x], z)
+
+        for target, ctx in ctx_list():
+            for kind in ["graph", "debug"]:
+                intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                op_res = intrp.evaluate(func)(x_data)
+                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-4)
+    for layout in ["NHWC", "NCHW"]:
+        verify_space_to_depth((1, 4, 4, 4), 2, layout)
+
+
 if __name__ == "__main__":
     test_resize_infer_type()
     test_resize()
@@ -586,3 +649,5 @@ if __name__ == "__main__":
     test_yolo_reorg()
     test_non_max_suppression()
     test_deformable_conv2d()
+    test_depth_to_space()
+    test_space_to_depth()
\ No newline at end of file

topi/python/topi/nn/__init__.py

@@ -42,3 +42,5 @@ from .batch_matmul import *
 from .sparse import *
 from .pad import *
 from .fifo_buffer import *
+from .depth_to_space import *
+from .space_to_depth import *

topi/python/topi/nn/depth_to_space.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""TVM operator depth_to_space compute."""
+from __future__ import absolute_import
+import tvm
+from .. import tag
+
+
+def depth_to_space(data, block_size, layout='NCHW', mode='DCR'):
+    """Perform depth to space transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        4-D tensor in either NCHW or NHWC layout.
+
+    block_size : int
+        Size of blocks to compose from channel dimension.
+
+    layout : string
+        Either NCHW or NHWC, indicating data layout.
+
+    mode : string
+        Either DCR or CDR, indicates how channels should be accessed.
+        In DCR, channels are interwoven in the Tensorflow style while
+        in CDR channels are accessed sequentially as in Pytorch.
+
+    Returns
+    -------
+    output : tvm.Tensor
+        Output of shape [N, C / block_size**2, H * block_size, W * block_size]
+    """
+    if layout == 'NCHW':
+        in_n, in_c, in_h, in_w = data.shape
+        channel_factor = tvm.truncdiv(in_c, (block_size * block_size))
+        output_shape = [in_n, channel_factor,
+                        in_h * block_size, in_w * block_size]
+    elif layout == 'NHWC':
+        in_n, in_h, in_w, in_c = data.shape
+        channel_factor = tvm.truncdiv(in_c, (block_size * block_size))
+        output_shape = [in_n, in_h * block_size,
+                        in_w * block_size, channel_factor]
+    else:
+        raise ValueError("Only NCHW and NHWC layouts are currently supported.")
+
+    def _get_indices(*indices):
+        if layout == 'NCHW':
+            n, c, y, x = indices
+        elif layout == 'NHWC':
+            n, y, x, c = indices
+        return n, c, y, x
+
+    def _get_pixel(n, c, y, x):
+        block_x = tvm.truncdiv(x, block_size)
+        block_y = tvm.truncdiv(y, block_size)
+        idx_x = tvm.truncmod(x, block_size)
+        idx_y = tvm.truncmod(y, block_size)
+        if mode == "DCR":
+            channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c
+        else:
+            channel_idx = (c * block_size * block_size) + ((block_size * idx_y) + idx_x)
+
+        if layout == 'NCHW':
+            output = data(n, channel_idx, block_y, block_x)
+        else:
+            output = data(n, block_y, block_x, channel_idx)
+        return output
+
+    def _compute(*indices):
+        n, c, y, x = _get_indices(*indices)
+        return _get_pixel(n, c, y, x)
+
+    return tvm.compute(output_shape, _compute, name='depth_to_space', tag=tag.INJECTIVE)

topi/python/topi/nn/space_to_depth.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""TVM operator space_to_depth compute."""
+from __future__ import absolute_import
+import tvm
+from .. import tag
+
+
+def space_to_depth(data, block_size, layout='NCHW'):
+    """Perform space to depth transformation on the data
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        4-D tensor in either NCHW or NHWC layout.
+
+    block_size : int
+        Size of blocks to decompose into channel dimension.
+
+    layout : string
+        Either NCHW or NHWC, indicating data layout.
+
+    Returns
+    -------
+    output : tvm.Tensor
+        Output of shape [N, C * block_size**2, H / block_size, W / block_size]
+    """
+
+    if layout == 'NCHW':
+        in_n, in_c, in_h, in_w = data.shape
+        output_shape = [in_n, in_c * block_size * block_size,
+                        tvm.truncdiv(in_h, block_size), tvm.truncdiv(in_w, block_size)]
+    elif layout == 'NHWC':
+        in_n, in_h, in_w, in_c = data.shape
+        output_shape = [in_n, tvm.truncdiv(in_h, block_size), tvm.truncdiv(
+            in_w, block_size), in_c * block_size * block_size]
+    else:
+        raise ValueError("Only NCHW and NHWC layouts are currently supported.")
+
+    def _get_indices(*indices):
+        if layout == 'NCHW':
+            n, c, y, x = indices
+        elif layout == 'NHWC':
+            n, y, x, c = indices
+        return n, c, y, x
+
+    def _get_pixel(n, c, y, x):
+        block_offset = tvm.truncdiv(c, in_c)
+        channel_idx = tvm.truncmod(c, in_c)
+        x_idx = tvm.truncmod(block_offset, block_size)
+        y_idx = tvm.truncdiv(block_offset, block_size)
+
+        if layout == 'NCHW':
+            output = data(n, channel_idx, y_idx +
+                          (y * block_size), x_idx + (x * block_size))
+        else:
+            output = data(n, y_idx + (y * block_size), x_idx +
+                          (x * block_size), channel_idx)
+        return output
+
+    def _compute(*indices):
+        n, c, y, x = _get_indices(*indices)
+        return _get_pixel(n, c, y, x)
+
+    return tvm.compute(output_shape, _compute, name='space_to_depth', tag=tag.INJECTIVE)

topi/python/topi/testing/__init__.py

@@ -46,3 +46,5 @@ from .sequence_mask_python import sequence_mask
 from .pool3d_python import pool3d_ncdhw_python
 from .pool_grad_python import pool_grad_nchw
 from .one_hot import one_hot
+from .depth_to_space import depth_to_space_python
+from .space_to_depth import space_to_depth_python

topi/python/topi/testing/depth_to_space.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals
+"""Depth to space in python"""
+import numpy as np
+
+
+def depth_to_space_python(data, block_size, mode='DCR'):
+    """Depth to Space operator in python for NCHW layout.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    block_size : int
+        Size of blocks to convert channel pixels into.
+
+    Returns
+    -------
+    d2s_out : np.ndarray
+        4-D with shape [batch, in_channel / (block_size * block_size),
+                        out_height * block_size, out_width * block_size]
+    """
+    in_n, in_c, in_h, in_w = data.shape
+    new_h = int(in_h * block_size)
+    new_w = int(in_h * block_size)
+    new_c = int(in_c / (block_size * block_size))
+
+    if mode == 'DCR':
+        expanded = np.reshape(
+            data, newshape=[in_n, block_size, block_size, new_c, in_h, in_w])
+        transposed = np.transpose(expanded, axes=[0, 3, 4, 1, 5, 2])
+    else:
+        expanded = np.reshape(
+            data, newshape=(in_n, new_c, block_size, block_size, in_h, in_w))
+        transposed = np.transpose(expanded, axes=(0, 1, 4, 2, 5, 3))
+    newshape = [in_n, new_c, new_h, new_w]
+    d2s_out = np.reshape(transposed, newshape=newshape)
+    return d2s_out

topi/python/topi/testing/space_to_depth.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals
+"""Space to depth in python"""
+import numpy as np
+
+
+def space_to_depth_python(data, block_size):
+    """Space to Depth operator in python for NCHW layout.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        4-D with shape [batch, in_channel, in_height, in_width]
+
+    block_size : int
+        Size of spatial blocks to decompose into channels.
+
+    Returns
+    -------
+    d2s_out : np.ndarray
+        4-D with shape [batch, in_channel * (block_size * block_size),
+                        out_height / block_size, out_width / block_size]
+    """
+    in_n, in_c, in_h, in_w = data.shape
+    new_h = int(in_h / block_size)
+    new_w = int(in_h / block_size)
+    new_c = int(in_c * (block_size * block_size))
+
+    expanded = np.reshape(
+        data, newshape=[in_n, in_c, new_h, block_size, new_w, block_size])
+    transposed = np.transpose(expanded, axes=[0, 3, 5, 1, 2, 4])
+    newshape = [in_n, new_c, new_h, new_w]
+    d2s_out = np.reshape(transposed, newshape=newshape)
+    return d2s_out

topi/tests/python/test_topi_depth_to_space.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test code for depth to space"""
+import numpy as np
+import tvm
+import topi
+import topi.testing
+
+from common import get_all_backend
+
+
+def verify_depth_to_space(block_size, batch, in_channel, in_height, in_width, layout='NCHW', mode='DCR'):
+    out_channel = int(in_channel / (block_size * block_size))
+    out_height = int(in_height * block_size)
+    out_width = int(in_width * block_size)
+
+    if layout == 'NCHW':
+        in_shape = [batch, in_channel, in_height, in_width]
+        out_shape = [batch, out_channel, out_height, out_width]
+    elif layout == 'NHWC':
+        in_shape = [batch, in_height, in_width, in_channel]
+        out_shape = [batch, out_height, out_width, out_channel]
+    else:
+        raise NotImplementedError('Layout not supported {}'.format(layout))
+
+    A = tvm.placeholder(in_shape, name='A', dtype='float32')
+    dtype = A.dtype
+    a_np = np.random.uniform(size=in_shape).astype(dtype)
+
+    B = topi.nn.depth_to_space(A, block_size=block_size, layout=layout, mode=mode)
+    if layout == 'NHWC':
+        a_np = np.transpose(a_np, axes=[0, 3, 1, 2])
+    b_np = topi.testing.depth_to_space_python(a_np, block_size, mode=mode)
+    if layout == 'NHWC':
+        a_np = np.transpose(a_np, axes=[0, 2, 3, 1])
+        b_np = np.transpose(b_np, axes=[0, 2, 3, 1])
+
+    def check_device(device):
+        ctx = tvm.context(device, 0)
+        if not ctx.exist:
+            print("Skip because %s is not enabled" % device)
+            return
+        print("Running on target: %s" % device)
+        with tvm.target.create(device):
+            s = topi.generic.schedule_injective(B)
+        a = tvm.nd.array(a_np, ctx)
+        b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), ctx)
+        f = tvm.build(s, [A, B], device)
+        f(a, b)
+        tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-3, atol=1e-3)
+
+    for device in get_all_backend():
+        check_device(device)
+
+
+def test_depth_to_space():
+    for layout in ['NCHW', 'NHWC']:
+        for mode in ['DCR', 'CDR']:
+            # Simplest possible case
+            verify_depth_to_space(2, 1, 4, 1, 1, layout=layout, mode=mode)
+            # Average input size
+            verify_depth_to_space(2, 1, 32, 32, 32, layout=layout, mode=mode)
+            # Large block size
+            verify_depth_to_space(8, 1, 256, 32, 32, layout=layout, mode=mode)
+            # Large batch size
+            verify_depth_to_space(4, 8, 32, 32, 32, layout=layout, mode=mode)
+            # Large input size
+            verify_depth_to_space(4, 8, 32, 128, 128, layout=layout, mode=mode)
+
+
+if __name__ == "__main__":
+    test_depth_to_space()

topi/tests/python/test_topi_space_to_depth.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test code for space to depth"""
+import numpy as np
+import tvm
+import topi
+import topi.testing
+
+from common import get_all_backend
+
+
+def verify_space_to_depth(block_size, batch, in_channel, in_height, in_width, layout='NCHW'):
+    out_channel = int(in_channel * (block_size * block_size))
+    out_height = int(in_height / block_size)
+    out_width = int(in_width / block_size)
+
+    if layout == 'NCHW':
+        in_shape = [batch, in_channel, in_height, in_width]
+        out_shape = [batch, out_channel, out_height, out_width]
+    elif layout == 'NHWC':
+        in_shape = [batch, in_height, in_width, in_channel]
+        out_shape = [batch, out_height, out_width, out_channel]
+    else:
+        raise NotImplementedError('Layout not supported {}'.format(layout))
+
+    A = tvm.placeholder(in_shape, name='A', dtype='float32')
+    dtype = A.dtype
+    a_np = np.random.uniform(size=in_shape).astype(dtype)
+
+    B = topi.nn.space_to_depth(A, block_size=block_size, layout=layout)
+    if layout == 'NHWC':
+        a_np = np.transpose(a_np, axes=[0, 3, 1, 2])
+    b_np = topi.testing.space_to_depth_python(a_np, block_size)
+    if layout == 'NHWC':
+        a_np = np.transpose(a_np, axes=[0, 2, 3, 1])
+        b_np = np.transpose(b_np, axes=[0, 2, 3, 1])
+
+    def check_device(device):
+        ctx = tvm.context(device, 0)
+        if not ctx.exist:
+            print("Skip because %s is not enabled" % device)
+            return
+        print("Running on target: %s" % device)
+        with tvm.target.create(device):
+            s = topi.generic.schedule_injective(B)
+        a = tvm.nd.array(a_np, ctx)
+        b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), ctx)
+        f = tvm.build(s, [A, B], device)
+        f(a, b)
+        tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-3, atol=1e-3)
+
+    for device in get_all_backend():
+        check_device(device)
+
+
+def test_space_to_depth():
+    for layout in ['NCHW', 'NHWC']:
+        # Simplest possible case
+        verify_space_to_depth(2, 1, 1, 2, 2, layout=layout)
+        # Average input size
+        verify_space_to_depth(2, 1, 32, 32, 32, layout=layout)
+        # Large block size
+        verify_space_to_depth(8, 1, 32, 64, 64, layout=layout)
+        # Large batch size
+        verify_space_to_depth(4, 8, 32, 32, 32, layout=layout)
+        # Large input size
+        verify_space_to_depth(4, 8, 32, 128, 128, layout=layout)
+
+
+if __name__ == "__main__":
+    test_space_to_depth()