SparseFillEmptyRows Op (apache#7442)

* Initial Commit

* Fix formats

* Remove comments

* Black

* THreeops

* Add Frontend Code

* Add Default Value to feed dict

* Add Frontend Code

* New test Cases and new code to handle them

* Add Python Implementation'
'

* Remove stuff

* Remove unused imports

* Pylint

* Pylint

* PyLint Shape Func

* Make tests cpu only

* Add unsorted tests

* Add frontend code

* Row Major Sorting Only Test

* Handle Dynamic Shapes

* Add dynamic input shapes

* Dynamic Shape Tests

* Add documentation

* Dtypes

* PR Comments

* Added comments and changed naming

* Add comments

* Comments to Shape Func

* Documentation

* PR Changes

* PR Comments

* Resolve input and output dtype compat

Co-authored-by: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>

python/tvm/relay/frontend/tensorflow.py

@@ -998,6 +998,28 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _sparse_fill_empty_rows():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 4, "There should be 4 input tensors"
+        sparse_indices = inputs[0]
+        sparse_values = inputs[1]
+        sparse_indices_num_cols = _infer_shape(sparse_indices, mod)[1]
+        first_column = _op.split(sparse_indices, sparse_indices_num_cols, axis=1)[0]
+        sorted_indices = _op.argsort(_op.squeeze(first_column))
+        sorted_sparse_indices = _op.take(sparse_indices, sorted_indices, axis=0)
+        sorted_sparse_values = _op.take(sparse_values, sorted_indices, axis=0)
+        new_sparse_indices, new_sparse_values, empty_row_indicator = _op.sparse_fill_empty_rows(
+            sorted_sparse_indices, sorted_sparse_values, inputs[2], inputs[3]
+        )
+
+        return _expr.TupleWrapper(
+            _expr.Tuple([new_sparse_indices, new_sparse_values, empty_row_indicator]),
+            3,
+        )
+
+    return _impl
+
+
 def _identity():
     def _impl(inputs, attr, params, mod):
         return inputs[0]
@@ -2455,6 +2477,7 @@ def _impl(inputs, attr, params, mod):
     "SpaceToDepth": _space_to_depth(),
     "SparseToDense": _sparse_to_dense(),
     "SparseTensorDenseMatMul": _sparse_tensor_dense_matmul(),
+    "SparseFillEmptyRows": _sparse_fill_empty_rows(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),

python/tvm/relay/op/_transform.py

@@ -15,7 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 """Backend compiler related feature registration"""
-# pylint: disable=invalid-name,unused-argument, len-as-condition, too-many-nested-blocks, too-many-local-variables, too-many-arguments
+# pylint: disable=invalid-name,unused-argument, len-as-condition, too-many-nested-blocks,
+# pylint: disable=too-many-local-variables, too-many-arguments, no-else-return
+
 from __future__ import absolute_import
 import tvm
 from tvm import te
@@ -94,6 +96,24 @@ def compute_scatter(attrs, inputs, output_type):
 
 _reg.register_strategy("scatter", strategy.scatter_strategy)
 
+# sparse_fill_empty_rows
+@_reg.register_compute("sparse_fill_empty_rows")
+def compute_sparse_fill_empty_rows(attrs, inputs, output_type):
+    """Compute definition of sparse_fill_empty_rows"""
+
+    return topi.sparse_fill_empty_rows(
+        inputs[0],
+        inputs[1],
+        inputs[2],
+        inputs[3],
+        output_type.fields[0].shape,
+        output_type.fields[1].shape,
+        output_type.fields[2].shape,
+    )
+
+
+_reg.register_strategy("sparse_fill_empty_rows", strategy.sparse_fill_empty_rows_strategy)
+
 # scatter_add
 @_reg.register_compute("scatter_add")
 def compute_scatter_add(attrs, inputs, output_type):
@@ -456,6 +476,47 @@ def argwhere_shape_func(attrs, inputs, out_ndims):
 _reg.register_shape_func("interpolate", False, elemwise_shape_func)
 
 
+@script
+def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape):
+
+    new_sparse_indices_shape = output_tensor((2,), "int64")
+    new_sparse_values_shape = output_tensor((1,), "int64")
+    empty_row_indicator_shape = output_tensor((1,), "int64")
+    num_dense_rows = int64(dense_shape[0])
+
+    if int64(sparse_indices.shape[0]) == int64(0):  # Handle Empty Case
+        #  Total rows will equal dense_shape[0]
+        new_sparse_indices_shape[0] = num_dense_rows
+        new_sparse_indices_shape[1] = int64(sparse_indices.shape[1])
+        new_sparse_values_shape[0] = num_dense_rows
+        empty_row_indicator_shape[0] = num_dense_rows
+        return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape)
+
+    else:
+        count = int64(sparse_indices.shape[0])  # Add count of all rows already in sparse_indices
+        for i in range(1, int64(sparse_indices.shape[0])):
+            index = int64(sparse_indices[i, 0])
+            prev_index = int64(sparse_indices[i - 1, 0] + 1)
+
+            if index > prev_index:
+                count += index - prev_index  # Add count of all rows between two consecutive indices
+
+        count += int64(sparse_indices[0, 0])  # Add count from 0 to first row id in sparse_indices
+        count += int64(
+            num_dense_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]
+        )  # Add count from last row id to dense_shape - 1
+        new_sparse_indices_shape[0] = int64(count)
+        new_sparse_indices_shape[1] = int64(sparse_indices.shape[1])
+        new_sparse_values_shape[0] = int64(count)
+        empty_row_indicator_shape[0] = num_dense_rows
+        return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape)
+
+
+@_reg.register_shape_func("sparse_fill_empty_rows", True)
+def sparse_fill_empty_rows_func(attrs, inputs, _):
+    return _sparse_fill_empty_rows_shape_func(inputs[0], inputs[2])
+
+
 @script
 def _layout_transform_shape_func(
     data_shape, out_layout_len, dst_equal_list, dst_mul_list, dst_div_list, dst_mix_list

python/tvm/relay/op/strategy/generic.py

@@ -1070,6 +1070,35 @@ def roi_align_strategy(attrs, inputs, out_type, target):
     return strategy
 
 
+# sparse_fill_empty_rows
+@override_native_generic_func("sparse_fill_empty_rows_strategy")
+def sparse_fill_empty_rows_strategy(attrs, outs, out_type, target):
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_sparse_fill_empty_rows(topi.sparse_fill_empty_rows),
+        wrap_topi_schedule(topi.generic.schedule_sparse_fill_empty_rows),
+        name="sparse_fill_empty_rows.generic",
+    )
+    return strategy
+
+
+def wrap_compute_sparse_fill_empty_rows(topi_compute):
+    """Wrap sparse_fill_empty_rows compute"""
+
+    def _compute_sparse_fill_empty_rows(attrs, inputs, output_type):
+        return topi_compute(
+            inputs[0],
+            inputs[1],
+            inputs[2],
+            inputs[3],
+            output_type.fields[0].shape,
+            output_type.fields[1].shape,
+            output_type.fields[2].shape,
+        )
+
+    return _compute_sparse_fill_empty_rows
+
+
 # roi_pool
 @generic_func
 def schedule_roi_pool(attrs, outs, target):

