[TOPI][Relay][OP] Add a strided_set operation.

python/tvm/relay/op/_transform.py

@@ -48,6 +48,7 @@
 _reg.register_schedule("cast_like", schedule_injective)
 _reg.register_schedule("reinterpret", schedule_injective)
 _reg.register_schedule("strided_slice", schedule_injective)
+_reg.register_schedule("strided_set", schedule_injective)
 _reg.register_schedule("slice_like", schedule_injective)
 _reg.register_schedule("split", schedule_injective)
 _reg.register_schedule("take", schedule_injective)
@@ -304,6 +305,28 @@ def compute_argwhere(attrs, inputs, output_type, _):
     new_output_type = tvm.relay.ty.TensorType(output_shape, "int32")
     return [topi.argwhere(new_output_type, inputs[0])]
 
+@_reg.register_compute("strided_set")
+def compute_strided_set(attrs, inputs, output_type, _):
+    """Compute definition of strided_set"""
+    begin = attrs.begin
+    end = attrs.end
+    strides = attrs.strides
+    n = len(inputs[0].shape)
+    strides = list(strides) + [1] * (n - len(strides))
+    lb = len(begin)
+    if lb < n:
+        begin = list(begin)
+        for i in range(lb, n):
+            begin.append(0 if strides[i] >= 0 else inputs.shape[i])
+    le = len(end)
+    if le < n:
+        end = list(end)
+        for i in range(le, n):
+            lim = inputs[0].shape[i] + 1
+            end.append(lim if strides[i] >= 0 else -lim)
+
+    return [topi.strided_set(inputs[0], inputs[1], begin, end, strides)]
+
 @script
 def _layout_transform_shape_func(data_shape,
                                  out_layout_len,

python/tvm/relay/op/transform.py

@@ -631,6 +631,36 @@ def strided_slice(data, begin, end, strides=None):
     return _make.strided_slice(data, list(begin), list(end), list(strides))
 
 
+def strided_set(data, v, begin, end, strides=None):
+    """Strided set of an array.
+
+    Parameters
+    ----------
+    data : relay.Expr
+        The source array to be sliced.
+
+    v : relay.Expr
+        The data to be set.
+
+    begin: list of int
+        The indices to begin with in the slicing.
+
+    end: list of int
+        Indices indicating end of the slice.
+
+    strides: list of int, optional
+        Specifies the stride values, it can be negative in that case,
+        the input tensor will be reversed in that particular axis.
+
+    Returns
+    -------
+    ret : relay.Expr
+        The computed result.
+    """
+    strides = strides or []
+    return _make.strided_set(data, v, list(begin), list(end), list(strides))
+
+
 def slice_like(data, shape_like, axes=None):
     """Slice the first input with respect to the second input.
 

src/relay/op/tensor/transform.cc

@@ -2049,6 +2049,56 @@ Examples::
 .set_attr<TOpPattern>("TOpPattern", kInjective)
 .set_attr<FInferCorrectLayout>("FInferCorrectLayout", StridedSliceInferCorrectLayout);
 
+// strided_set
+bool StridedSetRel(const Array<Type>& types,
+                   int num_inputs,
+                   const Attrs& attrs,
+                   const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 3);
+  reporter->Assign(types[2], types[0]);
+  return true;
+}
+
+Expr MakeStridedSet(Expr data,
+                    Expr v,
+                    Array<Integer> begin,
+                    Array<Integer> end,
+                    Array<Integer> strides) {
+  auto attrs = make_node<StridedSliceAttrs>();
+  attrs->begin = std::move(begin);
+  attrs->end = std::move(end);
+  attrs->strides = std::move(strides);
+  static const Op& op = Op::Get("strided_set");
+  return CallNode::make(op, {data, v}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op._make.strided_set")
+.set_body_typed(MakeStridedSet);
+
+
+RELAY_REGISTER_OP("strided_set")
+  .describe(R"code(Strided set of an array.
+Example::
+
+  x = [[  1.,   4.,   7.,  10.],
+       [  2.,   5.,   8.,  11.],
+       [  3.,   6.,   9.,  12.]]
+
+  v = [[ 11., 22., 33.]
+       [ 44., 55., 66.]]
+
+  strided_set(x, v, begin=[0, 1], end=[2, 4], stride=[1, 1]) = \
+      [[  1.,  11.,  22.,  33.],
+       [  2.,  44.,  55.,  66.],
+       [  3.,   6.,   9.,  12.]]
+)code" TVM_ADD_FILELINE)
+.set_num_inputs(2)
+.add_argument("data", "Tensor", "The input tensor.")
+.add_argument("v", "Tensor", "The data to set.")
+.set_support_level(4)
+.set_attrs_type<StridedSliceAttrs>()
+.set_attr<TOpPattern>("TOpPattern", kInjective)
+.add_type_rel("StridedSet", StridedSetRel);
 
 // relay.split
 TVM_REGISTER_NODE_TYPE(SplitAttrs);

tests/python/relay/test_op_level4.py

@@ -300,8 +300,45 @@ def verify(dshape, begin, end, strides, output, test_ref=True):
     verify((3, 4, 3), [1, 1], [4, 4, 3], None, (2, 3, 3))
 
 
+def test_strided_set():
+    def verify(dshape, begin, end, strides, vshape, test_ref=True):
+        x = relay.var("x", relay.TensorType(dshape, "float32"))
+        v = relay.var("v", relay.TensorType(vshape, "float32"))
+        z = relay.strided_set(x, v, begin=begin, end=end, strides=strides)
+        func = relay.Function([x, v], z)
+        func = run_infer_type(func)
+        text = func.astext()
+        assert "begin=" in text
+        assert "end=" in text
+        print(text)
+        assert func.body.checked_type == relay.ty.TensorType(dshape, "float32")
+        if not test_ref:
+            return
+        x_data = np.random.uniform(size=dshape).astype("float32")
+        v_data = np.random.uniform(size=vshape).astype("float32")
+        ref_res = topi.testing.strided_set_python(
+            x_data, v_data, begin, end, strides)
+        for target, ctx in ctx_list():
+            intrp = relay.create_executor("graph", ctx=ctx, target=target)
+            op_res = intrp.evaluate(func)(x_data, v_data)
+            tvm.testing.assert_allclose(op_res.asnumpy(), ref_res)
+
+    d1, d2, d3, d4 = tvm.var("d1"), tvm.var("d2"), tvm.var("d3"), tvm.var("d4")
+    verify((d1, d2, 3), [None, None, 1], [None, None, 2], None, (d1, d2, 1), False)
+    verify((3, 4, 3), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2))
+    verify((3, 4, 3), [1, 1, 0], [4, 4, 3], [2, 1, 1], (1, 3, 3))
+    verify((3, 4, 3), [1, -1, 0], [4, -5, 3], [2, -1, 1], (1, 4, 3))
+    verify((3, 4, 3), [1, 0, 0], [2, 2, 3], [1, 1, 2], (1, 2, 2))
+    verify((3, 4, 3), [1, -1, 0], [2, -3, 3], [1, -1, 1], (1, 2, 3))
+    verify((3, 4, 3), [1, 1, 0], [4, 4, 3], None, (2, 3, 3))
+    verify((3, 4, 3), [1, 1, 0], [4, 1000, 3], None, (2, 3, 3))
+    verify((3, 4, 3), [1, 1, 0], [4, 4], None, (2, 3, 3))
+    verify((3, 4, 3), [1, 1], [4, 4, 3], None, (2, 3, 3))
+
+
 if __name__ == "__main__":
     test_strided_slice()
+    test_strided_set()
     test_binary_op()
     test_cmp_type()
     test_binary_int_broadcast()

topi/python/topi/testing/__init__.py

@@ -20,7 +20,7 @@
 from .lrn_python import lrn_python
 from .l2_normalize_python import l2_normalize_python
 from .gather_nd_python import gather_nd_python
-from .strided_slice_python import strided_slice_python
+from .strided_slice_python import strided_slice_python, strided_set_python
 from .batch_matmul import batch_matmul
 from .slice_axis_python import slice_axis_python
 from .sequence_mask_python import sequence_mask

topi/python/topi/testing/strided_slice_python.py

@@ -14,7 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""gather_nd in python"""
+"""strided_slice/set in python"""
+
 
 def strided_slice_python(data, begin, end, strides):
     """Python version of strided slice operator.
@@ -46,3 +47,40 @@ def strided_slice_python(data, begin, end, strides):
             end[i] if i < len(end) else None,
             strides[i] if i < len(strides) else None))
     return data[tuple(slices)]
+
+
+def strided_set_python(data, v, begin, end, strides):
+    """Python version of strided slice operator.
+
+    Parameters
+    ----------
+    data : numpy.ndarray
+        Input data
+
+    v : numpy.ndarray
+        Value data
+
+    begin : list
+        Begining of the slices.
+
+    end : list
+        End of the slices.
+
+    strides : list
+        The stride of each slice.
+
+    Returns
+    -------
+    result : numpy.ndarray
+        The updated result.
+    """
+    strides = [] if strides is None else strides
+    slices = []
+    res = data.copy()
+    for i in range(len(data.shape)):
+        slices.append(slice(
+            begin[i] if i < len(begin) else None,
+            end[i] if i < len(end) else None,
+            strides[i] if i < len(strides) else None))
+    res[tuple(slices)] = v
+    return res

topi/python/topi/transform.py

@@ -20,6 +20,8 @@
 import tvm
 import topi
 from . import cpp
+from . import tag
+from .util import within_index, make_idx
 
 
 def expand_dims(a, axis, num_newaxis=1):
@@ -155,6 +157,74 @@ def strided_slice(a, begin, end, strides=None):
         strides = []
     return cpp.strided_slice(a, begin, end, strides)
 
+@tvm.tag_scope(tag=tag.INJECTIVE+",strided_set")
+def strided_set(a, v, begin, end, strides=None):
+    """Set slice of an array.
+
+    Parameters
+    ----------
+    a : tvm.Tensor
+        The tensor to be sliced.
+
+    v : tvm.Tensor
+        The values to set
+
+    begin: list of Expr
+        The indices to begin with in the slicing.
+
+    end: list of Expr
+        Indicies indicating end of the slice.
+
+    strides: list of Expr, optional
+        Specifies the stride values, it can be negative
+        in that case, the input tensor will be reversed
+        in that particular axis.
+
+    Returns
+    -------
+    ret : tvm.Tensor
+    """
+    n = len(a.shape)
+    if strides is None:
+        strides = [1] * n
+
+    if len(begin) != n:
+        raise ValueError("size mismatch")
+    if len(end) != n:
+        raise ValueError("size mismatch")
+    if len(strides) != n:
+        raise ValueError("size mismatch")
+
+    begin = list(map(tvm.convert, begin))
+    end = list(map(tvm.convert, end))
+    strides = list(map(tvm.convert, strides))
+
+    def _max(a, b):
+        return tvm.expr.Select(a > b, a, b)
+
+    # Convert negative indexes
+    for i in range(n):
+        begin[i] = tvm.if_then_else(begin[i] < 0,
+                                    begin[i] + a.shape[i],
+                                    begin[i])
+        end[i] = tvm.if_then_else(end[i] < 0,
+                                  end[i] + a.shape[i],
+                                  end[i])
+
+    def _select(*indices):
+        from_val = []
+        index_tuple = []
+        for i in range(n):
+            from_val.append(
+       	        within_index(begin[i], end[i], strides[i], indices[i]))
+            index_tuple.append(
+                make_idx(begin[i], end[i], strides[i], a.shape[i], indices[i]))
+        return tvm.if_then_else(tvm.all(*from_val),
+       	       	                v(*index_tuple),
+                                a(*indices))
+
+    return tvm.compute(a.shape, _select, name="strided_set")
+
 
 def reshape(a, newshape):
     """Reshape the array

topi/python/topi/util.py

@@ -345,3 +345,75 @@ def get_shape(src_shape, src_layout, dst_layout):
         tvm.convert([i for i in range(len(src_layout))]))
 
     return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices]))
