You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/03/20 22:01:53 UTC
[GitHub] [incubator-tvm] Laurawly opened a new pull request #5116: Add
thrust support for nms
Laurawly opened a new pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116
@kazum @masahi @icemelon9 Please feel free to raise any suggests or comments.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r395939260
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
if we replace `current_values` with `n_values`, the sorting code above and below looks identical with `thrust_sort`
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] kazum commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
kazum commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396213537
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -35,11 +35,12 @@ namespace contrib {
using namespace runtime;
// Performs sorting along axis -1 and returns both sorted values and indices.
-template<typename DataType, typename IndicesType>
+template<typename DataType, typename IndicesType, typename F>
void thrust_sort(DLTensor* input,
DLTensor* out_values,
DLTensor* out_indices,
- bool is_ascend) {
+ bool is_ascend,
+ const F &get_sort_len) {
Review comment:
I think `std::function<int(int)>` looks easier to read than introducing typename F.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396131139
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
Shall I add `assert 0 <= axis < ndim` after this line?
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on issue #5116: Add thrust support
for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on issue #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#issuecomment-602303395
> Lemme know if the changes look good to you
Yes this looks great! Thanks
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396131529
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
In original sort_nms, we do support `axis` other than -1, I'll add your swap operations in this operator.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on issue #5116: Add thrust support
for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on issue #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#issuecomment-601960643
@Laurawly Can we refactor `thrust_sort` and `thrust_sort_nms`? I see considerable duplicate.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on issue #5116: Add thrust
support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on issue #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#issuecomment-602302249
> @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.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396130839
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
We don't only support sorting along the last axis.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi merged pull request #5116: Add thrust
support for nms
Posted by GitBox <gi...@apache.org>.
masahi merged pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396053175
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
`thrust_sort_common` should take a function (lambda) as an argument, which takes current iter (`i` in the code) and returns:
* For existing `thrust_sort`, always return `n_values`
* For `thrust_sort_nms`, returns `valid_count_ptr[i]` (`current_values`)
Does this make sense?
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] kazum commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
kazum commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r395956862
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,51 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ data_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "data_buf",
+ data_alignment=8)
+ valid_count_buf = tvm.tir.decl_buffer(valid_count.shape, valid_count.dtype,
+ "valid_count_buf", data_alignment=4)
+ out_bufs = [
+ tvm.tir.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8),
+ tvm.tir.decl_buffer(data.shape, "int32", "indices_buf", data_alignment=8)
+ ]
+ out = te.extern([data.shape, data.shape],
+ [data, valid_count],
+ lambda ins, outs: tvm.tir.call_packed(
+ "tvm.contrib.thrust.sort_nms", ins[0], ins[1], outs[0], outs[1], is_ascend),
Review comment:
thrust_sort_nms() supports only axis -1.
Should we add `assert(axis == -1)`?
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on issue #5116: Add thrust support
for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on issue #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#issuecomment-602147401
@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);
}
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r395936924
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
Review comment:
Align the indent
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396131139
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
Shall I add `assert 0 <= axis < ndim` after this line?
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396054117
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
+ thrust::sequence(indices_ptr, indices_ptr + current_values);
+ if (is_ascend) {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr);
+ } else {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr,
+ thrust::greater<DataType>());
+ }
+ values_ptr += current_values;
+ indices_ptr += current_values;
+ }
+}
+
+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);
+
+ if (data_dtype == "float32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<float, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<float, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<float, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<float, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "float64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<double, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<double, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<double, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<double, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int32_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int32_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int32_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int32_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int64_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int64_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int64_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int64_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else {
+ LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
+ }
+});
+
Review comment:
I mean we shouldn't duplicate this if/else blocks. See my comment below.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396164708
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -53,81 +54,114 @@ void thrust_sort(DLTensor* input,
thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
for (int i = 0 ; i < n_iter; ++i) {
- thrust::sequence(indices_ptr, indices_ptr + n_values);
+ int current_values = get_sort_len(i);
Review comment:
I prefer renaming `current_values` to `n_values` 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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396052230
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
+ thrust::sequence(indices_ptr, indices_ptr + current_values);
+ if (is_ascend) {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr);
+ } else {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr,
+ thrust::greater<DataType>());
+ }
+ values_ptr += current_values;
+ indices_ptr += current_values;
+ }
+}
+
+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);
+
+ if (data_dtype == "float32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<float, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<float, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<float, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<float, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "float64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<double, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<double, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<double, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<double, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int32_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int32_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int32_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int32_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int64_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int64_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int64_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int64_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else {
+ LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
+ }
+});
+
Review comment:
Could you elaborate more on this?
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396053175
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
`thrust_sort_common` should take a function (lambda) as an argument, which takes current iter (`i` in the code) and returns:
* For existing `thrust_sort`, always return `n_values`
* For `thrust_sort_nms`, returns `valid_count_ptr[i]` (`current_values`)
Does this make sense? I remember NVCC should not have any problem compiling C++ lambdas.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396052223
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
`thrust_sort_nms` needs one more input tensor than `thrust_sort`, I'm not sure how to pass the `valid_count` tensor without creating a new function.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396053175
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
`thrust_sort_common` should take a function (lambda) as an argument, which takes current iter (`i` in the code) and returns:
* For existing `thrust_sort`, always return `n_values`
* For `thrust_sort_nms`, returns `valid_count_ptr[i]` (`current_values`)
Does this make sense? I remeber NVCC should not have any problem compiling C++ lambdas.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396164360
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -53,81 +54,112 @@ void thrust_sort(DLTensor* input,
thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
for (int i = 0 ; i < n_iter; ++i) {
- thrust::sequence(indices_ptr, indices_ptr + n_values);
+ int current_values = get_sort_len(i);
+ thrust::sequence(indices_ptr, indices_ptr + current_values);
if (is_ascend) {
- thrust::sort_by_key(values_ptr, values_ptr + n_values, indices_ptr);
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr);
} else {
- thrust::sort_by_key(values_ptr, values_ptr + n_values, indices_ptr,
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr,
thrust::greater<DataType>());
}
- values_ptr += n_values;
- indices_ptr += n_values;
+ values_ptr += current_values;
+ indices_ptr += current_values;
}
}
-TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort")
-.set_body([](TVMArgs args, TVMRetValue* ret) {
- CHECK_GE(args.num_args, 4);
- DLTensor* input = args[0];
- DLTensor* values_out = args[1];
- DLTensor* indices_out = args[2];
- bool is_ascend = args[3];
-
- auto data_dtype = DLDataType2String(input->dtype);
- auto out_dtype = DLDataType2String(indices_out->dtype);
-
+template<typename F>
+void thrust_sort_common(DLTensor* input,
+ DLTensor* values_out,
+ DLTensor* indices_out,
+ bool is_ascend,
+ F &get_sort_len,
+ std::string data_dtype,
+ std::string out_dtype) {
if (data_dtype == "float32") {
if (out_dtype == "int32") {
- thrust_sort<float, int32_t>(input, values_out, indices_out, is_ascend);
+ thrust_sort<F, float, int32_t>(input, values_out, indices_out, is_ascend, get_sort_len);
Review comment:
I think if you bring `typename F` after `typename, DataType, typename IndicesType` in `thrust_sort`, you can remove explicit `F` in `thrust_sort<F, ...>`
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] Laurawly commented on a change in pull request
#5116: Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396130839
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
We don't only support sorting along the last axis.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r395939482
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
+ thrust::sequence(indices_ptr, indices_ptr + current_values);
+ if (is_ascend) {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr);
+ } else {
+ thrust::sort_by_key(values_ptr, values_ptr + current_values, indices_ptr,
+ thrust::greater<DataType>());
+ }
+ values_ptr += current_values;
+ indices_ptr += current_values;
+ }
+}
+
+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);
+
+ if (data_dtype == "float32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<float, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<float, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<float, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<float, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "float64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<double, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<double, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<double, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<double, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int32") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int32_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int32_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int32_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int32_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else if (data_dtype == "int64") {
+ if (out_dtype == "int32") {
+ thrust_sort_nms<int64_t, int32_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "int64") {
+ thrust_sort_nms<int64_t, int64_t>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float32") {
+ thrust_sort_nms<int64_t, float>(input, valid_count, values_out, indices_out, is_ascend);
+ } else if (out_dtype == "float64") {
+ thrust_sort_nms<int64_t, double>(input, valid_count, values_out, indices_out, is_ascend);
+ } else {
+ LOG(FATAL) << "Unsupported output dtype: " << out_dtype;
+ }
+ } else {
+ LOG(FATAL) << "Unsupported input dtype: " << data_dtype;
+ }
+});
+
Review comment:
Likely we can share this if/else with `thrust_sort` using template or lambda
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396052831
##########
File path: src/runtime/contrib/thrust/thrust.cu
##########
@@ -65,6 +65,105 @@ void thrust_sort(DLTensor* input,
}
}
+// Performs sorting along axis -1 and returns both sorted values and indices.
+template<typename DataType, typename IndicesType>
+void thrust_sort_nms(DLTensor* input,
+ DLTensor* valid_count,
+ DLTensor* out_values,
+ DLTensor* out_indices,
+ bool is_ascend) {
+ thrust::device_ptr<IndicesType> valid_count_ptr(static_cast<IndicesType *>(valid_count->data));
+ thrust::device_ptr<DataType> data_ptr(static_cast<DataType *>(input->data));
+ thrust::device_ptr<DataType> values_ptr(static_cast<DataType *>(out_values->data));
+ thrust::device_ptr<IndicesType> indices_ptr(static_cast<IndicesType *>(out_indices->data));
+
+ int n_values = input->shape[input->ndim - 1];
+ int n_iter = 1;
+ for (int i = 0; i < input->ndim - 1; ++i) {
+ n_iter *= input->shape[i];
+ }
+
+ thrust::copy(data_ptr, data_ptr + n_iter * n_values, values_ptr);
+
+ for (int i = 0 ; i < n_iter; ++i) {
+ int current_values = valid_count_ptr[i];
Review comment:
I imagine we can create another function `thrust_sort_common`, and put all common code in it. Then have existing `thrust_sort` and `thrust_sort_nms` dispatch into it appropriately.
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] kazum commented on a change in pull request #5116:
Add thrust support for nms
Posted by GitBox <gi...@apache.org>.
kazum commented on a change in pull request #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#discussion_r396053698
##########
File path: topi/python/topi/cuda/sort.py
##########
@@ -237,6 +237,53 @@ def sort_nms_ir(data, valid_count, output, axis, is_ascend):
return ib.get()
+def argsort_nms_thrust(data, valid_count, axis=-1, is_ascend=1, dtype="float32"):
+ """Performs sorting along the given axis and returns an array of indicies
+ having same shape as an input array that index data in sorted order.
+
+ Parameters
+ ----------
+ data: tvm.te.Tensor
+ The input array.
+
+ valid_count : tvm.te.Tensor, optional
+ The number of valid elements to be sorted.
+
+ axis : int, optional
+ Axis long which to sort the input tensor.
+
+ is_ascend : boolean, optional
+ Whether to sort in ascending or descending order.
+
+ dtype : string, optional
+ DType of the output indices.
+
+ Returns
+ -------
+ out : tvm.te.Tensor
+ The output of this function.
+ """
+ if axis < 0:
+ axis = len(data.shape) + axis
Review comment:
Add check for the axis. For example,
```
assert axis==len(data.shape)-1, "Supports sorting along the last axis only"
```
----------------------------------------------------------------
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
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on issue #5116: Add thrust support
for nms
Posted by GitBox <gi...@apache.org>.
masahi commented on issue #5116: Add thrust support for nms
URL: https://github.com/apache/incubator-tvm/pull/5116#issuecomment-602921562
Thanks @Laurawly @kazum
----------------------------------------------------------------
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
With regards,
Apache Git Services