From commits-return-9160-archive-asf-public=cust-asf.ponee.io@tvm.apache.org Sun Mar 22 04:34:18 2020 Return-Path: X-Original-To: archive-asf-public@cust-asf.ponee.io Delivered-To: archive-asf-public@cust-asf.ponee.io Received: from mail.apache.org (hermes.apache.org [207.244.88.153]) by mx-eu-01.ponee.io (Postfix) with SMTP id DA96318064E for ; Sun, 22 Mar 2020 05:34:17 +0100 (CET) Received: (qmail 52321 invoked by uid 500); 22 Mar 2020 04:34:17 -0000 Mailing-List: contact commits-help@tvm.apache.org; run by ezmlm Precedence: bulk List-Help: List-Unsubscribe: List-Post: List-Id: Reply-To: dev@tvm.apache.org Delivered-To: mailing list commits@tvm.apache.org Received: (qmail 52312 invoked by uid 99); 22 Mar 2020 04:34:16 -0000 Received: from ec2-52-202-80-70.compute-1.amazonaws.com (HELO gitbox.apache.org) (52.202.80.70) by apache.org (qpsmtpd/0.29) with ESMTP; Sun, 22 Mar 2020 04:34:16 +0000 From: GitBox To: commits@tvm.apache.org Subject: [GitHub] [incubator-tvm] masahi commented on a change in pull request #5116: Add thrust support for nms Message-ID: <158485165632.17100.13058358569382520179.gitbox@gitbox.apache.org> References: In-Reply-To: Date: Sun, 22 Mar 2020 04:34:16 -0000 Content-Type: text/plain; charset=utf-8 Content-Transfer-Encoding: 8bit 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 +void thrust_sort_nms(DLTensor* input, + DLTensor* valid_count, + DLTensor* out_values, + DLTensor* out_indices, + bool is_ascend) { + thrust::device_ptr valid_count_ptr(static_cast(valid_count->data)); + thrust::device_ptr data_ptr(static_cast(input->data)); + thrust::device_ptr values_ptr(static_cast(out_values->data)); + thrust::device_ptr indices_ptr(static_cast(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