tkonolige commented on a change in pull request #7149:
URL: https://github.com/apache/tvm/pull/7149#discussion_r547482617
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
Review comment:
```suggestion
ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but "
<< types.size() << " were provided.";
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
Review comment:
```suggestion
ICHECK_NOTNULL(param);
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ Array<PrimExpr> new_data_shape;
+ new_data_shape.push_back(tvm::max(indices->shape[0], param->num_segments));
+ for (int i = 1; i < static_cast<int>(data->shape.size()); ++i) {
+ new_data_shape.push_back(data->shape[i]);
+ }
+ std::vector<Type> fields;
+ fields.push_back(TensorType(new_data_shape, data->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> SparseSegmentSumCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input but
provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
Review comment:
```suggestion
ICHECK_NOTNULL(param);
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ Array<PrimExpr> new_data_shape;
+ new_data_shape.push_back(tvm::max(indices->shape[0], param->num_segments));
+ for (int i = 1; i < static_cast<int>(data->shape.size()); ++i) {
+ new_data_shape.push_back(data->shape[i]);
+ }
+ std::vector<Type> fields;
+ fields.push_back(TensorType(new_data_shape, data->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> SparseSegmentSumCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input but
provided "
+ << inputs.size();
Review comment:
```suggestion
ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input
but " << inputs.size() << " were provided.";
```
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1320,3 +1320,55 @@ def adv_index(inputs):
Output tensor.
"""
return _make.adv_index(Tuple(inputs))
+
+
+def sparse_segment_sum(data, indices, segment_ids, num_segments=None):
+ """
+ Compute the sparse segment sum on the indices over the segment_ids
Review comment:
I think we need a lot more documentation here. Please include a
description of what a segment sum is and how a sparse segment sum is different
from a dense one.
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ Array<PrimExpr> new_data_shape;
+ new_data_shape.push_back(tvm::max(indices->shape[0], param->num_segments));
+ for (int i = 1; i < static_cast<int>(data->shape.size()); ++i) {
+ new_data_shape.push_back(data->shape[i]);
+ }
+ std::vector<Type> fields;
+ fields.push_back(TensorType(new_data_shape, data->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> SparseSegmentSumCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input but
provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ return {topi::SparseSegmentSum(inputs[0], inputs[1], inputs[2],
param->num_segments)};
+}
+
+Expr MakeSparseSegmentSum(Expr data, Expr indices, Expr segment_ids, int
num_segments) {
+ auto attrs = make_object<SparseSegmentSumAttrs>();
+ attrs->num_segments = std::move(num_segments);
+ static const Op& op = Op::Get("sparse_segment_sum");
+ return Call(op, {data, indices, segment_ids}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.sparse_segment_sum").set_body_typed(MakeSparseSegmentSum);
+
+RELAY_REGISTER_OP("sparse_segment_sum")
+ .describe(R"code(Return sparse segment sum of the tensor given segments
+)code" TVM_ADD_FILELINE)
+ .set_num_inputs(3)
+ .set_attrs_type<SparseSegmentSumAttrs>()
+ .add_argument("data", "Tensor", "The first tensor")
+ .add_argument("indices", "Tensor", "The second tensor")
+ .add_argument("segment_ids", "Tensor", "The third tensor")
+ .add_type_rel("sparse_segment_sum", SparseSegmentSumRel)
+ .set_attr<TOpPattern>("TOpPattern", kInjective)
Review comment:
@mbrookhart Is `kInjective` the correct pattern here?
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ Array<PrimExpr> new_data_shape;
+ new_data_shape.push_back(tvm::max(indices->shape[0], param->num_segments));
+ for (int i = 1; i < static_cast<int>(data->shape.size()); ++i) {
+ new_data_shape.push_back(data->shape[i]);
+ }
+ std::vector<Type> fields;
+ fields.push_back(TensorType(new_data_shape, data->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> SparseSegmentSumCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input but
provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ return {topi::SparseSegmentSum(inputs[0], inputs[1], inputs[2],
param->num_segments)};
+}
+
+Expr MakeSparseSegmentSum(Expr data, Expr indices, Expr segment_ids, int
num_segments) {
+ auto attrs = make_object<SparseSegmentSumAttrs>();
+ attrs->num_segments = std::move(num_segments);
+ static const Op& op = Op::Get("sparse_segment_sum");
+ return Call(op, {data, indices, segment_ids}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.sparse_segment_sum").set_body_typed(MakeSparseSegmentSum);
+
+RELAY_REGISTER_OP("sparse_segment_sum")
+ .describe(R"code(Return sparse segment sum of the tensor given segments
+)code" TVM_ADD_FILELINE)
+ .set_num_inputs(3)
+ .set_attrs_type<SparseSegmentSumAttrs>()
+ .add_argument("data", "Tensor", "The first tensor")
+ .add_argument("indices", "Tensor", "The second tensor")
+ .add_argument("segment_ids", "Tensor", "The third tensor")
Review comment:
Please make the description better here.
----------------------------------------------------------------
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:
us...@infra.apache.org
