Add thrust support for nms

[Laurawly]

@kazum @masahi @icemelon9 Please feel free to raise any suggestions or comments.

[masahi]

@Laurawly Can we refactor thrust_sort and thrust_sort_nms? I see considerable duplicate.

[kazum]

Add check for the axis. For example,

assert axis==len(data.shape)-1, "Supports sorting along the last axis only"

[Laurawly]

In original sort_nms, we do support axis other than -1, I'll add your swap operations in this operator.

[masahi]

@Laurawly Can we have implementation like this? Big if/else blocks should be put inside thrust_sort_common.

TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort")
.set_body([](TVMArgs args, TVMRetValue* ret) {
  CHECK_GE(args.num_args, 5);
  DLTensor* input = args[0];
  DLTensor* valid_count = args[1];
  DLTensor* values_out = args[2];
  DLTensor* indices_out = args[3];
  bool is_ascend = args[4];

  auto data_dtype = DLDataType2String(input->dtype);
  auto out_dtype = DLDataType2String(indices_out->dtype);

  int n_values = input->shape[input->ndim - 1];
  auto get_sort_len = [=](int i) { return n_values; };
  thrust_sort_common(input, values_out, indices_out, is_ascend, get_sort_len, data_type, out_dtype);
}

TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort_nms")
.set_body([](TVMArgs args, TVMRetValue* ret) {
  CHECK_GE(args.num_args, 5);
  DLTensor* input = args[0];
  DLTensor* valid_count = args[1];
  DLTensor* values_out = args[2];
  DLTensor* indices_out = args[3];
  bool is_ascend = args[4];

  auto data_dtype = DLDataType2String(input->dtype);
  auto out_dtype = DLDataType2String(indices_out->dtype);

  thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
  auto get_sort_len = [&valid_count_ptr](int i) { return valid_count_ptr[i] };
  thrust_sort_common(input, values_out, indices_out, is_ascend, get_sort_len, data_type, out_dtype);
}

[Laurawly]

@Laurawly Can we have implementation like this? Big if/else blocks should be put inside thrust_sort_common.

TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort")
.set_body([](TVMArgs args, TVMRetValue* ret) {
  CHECK_GE(args.num_args, 5);
  DLTensor* input = args[0];
  DLTensor* valid_count = args[1];
  DLTensor* values_out = args[2];
  DLTensor* indices_out = args[3];
  bool is_ascend = args[4];

  auto data_dtype = DLDataType2String(input->dtype);
  auto out_dtype = DLDataType2String(indices_out->dtype);

  int n_values = input->shape[input->ndim - 1];
  auto get_sort_len = [=](int i) { return n_values; };
  thrust_sort_common(input, values_out, indices_out, is_ascend, get_sort_len, data_type, out_dtype);
}

TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort_nms")
.set_body([](TVMArgs args, TVMRetValue* ret) {
  CHECK_GE(args.num_args, 5);
  DLTensor* input = args[0];
  DLTensor* valid_count = args[1];
  DLTensor* values_out = args[2];
  DLTensor* indices_out = args[3];
  bool is_ascend = args[4];

  auto data_dtype = DLDataType2String(input->dtype);
  auto out_dtype = DLDataType2String(indices_out->dtype);

  thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
  auto get_sort_len = [&valid_count_ptr](int i) { return valid_count_ptr[i] };
  thrust_sort_common(input, values_out, indices_out, is_ascend, get_sort_len, data_type, out_dtype);
}

Lemme know if the changes look good to you.

[masahi]

Lemme know if the changes look good to you

Yes this looks great! Thanks

[kazum]

Looks good to me, thanks!

[masahi]

Thanks @Laurawly @kazum