+
+
+def within_index(b, e, s, i):
+    """Return a boolean value that indicates if i is within the given index.
+
+    Parameter
+    ---------
+    b : Expr
+      beginning of the index
+
+    e : Expr
+      end of the index
+
+    s : Expr
+      strides of index
+
+    i : Expr
+      array position
+
+    Returns
+    -------
+    selected: Expr
+        bool expression that is True is the array position would be selected
+        by the index and False otherwise
+    """
+    bc = tvm.expr.Select(s < 0, i <= e, i < b)
+    ec = tvm.expr.Select(s < 0, i > b, i >= e)
+    ss = tvm.if_then_else(s < 0,
+                          ((i - e) + (e % tvm.abs(s)) + 1) % tvm.abs(s),
+                          (i - b) % s)
+    return tvm.expr.Select(tvm.expr.Or(bc, ec), tvm.const(False), ss.equal(0))
+
+
+def make_idx(b, e, s, z, i):
+    """Return the array position in the selection that corresponds to an
+    array position in the full array.
+
+    The returned value is only meaningful if within_index() returns True
+    for the same set of parameters.
+
+    Parameter
+    ---------
+    b : Expr
+      beginning of the index
+
+    e : Expr
+      end of the index
+
+    s : Expr
+      strides of index
+
+    z : Expr
+      size of the indexed dimension
+
+    i : Expr
+      array position
+
+    Returns
+    -------
+    postion: Expr
+        int expression that corresponds to an array position in the selection.
+    """
+    bc = tvm.expr.Select(s < 0, i <= e, i < b)
+    ec = tvm.expr.Select(s < 0, i > b, i >= e)
+
+    # Clamp to array size
+    b = tvm.expr.Select(z < b, z - 1, b)
+
+    ss = tvm.if_then_else(s < 0,
+                          (b - i) // tvm.abs(s),
+                          (i - b) // s)
+    return tvm.if_then_else(tvm.expr.Or(bc, ec), 88, ss)