Sparse fill empty rows op

include/tvm/relay/attrs/transform.h

@@ -401,6 +401,14 @@ struct SparseToDenseAttrs : public tvm::AttrsNode<SparseToDenseAttrs> {
   }
 };  // struct SparseToDenseAttrs
 
+/*! \brief Attributes used in sparsefillemptyRows operator */
+struct SparseFillEmptyRowsAttrs : public tvm::AttrsNode<SparseFillEmptyRowsAttrs> {
+  Array<Integer> dense_shape;
+
+  TVM_DECLARE_ATTRS(SparseFillEmptyRowsAttrs, "relay.attrs.SparseFillEmptyRowsAttrs") {
+    TVM_ATTR_FIELD(dense_shape).describe("Shape of the dense output tensor");
+  }
+};  // struct SparseFillEmptyRowsAttrs
 /*! \brief Attributes for ndarray_size operator */
 struct NdarraySizeAttrs : public tvm::AttrsNode<NdarraySizeAttrs> {
   DataType dtype;

include/tvm/topi/transform.h

@@ -1386,6 +1386,96 @@ inline Array<Tensor> meshgrid(const Array<Tensor>& inputs, const std::string& in
   return result;
 }
 
+/*!
+ * \brief Fill Empty rows of a sparse tensor with default value
+ *
+ * \param sparse_indices Indices where values of the dense tensor exist
+ * \param sparse_values Values at the above indices respectively
+ * \param default_value Default value at to be used at empty rows
+ * \param dense_shape Dense shape of the sparse tensor
+ * \param name The name of the operation
+ * \param tag The tag to mark the operation
+ *
+ * \return A Tensor whose op member is the SparseFillEmptyRows operation
+ */
+inline Array<Tensor> SparseFillEmptyRows(const Tensor& sparse_indices, const Tensor& sparse_values,
+                                         const Tensor& default_value,
+                                         const Array<Integer>& dense_shape,
+                                         const std::string name = "T_sparse_fill_empty_rows",
+                                         std::string tag = kInjective) {
+  Array<Tensor> result;
+  Array<PrimExpr> sp_ordered_output_shape;
+  sp_ordered_output_shape.push_back(dense_shape[0] + sparse_indices->shape[0]);
+  if (sparse_indices->shape.size() > 1) {
+    sp_ordered_output_shape.push_back(sparse_indices->shape[1]);
+  }
+  auto empty_row_indicator =
+      compute(Array<PrimExpr>{dense_shape[0]}, [&](const Array<Var>& indices) {
+        PrimExpr ret = PrimExpr(Bool(1));
+        for (int i = 0; i < GetConstInt(sparse_indices->shape[0]); ++i) {
+          PrimExpr sparse_index;
+          if (sparse_indices->shape.size() == 1) {
+            sparse_index = sparse_indices[i];
+          } else {
+            sparse_index = sparse_indices[i][0];
+          }
+          ret = if_then_else(sparse_index == indices[0], PrimExpr(Bool(0)), ret);
+        }
+        return ret;
+      });
+  result.push_back(compute(
+      sp_ordered_output_shape,
+      [&](const Array<Var>& indices) {
+        PrimExpr ret = -1;
+        ret = if_then_else(indices[0] < sparse_indices->shape[0], sparse_indices(indices), ret);
+        PrimExpr empty_row_count = 0;
+        for (int i = 0; i < static_cast<int>(dense_shape[0]); ++i) {
+          empty_row_count =
+              if_then_else(empty_row_indicator[i], empty_row_count + 1, empty_row_count);
+          PrimExpr at_correct_index =
+              (indices[0] == (sparse_indices->shape[0] + empty_row_count - 1));
+          PrimExpr condition = at_correct_index && empty_row_indicator[i];
+
+          ret = if_then_else(condition, i, ret);
+          if (indices.size() > 1) {
+            ret = if_then_else(condition && indices[1] > 0, 0, ret);
+          }
+        }
+        return ret;
+      },
+      name, tag));
+  result.push_back(empty_row_indicator);
+  result.push_back(compute(
+      Array<PrimExpr>{sp_ordered_output_shape[0]},
+      [&](const Array<Var>& indices) {
+        PrimExpr ret = -1;
+        ret = if_then_else(indices[0] < sparse_values->shape[0], sparse_values(indices), ret);
+        PrimExpr empty_row_count = 0;
+        for (int i = 0; i < static_cast<int>(dense_shape[0]); ++i) {
+          empty_row_count =
+              if_then_else(empty_row_indicator[i], empty_row_count + 1, empty_row_count);
+          PrimExpr condition =
+              (indices[0] == sparse_values->shape[0] + empty_row_count - 1) && empty_row_count > 0;
+          ret = if_then_else(condition, default_value[0], ret);
+        }
+        return ret;
+      },
+      name, tag));
+  result.push_back(compute(
+      Array<PrimExpr>{1},
+      [&](const Array<Var>& indices) {
+        PrimExpr new_sparse_values_len = sparse_values->shape[0];
+        PrimExpr empty_row_count = 0;
+        for (int i = 0; i < static_cast<int>(dense_shape[0]); ++i) {
+          new_sparse_values_len = if_then_else(empty_row_indicator[i], new_sparse_values_len + 1,
+                                               new_sparse_values_len);
+        }
+        return new_sparse_values_len;
+      },
+      name, tag));
+  return result;
+}
+
 /*!
  * \brief Transform the layout according to \p src_layout and \p dst_layout
  * \param src the source input.

python/tvm/relay/op/_transform.py

@@ -63,6 +63,7 @@
 _reg.register_injective_schedule("sparse_to_dense")
 _reg.register_injective_schedule("matrix_set_diag")
 _reg.register_injective_schedule("adv_index")
+_reg.register_injective_schedule("sparse_fill_empty_rows")
 
 # concatenate
 _reg.register_schedule("concatenate", strategy.schedule_concatenate)

python/tvm/relay/op/transform.py

@@ -1320,3 +1320,83 @@ def adv_index(inputs):
         Output tensor.
     """
     return _make.adv_index(Tuple(inputs))
+
+
+def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value):
+    """
+    Fill first column of the empty rows with default values for a sparse array.
+    It returns a TupleWrapper with four outputs
+
+    Parameters
+    ----------
+    sparse_indices : relay.Expr
+        A 2-D tensor[N, n_dim] 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
+
+    sparse_values : relay.Expr
+        A 1-D tensor[N] containing the sparse values for the sparse indices.
+
+    dense_shape : relay.Expr
+        A list of integers. Shape of the dense output tensor.
+
+    default_value : relay.Expr
+        A 0-D tensor containing the default value for the remaining locations.
+        Defaults to 0.
+
+    Returns
+    -------
+    new_sparse_indices : relay.Expr
+        A 2-D tensor[N + dense_shape[0], n_dim] of integers containing location of new sparse
+        indices where N is the number of sparse values. It is filled with -1 at irrelevant indices
+        which will be sliced in a future op discarding non-useful elements. This is done since the
+        real rows of new_sparse_indices depends on the input.
+
+    empty_row_indicator : relay.Expr
+        A 1-D Boolean tensor[dense_shape[0]] indicating whether the particular row is empty
+
+    new_sparse_values : relay.Expr
+        A 1-D tensor[dense_shape[0]] containing the sparse values for the sparse indices. It is
+        filled with -1 at to_be_discarded indices
+
+    slice_element_index : relay.Expr
+        A 1-D tensor containing the amount of elements in the sparse_indices and new_sparse_values
+        expression to be sliced in a future op discarding non-useful elements in new_sparse_indices
+        and new_sparse_values
+
+    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],
+                             [2, 0],
+                             [3, 1],
+                             [1, 0],
+                             [4, 0],
+                             [-1, -1],
+                             [-1, -1],
+                             [-1, -1]]
+
+        empty_row_indicator = [False, True, False, False, True]
+
+        new_sparse_values = [1, 2, 3, 4, 10, 10, -1, -1, -1]
+
+        slice_element_index = [6]
+
+    """
+    return TupleWrapper(
+        _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 4
+    )

python/tvm/topi/transform.py

@@ -931,3 +931,81 @@ def adv_index(data, indices):
         Output tensor
     """
     return cpp.adv_index(data, indices)
+
+
+def sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape):
+    """
+    Fill first column of the empty rows with default values for a sparse array.
+    It returns a TupleWrapper with four outputs
+
+    Parameters
+    ----------
+    sparse_indices : relay.Expr
+        A 2-D tensor[N, n_dim] 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
+
+    sparse_values : relay.Expr
+        A 1-D tensor[N] containing the sparse values for the sparse indices.
+
+    dense_shape : relay.Expr
+        A list of integers. Shape of the dense output tensor.
+
+    default_value : relay.Expr
+        A 0-D tensor containing the default value for the remaining locations.
+        Defaults to 0.
+
+    Returns
+    -------
+    new_sparse_indices : relay.Expr
+        A 2-D tensor[N + dense_shape[0], n_dim] of integers containing location of new sparse
+        indices where N is the number of sparse values. It is filled with -1 at irrelevant indices
+        which will be sliced in a future op discarding non-useful elements. This is done since the
+        real rows of new_sparse_indices depends on the input.
+
+    empty_row_indicator : relay.Expr
+        A 1-D Boolean tensor[dense_shape[0]] indicating whether the particular row is empty
+
+    new_sparse_values : relay.Expr
+        A 1-D tensor[dense_shape[0]] containing the sparse values for the sparse indices. It is
+        filled with -1 at to_be_discarded indices
+
+    slice_element_index : relay.Expr
+        A 1-D tensor containing the amount of elements in the sparse_indices and new_sparse_values
+        expression to be sliced in a future op discarding non-useful elements in new_sparse_indices
+        and new_sparse_values
+
+    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],
+                             [2, 0],
+                             [3, 1],
+                             [1, 0],
+                             [4, 0],
+                             [-1, -1],
+                             [-1, -1],
+                             [-1, -1]]
+
+        empty_row_indicator = [False, True, False, False, True]
+
+        new_sparse_values = [1, 2, 3, 4, 10, 10, -1, -1, -1]
+
+        slice_element_index = [6]
+
+    """
+    return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape)

src/relay/op/tensor/transform.cc

@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
     .set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
     .set_attr<TOpPattern>("TOpPattern", kInjective);
 
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                            const TypeReporter& reporter) {
+  // types: [ sparse_indices, sparse_values, default_values, result]
+  ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments but provided "
+                             << types.size();
+  std::vector<Type> fields;
+  auto sparse_indices = types[0].as<TensorTypeNode>();
+  auto default_value = types[2].as<TensorTypeNode>();
+  const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+  ICHECK(param != nullptr);
+
+  Array<IndexExpr> sp_ordered_output_shape;
+  sp_ordered_output_shape.push_back(param->dense_shape[0] + sparse_indices->shape[0]);
+  if (sparse_indices->shape.size() > 1) {
+    sp_ordered_output_shape.push_back(sparse_indices->shape[1]);
+  }
+  fields.push_back(TensorType(sp_ordered_output_shape, sparse_indices->dtype));
+  fields.push_back(TensorType(Array<PrimExpr>{param->dense_shape[0]}, tvm::DataType::Bool()));
+  fields.push_back(TensorType(Array<PrimExpr>{sp_ordered_output_shape[0]}, default_value->dtype));
+  fields.push_back(TensorType(Array<PrimExpr>{1}, tvm::DataType::Int(32)));
+  reporter->Assign(types[3], TupleType(Array<Type>(fields)));
+  return true;
+}
+
+Array<te::Tensor> SparseFillEmptyRowsCompute(const Attrs& attrs, const Array<te::Tensor>& inputs,
+                                             const Type& out_type) {
+  ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3 arguments but provided "
+                              << inputs.size();
+  const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+  ICHECK(param != nullptr);
+  return {topi::SparseFillEmptyRows(inputs[0], inputs[1], inputs[2], param->dense_shape)};
+}
+
+Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr default_value,
+                             Array<Integer> dense_shape) {
+  auto attrs = make_object<SparseFillEmptyRowsAttrs>();
+  attrs->dense_shape = std::move(dense_shape);
+  static const Op& op = Op::Get("sparse_fill_empty_rows");
+  return Call(op, {sparse_indices, sparse_values, default_value}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.sparse_fill_empty_rows")
+    .set_body_typed(MakeSparseFillEmptyRows);
+
+RELAY_REGISTER_OP("sparse_fill_empty_rows")
+    .describe(R"code(Return representation of a sparse tensor with empty rows filled with default 
+    value.)code" TVM_ADD_FILELINE)
+    .set_num_inputs(3)
+    .set_attrs_type<SparseFillEmptyRowsAttrs>()
+    .add_argument("sparse_indices", "Tensor", "The first tensor")
+    .add_argument("sparse_values", "Tensor", "The second tensor")
+    .add_argument("default_value", "Tensor", "The third tensor")
+    .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel)
+    .set_support_level(3)
+    .set_attr<TOpPattern>("TOpPattern", kInjective)
+    .set_attr<FTVMCompute>("FTVMCompute", SparseFillEmptyRowsCompute);
+
 // tile operator
 TVM_REGISTER_NODE_TYPE(TileAttrs);