python/tvm/relay/op/transform.py

@@ -1377,6 +1377,78 @@ def adv_index(inputs):
     return _make.adv_index(Tuple(inputs))
 
 
+def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value):
+    """
+    Fill rows in a sparse matrix that do no contain any values. Values are placed in the first
+    column of empty rows. The sparse array is in COO format.
+    It returns a TupleWrapper with 3 outputs
+    Parameters
+    ----------
+    sparse_indices : relay.Expr
+        A 2-D tensor[N, ndims] of integers containing location of sparse values, where N is
+        the number of sparse values and n_dim is the number of dimensions of the dense_shape.
+        The first column of this relay parameter must be sorted in ascending order.
+    sparse_values : relay.Expr
+        A 1-D tensor[N] containing the sparse values for the sparse indices.
+    dense_shape : relay.Expr
+        A 1-D tensor[ndims] which contains shape of the dense output tensor.
+    default_value : relay.Expr
+        A 1-D tensor[1] containing the default value for the remaining locations.
+    Returns
+    -------
+    new_sparse_indices : relay.Expr
+        A 2-D tensor[?, ndims] of integers containing location of new sparse
+        indices. The first column outputs must be sorted in ascending order.
+    new_sparse_values : relay.Expr
+        A 1-D tensor[?] containing the sparse values for the sparse indices.
+    empty_row_indicator : relay.Expr
+        A 1-D tensor[dense_shape[0]] filled with zeros and ones
+        indicating whether the particular row is empty or full respectively
+
+    Note
+    ----
+    This op exactly follows the documentation here:
+    https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows
+    There are two exceptions:
+    1. Input Sparse Indices are expected to be in row-major order.
+    2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False).
+
+    Examples
+    -------
+    .. code-block:: python
+        sparse_indices = [[0, 1],
+                         [0, 3],
+                         [2, 0],
+                         [3, 1]]
+        sparse_values = [1, 2, 3, 4]
+        default_value = [10]
+        dense_shape = [5, 6]
+        new_sparse_indices, empty_row_indicator, new_sparse_values, slice_element_index =
+                            relay.sparse_fill_empty_rows(
+                            sparse_indices,
+                            sparse_values,
+                            default_value,
+                            dense_shape)
+        new_sparse_indices = [[0, 1],
+                             [0, 3],
+                             [1, 0],
+                             [2, 0],
+                             [3, 1],
+                             [4, 0]]
+        empty_row_indicator = [False, True, False, False, True]
+        new_sparse_values = [1, 2, 10, 3, 4, 10]
+
+    """
+    new_sparse_indices, new_sparse_values, empty_row_indicator = TupleWrapper(
+        _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 3
+    )
+    new_sparse_indices = cast_like(new_sparse_indices, sparse_indices)
+    new_sparse_values = cast_like(new_sparse_values, sparse_values)
+    empty_row_indicator = cast(empty_row_indicator, "bool")
+
+    return Tuple((new_sparse_indices, new_sparse_values, empty_row_indicator))
+
+
 def cumsum(data, axis=None, dtype=None, exclusive=None):
     """Numpy style cumsum op. Return the cumulative inclusive sum of the elements along
     a given axis.

python/tvm/topi/__init__.py

@@ -38,6 +38,7 @@
 from .broadcast import *
 from .sort import *
 from .scatter import *
+from .sparse_fill_empty_rows import *
 from .scatter_add import *
 from .argwhere import *
 from .interpolate import *

python/tvm/topi/generic/search.py

@@ -82,3 +82,7 @@ def schedule_interpolate(outs):
       The computation schedule for the op.
     """
     return _default_schedule(outs, False)
