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

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

codeislife99 commented on a change in pull request #7149:
URL: https://github.com/apache/tvm/pull/7149#discussion_r551089181



##########
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:
       Done. 




----------------------------------------------------------------
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