tkonolige commented on a change in pull request #7442:
URL: https://github.com/apache/tvm/pull/7442#discussion_r576945730
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1322,6 +1322,71 @@ def adv_index(inputs):
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.
Review comment:
```suggestion
Fill rows in a sparse matrix that do no contain any values. Values are
placed in the first column of empty rows.
```
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1322,6 +1322,71 @@ def adv_index(inputs):
return _make.adv_index(Tuple(inputs))
+def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape,
default_value):
Review comment:
Please document that this takes a sparse matrix in coordinate (COO)
format.
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1322,6 +1322,71 @@ def adv_index(inputs):
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 3 outputs
+ Parameters
+ ----------
+ sparse_indices : relay.Expr
+ A 2-D int64 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 inputs for this relay parameter must be in row-major order.
+ sparse_values : relay.Expr
+ A 1-D int64 tensor[N] containing the sparse values for the sparse
indices.
+ dense_shape : relay.Expr
+ A 1-D int64 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 int64 tensor[?, ndims] of integers containing location of new
sparse
+ indices
+ new_sparse_values : relay.Expr
+ A 1-D int64 tensor[?] containing the sparse values for the sparse
indices.
+ empty_row_indicator : relay.Expr
+ A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones
+ indicating whether the particular row is empty or full respectively
+
+ Note:
Review comment:
```suggestion
Note
----
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1584,6 +1584,51 @@ RELAY_REGISTER_OP("repeat")
.set_attr<FTVMCompute>("FTVMCompute", RepeatCompute)
.set_attr<TOpPattern>("TOpPattern", kBroadcast);
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [sparse_indices, sparse_values, dense_shape, default_value, result]
+ ICHECK_EQ(types.size(), 5) << "SparseFillEmptyRowsRel expects 5 inputs but "
<< types.size()
+ << "provided";
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto ndims = sparse_indices->shape[1];
+ fields.push_back(TensorType(Array<PrimExpr>{Any(), ndims},
tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ reporter->Assign(types[types.size() - 1], TupleType(Array<Type>(fields)));
+ return true;
+}
+
+Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr
dense_shape,
+ Expr default_value) {
+ static const Op& op = Op::Get("sparse_fill_empty_rows");
+ return Call(op, {sparse_indices, sparse_values, dense_shape, default_value},
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(Fill empty rows of a sparse tensor with a default value.)code"
TVM_ADD_FILELINE)
+ .set_num_inputs(4)
+ .add_argument("sparse_indices", "Tensor",
+ "A 2-D int64 tensor of shape [N, ndims], which specifies the
indices of the"
+ "elements in the sparse tensor that contain nonzero values.
COO Format")
+ .add_argument(
+ "sparse_values", "Tensor",
+ "A 1-D tensor[N] which supplies the values for each element in
indices. COO Format")
+ .add_argument("dense_shape", "Tensor",
+ "A 1-D int64 tensor of shape [ndims], which specifies the
dense_shape of the"
+ "sparse tensor. Takes a list indicating the number of
elements in each "
+ "dimension. COO Format")
+ .add_argument(
+ "default_value", "Tensor",
+ "The value to fill for empty rows, with the same type as
sparse_values. COO Format")
Review comment:
```suggestion
"The value to fill for empty rows, with the same type as
sparse_values.")
```
##########
File path: python/tvm/topi/sparse_fill_empty_rows.py
##########
@@ -0,0 +1,117 @@
+# 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
+
+
[email protected]
+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")
+ idx = 0
+
+ if int64(sparse_indices.shape[0]) == int64(0):
+ for i in range(0, int64(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:
+ for i in range(0, int64(sparse_indices[0, 0])):
+ new_sparse_indices[idx, 0] = int64(i)
+ for k in range(1, int64(new_sparse_indices_shape[1])):
+ new_sparse_indices[idx, k] = int64(0)
+
+ new_sparse_values[idx] = default_value_
+ empty_row_indicator[i] = int64(1)
+ idx += 1
+
+ for i in range(0, int64(sparse_indices.shape[0])):
+ index = int64(sparse_indices[i, 0])
+ if i == 0:
+ new_sparse_indices[idx, 0] = index
+ for k in range(1, int64(new_sparse_indices_shape[1])):
+ new_sparse_indices[idx, k] = int64(sparse_indices[i, k])
+ new_sparse_values[idx] = int64(sparse_values[i])
+ empty_row_indicator[index] = int64(0)
+ idx += 1
+ else:
+ prev_index = int64(sparse_indices[i - 1, 0] + 1)
+ for j in range(prev_index, index):
+ new_sparse_indices[idx, 0] = int64(j)
+ for k in range(1, int64(new_sparse_indices_shape[1])):
+ new_sparse_indices[idx, k] = int64(0)
+ empty_row_indicator[prev_index] = int64(1)
+ new_sparse_values[idx] = default_value_
+ idx += 1
+
+ new_sparse_indices[idx, 0] = index
+ for k in range(1, int64(new_sparse_indices_shape[1])):
+ new_sparse_indices[idx, k] = int64(sparse_indices[i, k])
+ new_sparse_values[idx] = int64(sparse_values[i])
+ empty_row_indicator[index] = int64(0)
+ idx += 1
+
+ for i in range(
+ int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1),
int64(dense_shape[0])
Review comment:
Sorry, I got it in my head that we were in CSR format. This looks
correct.
##########
File path: python/tvm/topi/sparse_fill_empty_rows.py
##########
@@ -0,0 +1,124 @@
+# 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
+
+
[email protected]
+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,
sparse_indices.dtype)
+ new_sparse_values = output_tensor(new_sparse_values_shape,
sparse_values.dtype)
+ 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, int64(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:
+ # Add rows with default value if first row id of sparse_indices is not
a zero
+ for i in range(0, int64(sparse_indices[0, 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)
+
+ new_sparse_values[new_sparse_indices_row_id] = default_value_
+ empty_row_indicator[i] = int64(1)
+ new_sparse_indices_row_id += 1
+
+ # Iterate through sparse_indices and add rows if/when required
+ for i in range(0, int64(sparse_indices.shape[0])):
+ row_id = int64(sparse_indices[i, 0])
+ if i == 0:
+ # Add first row of input 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
+ else:
+ prev_row_id = int64(sparse_indices[i - 1, 0] + 1)
+ # 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 row 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
Review comment:
You can simplify the logic here.
```suggestion
# 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
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1584,6 +1584,51 @@ RELAY_REGISTER_OP("repeat")
.set_attr<FTVMCompute>("FTVMCompute", RepeatCompute)
.set_attr<TOpPattern>("TOpPattern", kBroadcast);
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [sparse_indices, sparse_values, dense_shape, default_value, result]
+ ICHECK_EQ(types.size(), 5) << "SparseFillEmptyRowsRel expects 5 inputs but "
<< types.size()
+ << "provided";
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto ndims = sparse_indices->shape[1];
+ fields.push_back(TensorType(Array<PrimExpr>{Any(), ndims},
tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ reporter->Assign(types[types.size() - 1], TupleType(Array<Type>(fields)));
+ return true;
+}
+
+Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr
dense_shape,
+ Expr default_value) {
+ static const Op& op = Op::Get("sparse_fill_empty_rows");
+ return Call(op, {sparse_indices, sparse_values, dense_shape, default_value},
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(Fill empty rows of a sparse tensor with a default value.)code"
TVM_ADD_FILELINE)
+ .set_num_inputs(4)
+ .add_argument("sparse_indices", "Tensor",
+ "A 2-D int64 tensor of shape [N, ndims], which specifies the
indices of the"
+ "elements in the sparse tensor that contain nonzero values.
COO Format")
+ .add_argument(
+ "sparse_values", "Tensor",
+ "A 1-D tensor[N] which supplies the values for each element in
indices. COO Format")
+ .add_argument("dense_shape", "Tensor",
+ "A 1-D int64 tensor of shape [ndims], which specifies the
dense_shape of the"
+ "sparse tensor. Takes a list indicating the number of
elements in each "
+ "dimension. COO Format")
Review comment:
```suggestion
"dimension.")
```
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1322,6 +1322,71 @@ def adv_index(inputs):
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 3 outputs
+ Parameters
+ ----------
+ sparse_indices : relay.Expr
+ A 2-D int64 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 inputs for this relay parameter must be in row-major order.
Review comment:
Row major order means that the first dimension of the array indicates
the row indices. I think what you want to say here is that `sparse_indices`
must be sorted such that the first column is in ascending order.
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1584,6 +1584,47 @@ RELAY_REGISTER_OP("repeat")
.set_attr<FTVMCompute>("FTVMCompute", RepeatCompute)
.set_attr<TOpPattern>("TOpPattern", kBroadcast);
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [sparse_indices, sparse_values, dense_shape, default_value, result]
+ ICHECK_EQ(types.size(), 5);
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto ndims = sparse_indices->shape[1];
+ fields.push_back(TensorType(Array<PrimExpr>{Any(), ndims},
tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ fields.push_back(TensorType(Array<PrimExpr>{Any()}, tvm::DataType::Int(64)));
+ reporter->Assign(types[types.size() - 1], TupleType(Array<Type>(fields)));
+ return true;
+}
+
+Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr
dense_shape,
+ Expr default_value) {
+ static const Op& op = Op::Get("sparse_fill_empty_rows");
+ return Call(op, {sparse_indices, sparse_values, dense_shape, default_value},
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(4)
+ .add_argument("sparse_indices", "Tensor",
+ "A 2-D int64 tensor of shape [N, ndims], which specifies the
indices of the"
+ "elements in the sparse tensor that contain nonzero values")
+ .add_argument("sparse_values", "Tensor",
+ "A 1-D tensor[N] which supplies the values for each element
in indices")
+ .add_argument("dense_shape", "Tensor",
Review comment:
Sure, we can keep it.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]