+
+
+def schedule_sparse_fill_empty_rows(outs):
+    return _default_schedule(outs, False)

python/tvm/topi/sparse_fill_empty_rows.py

@@ -0,0 +1,109 @@
+# 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, WITHnew_sparse_indices 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=no-else-return, too-many-locals, too-many-arguments, too-many-branches
+# pylint: disable=undefined-variable, invalid-name
+"""SparseFillEmptyRows operator"""
+from ..te import hybrid
+
+
+@hybrid.script
+def _sparse_fill_empty_rows(
+    sparse_indices,
+    sparse_values,
+    dense_shape,
+    default_value,
+    new_sparse_indices_shape,
+    new_sparse_values_shape,
+    empty_row_indicator_shape,
+):
+    default_value_ = int64(default_value[0])
+    new_sparse_indices = output_tensor(new_sparse_indices_shape, "int64")
+    new_sparse_values = output_tensor(new_sparse_values_shape, "int64")
+    empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64")
+    new_sparse_indices_row_id = 0
+
+    if int64(sparse_indices.shape[0]) == int64(0):  # Handle Empty Case
+        #  Fill all rows with default values
+        for i in range(0, new_sparse_indices_shape[0]):
+            new_sparse_indices[i, 0] = int64(i)
+            new_sparse_values[i] = default_value_
+            empty_row_indicator[i] = int64(1)
+            for k in range(1, int64(new_sparse_indices_shape[1])):
+                new_sparse_indices[i, k] = int64(0)
+
+        return (new_sparse_indices, new_sparse_values, empty_row_indicator)
+
+    else:
+        # Iterate through sparse_indices and add rows if/when required
+        for i in range(0, int64(sparse_indices.shape[0])):
+            if i == 0:
+                prev_row_id = int64(0)
+            else:
+                prev_row_id = int64(sparse_indices[i - 1, 0] + 1)
+            row_id = int64(sparse_indices[i, 0])
+
+            # Since input is in row-major order, add rows between prev_row_id and row_id
+            for j in range(prev_row_id, row_id):
+                new_sparse_indices[new_sparse_indices_row_id, 0] = int64(j)
+                for k in range(1, int64(new_sparse_indices_shape[1])):
+                    new_sparse_indices[new_sparse_indices_row_id, k] = int64(0)
+                empty_row_indicator[prev_row_id] = int64(1)
+                new_sparse_values[new_sparse_indices_row_id] = default_value_
+                new_sparse_indices_row_id += 1
+
+            # Add current element to output
+            new_sparse_indices[new_sparse_indices_row_id, 0] = row_id
+            for k in range(1, int64(new_sparse_indices_shape[1])):
+                new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k])
+            new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i])
+            empty_row_indicator[row_id] = int64(0)
+            new_sparse_indices_row_id += 1
+
+        # Add rows with default value if last row id of sparse_indices is not dense_shape[0] - 1
+        for i in range(
+            int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0])
+        ):
+
+            new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i)
+            for k in range(1, int64(new_sparse_indices_shape[1])):
+                new_sparse_indices[new_sparse_indices_row_id, k] = int64(0)
+            empty_row_indicator[i] = int64(1)
+            new_sparse_values[new_sparse_indices_row_id] = default_value_
+            new_sparse_indices_row_id += 1
+
+        return (new_sparse_indices, new_sparse_values, empty_row_indicator)
+
+
+def sparse_fill_empty_rows(
+    sparse_indices,
+    sparse_values,
+    dense_shape,
+    default_value,
+    new_sparse_indices_shape,
+    new_sparse_values_shape,
+    empty_row_indicator_shape,
+):
+    return _sparse_fill_empty_rows(
+        sparse_indices,
+        sparse_values,
+        dense_shape,
+        default_value,
+        new_sparse_indices_shape,
+        new_sparse_values_shape,
+        empty_row_indicator_shape,
+    )