[GitHub] [tvm] tkonolige commented on a change in pull request #7149: Sparse segment sum op

[GitBox <gi...@apache.org>]

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:
users@infra.apache.org