perf(autoware_tensorrt_plugins): remove Thrust from sort kernels

Author: mojomex

perception/autoware_tensorrt_plugins/CMakeLists.txt

@@ -147,6 +147,25 @@ if(TRT_AVAIL AND CUDA_AVAIL AND SPCONV_AVAIL)
       spconv::spconv
   )
 
+  if(BUILD_TESTING)
+    find_package(ament_cmake_gtest REQUIRED)
+
+    ament_add_gtest(reference_kernels_test
+      test/reference_kernels_test.cpp
+    )
+    if(TARGET reference_kernels_test)
+      target_link_libraries(reference_kernels_test
+        CUDA::cudart
+        cuda_ops
+      )
+      target_include_directories(reference_kernels_test PRIVATE
+        include
+        ${CUDA_INCLUDE_DIRS}
+      )
+      target_compile_definitions(reference_kernels_test PRIVATE _GLIBCXX_USE_CXX11_ABI=1)
+    endif()
+  endif()
+
   install(
     TARGETS ${PROJECT_NAME}
     DESTINATION share/${PROJECT_NAME}/plugins

perception/autoware_tensorrt_plugins/include/autoware/unique_ops/unique.hpp

@@ -24,6 +24,7 @@ std::int64_t unique(
   std::int64_t * unique_counts, void * workspace, std::size_t num_input_elements,
   std::size_t workspace_size, cudaStream_t stream);
 
+std::size_t get_unique_temp_storage_size(std::size_t num_elements);
 std::size_t get_unique_workspace_size(std::size_t num_elements);
 
 #endif  // AUTOWARE__UNIQUE_OPS__UNIQUE_HPP_

perception/autoware_tensorrt_plugins/package.xml

@@ -21,6 +21,7 @@
   <depend>autoware_cuda_utils</depend>
 
   <test_depend>ament_cmake_ros</test_depend>
+  <test_depend>ament_cmake_gtest</test_depend>
   <test_depend>ament_lint_auto</test_depend>
   <test_depend>autoware_lint_common</test_depend>
 

perception/autoware_tensorrt_plugins/src/argsort_ops/argsort.cu

@@ -16,25 +16,52 @@
 
 #include <cub/cub.cuh>
 
-#include <thrust/device_ptr.h>
-#include <thrust/execution_policy.h>
-#include <thrust/sequence.h>
+namespace
+{
+
+constexpr int kThreadsPerBlock = 256;
+
+std::size_t align_up(const std::size_t size, const std::size_t alignment)
+{
+  return ((size + alignment - 1U) / alignment) * alignment;
+}
+
+__global__ void fill_iota(std::int64_t * output, const std::size_t num_elements)
+{
+  const auto index = static_cast<std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= num_elements) {
+    return;
+  }
+
+  output[index] = static_cast<std::int64_t>(index);
+}
+
+}  // namespace
 
 cudaError_t argsort(
   const std::int64_t * input_d, std::int64_t * output_d, void * workspace, std::size_t num_elements,
   std::size_t argsort_workspace_size, cudaStream_t stream)
 {
-  int workspace_offset = (argsort_workspace_size + sizeof(std::int64_t) - 1) / sizeof(std::int64_t);
-  thrust::device_ptr<std::int64_t> idx_ptr(
-    &reinterpret_cast<std::int64_t *>(workspace)[workspace_offset]);
+  if (num_elements == 0U) {
+    return cudaSuccess;
+  }
 
-  thrust::sequence(thrust::cuda::par.on(stream), idx_ptr, idx_ptr + num_elements, 0);
+  const auto scratch_offset = align_up(argsort_workspace_size, alignof(std::int64_t));
+  auto * input_idx_d =
+    reinterpret_cast<std::int64_t *>(reinterpret_cast<char *>(workspace) + scratch_offset);
+  auto * input_sorted_d = input_idx_d + num_elements;
 
-  std::int64_t * input_sorted_d = thrust::raw_pointer_cast(idx_ptr) + num_elements;
+  const auto num_blocks =
+    static_cast<unsigned int>((num_elements + kThreadsPerBlock - 1U) / kThreadsPerBlock);
+  fill_iota<<<num_blocks, kThreadsPerBlock, 0, stream>>>(input_idx_d, num_elements);
+  cudaError_t status = cudaGetLastError();
+  if (status != cudaSuccess) {
+    return status;
+  }
 
   return cub::DeviceRadixSort::SortPairs(
-    workspace, argsort_workspace_size, input_d, input_sorted_d, thrust::raw_pointer_cast(idx_ptr),
-    output_d, num_elements, 0, 64, stream);
+    workspace, argsort_workspace_size, input_d, input_sorted_d, input_idx_d, output_d, num_elements,
+    0, 64, stream);
 }
 
 std::size_t get_argsort_workspace_size(std::size_t num_elements)

perception/autoware_tensorrt_plugins/src/argsort_plugin.cpp

@@ -149,14 +149,11 @@ std::int32_t ArgsortPlugin::enqueue(
   cudaStream_t stream) noexcept
 {
   auto num_elements = static_cast<std::size_t>(input_desc[0].dims.d[0]);
-  if (max_num_elements_ < num_elements) {
-    max_num_elements_ = num_elements;
-    argsort_workspace_size_ = get_argsort_workspace_size(max_num_elements_);
-  }
+  const auto workspace_size = get_argsort_workspace_size(num_elements);
 
   return argsort(
     reinterpret_cast<std::int64_t const *>(inputs[0]), reinterpret_cast<std::int64_t *>(outputs[0]),
-    workspace, num_elements, argsort_workspace_size_, stream);
+    workspace, num_elements, workspace_size, stream);
 }
 
 std::int32_t ArgsortPlugin::onShapeChange(
@@ -183,8 +180,10 @@ std::size_t ArgsortPlugin::getWorkspaceSize(
   [[maybe_unused]] std::int32_t num_outputs) const noexcept
 {
   std::int64_t max_num_elements = inputs[0].max.d[0];
-  return get_argsort_workspace_size(max_num_elements) +
-         sizeof(std::int64_t) * 2 * (max_num_elements + 1);
+  const auto temp_size = get_argsort_workspace_size(max_num_elements);
+  const auto scratch_offset =
+    ((temp_size + alignof(std::int64_t) - 1U) / alignof(std::int64_t)) * alignof(std::int64_t);
+  return scratch_offset + sizeof(std::int64_t) * 2 * max_num_elements;
 }
 
 }  // namespace nvinfer1::plugin

perception/autoware_tensorrt_plugins/src/unique_ops/unique.cu

@@ -101,78 +101,152 @@
 
 #include <cub/cub.cuh>
 
-#include <thrust/adjacent_difference.h>
-#include <thrust/device_ptr.h>
-#include <thrust/execution_policy.h>
-#include <thrust/scan.h>
-#include <thrust/scatter.h>
-#include <thrust/sequence.h>
-#include <thrust/sort.h>
-#include <thrust/unique.h>
+#include <algorithm>
+#include <cstdint>
 
-std::int64_t unique(
-  const std::int64_t * input, std::int64_t * unique, std::int64_t * inverse_indices,
-  std::int64_t * unique_counts, void * workspace, std::size_t num_input_elements,
-  std::size_t unique_workspace_size, cudaStream_t stream)
+namespace
 {
-  auto policy = thrust::cuda::par.on(stream);
-
-  thrust::device_ptr<std::int64_t> idx_ptr(reinterpret_cast<std::int64_t *>(workspace));
 
-  thrust::sequence(policy, idx_ptr, idx_ptr + num_input_elements + 1, 0);
+constexpr int kThreadsPerBlock = 256;
 
-  std::int64_t * sorted_input = unique;
-  std::int64_t * sorted_idx = thrust::raw_pointer_cast(idx_ptr) + 2 * num_input_elements + 1;
-  std::int64_t * inv_loc_ptr = thrust::raw_pointer_cast(idx_ptr) + 3 * num_input_elements + 1;
+std::size_t align_up(const std::size_t size, const std::size_t alignment)
+{
+  return ((size + alignment - 1U) / alignment) * alignment;
+}
 
-  void * sort_workspace_ptr =
-    reinterpret_cast<void *>(thrust::raw_pointer_cast(idx_ptr) + 4 * num_input_elements + 1);
+std::size_t query_unique_temp_storage_size(const std::size_t num_elements)
+{
+  std::size_t sort_temp_size = 0;
+  std::size_t scan_temp_size = 0;
+  std::size_t unique_temp_size = 0;
 
-  auto sort_workspace_size =
-    unique_workspace_size - (4 * num_input_elements + 1) * sizeof(std::int64_t);
+  std::int64_t * int64_nullptr = nullptr;
+  std::int32_t * int32_nullptr = nullptr;
 
   cub::DeviceRadixSort::SortPairs(
-    sort_workspace_ptr, sort_workspace_size, input, sorted_input, thrust::raw_pointer_cast(idx_ptr),
-    sorted_idx, num_input_elements, 0, 64, stream);
-
-  auto equal = [] __device__(const std::int64_t a, const std::int64_t b) { return a == b; };
+    nullptr, sort_temp_size, int64_nullptr, int64_nullptr, int64_nullptr, int64_nullptr,
+    num_elements, 0, 64, nullptr);
+  cub::DeviceScan::InclusiveSum(
+    nullptr, scan_temp_size, int32_nullptr, int32_nullptr, num_elements, nullptr);
+  cub::DeviceSelect::UniqueByKey(
+    nullptr, unique_temp_size, int64_nullptr, int64_nullptr, int64_nullptr, int64_nullptr,
+    int64_nullptr, num_elements, nullptr);
+
+  return std::max(sort_temp_size, std::max(scan_temp_size, unique_temp_size));
+}
 
-  auto not_equal = [] __device__(const std::int64_t a, const std::int64_t b) { return a != b; };
+__global__ void mark_run_starts(
+  const std::int64_t * sorted_input, std::int32_t * run_ids, const std::size_t num_input_elements)
+{
+  const auto index = static_cast<std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= num_input_elements) {
+    return;
+  }
 
-  thrust::adjacent_difference(
-    policy, sorted_input, sorted_input + num_input_elements, inv_loc_ptr, not_equal);
+  run_ids[index] = (index == 0U || sorted_input[index] != sorted_input[index - 1U]) ? 1 : 0;
+}
 
-  cudaMemsetAsync(inv_loc_ptr, 0, sizeof(int64_t), stream);
+__global__ void fill_iota(std::int64_t * output, const std::size_t num_input_elements)
+{
+  const auto index = static_cast<std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= num_input_elements) {
+    return;
+  }
 
-  thrust::inclusive_scan(policy, inv_loc_ptr, inv_loc_ptr + num_input_elements, inv_loc_ptr);
-  thrust::scatter(
-    policy, inv_loc_ptr, inv_loc_ptr + num_input_elements, sorted_idx, inverse_indices);
+  output[index] = static_cast<std::int64_t>(index);
+}
 
-  std::int64_t num_out;
+__global__ void scatter_inverse_indices(
+  const std::int64_t * sorted_idx, const std::int32_t * run_ids, std::int64_t * inverse_indices,
+  const std::size_t num_input_elements)
+{
+  const auto index = static_cast<std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= num_input_elements) {
+    return;
+  }
 
-  std::int64_t * range_ptr = idx_ptr.get();
-  num_out =
-    thrust::unique_by_key(policy, sorted_input, sorted_input + num_input_elements, range_ptr, equal)
-      .first -
-    sorted_input;
+  inverse_indices[sorted_idx[index]] = static_cast<std::int64_t>(run_ids[index] - 1);
+}
 
-  cudaMemcpyAsync(
-    range_ptr + num_out * sizeof(int64_t), &num_input_elements, sizeof(std::int64_t),
-    cudaMemcpyHostToDevice, stream);
+__global__ void write_unique_offset_sentinel(
+  std::int64_t * unique_offsets, const std::int64_t * num_unique,
+  const std::size_t num_input_elements)
+{
+  unique_offsets[*num_unique] = static_cast<std::int64_t>(num_input_elements);
+}
 
-  thrust::adjacent_difference(policy, range_ptr + 1, range_ptr + num_out + 1, unique_counts);
+__global__ void write_unique_counts(
+  const std::int64_t * unique_offsets, const std::int64_t * num_unique, std::int64_t * unique_counts)
+{
+  const auto index = static_cast<std::size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (index >= static_cast<std::size_t>(*num_unique)) {
+    return;
+  }
 
-  return num_out;
+  unique_counts[index] = unique_offsets[index + 1U] - unique_offsets[index];
 }
 
-std::size_t get_unique_workspace_size(std::size_t num_elements)
+}  // namespace
+
+std::int64_t unique(
+  const std::int64_t * input, std::int64_t * unique, std::int64_t * inverse_indices,
+  std::int64_t * unique_counts, void * workspace, std::size_t num_input_elements,
+  std::size_t unique_workspace_size, cudaStream_t stream)
 {
-  std::size_t temp_size = 0;
-  std::int64_t * int64_nullptr = nullptr;
+  if (num_input_elements == 0U) {
+    return 0;
+  }
 
+  const auto temp_storage_size = get_unique_temp_storage_size(num_input_elements);
+  const auto scratch_offset = align_up(temp_storage_size, alignof(std::int64_t));
+  auto * scratch = reinterpret_cast<char *>(workspace) + scratch_offset;
+
+  auto * input_positions = reinterpret_cast<std::int64_t *>(scratch);
+  auto * sorted_input = input_positions + num_input_elements;
+  auto * unique_offsets = sorted_input + num_input_elements;
+  auto * num_unique_d = unique_offsets + num_input_elements + 1U;
+  auto * run_ids = reinterpret_cast<std::int32_t *>(num_unique_d + 1U);
+
+  const auto num_blocks =
+    static_cast<unsigned int>((num_input_elements + kThreadsPerBlock - 1U) / kThreadsPerBlock);
+
+  fill_iota<<<num_blocks, kThreadsPerBlock, 0, stream>>>(input_positions, num_input_elements);
   cub::DeviceRadixSort::SortPairs(
-    nullptr, temp_size, int64_nullptr, int64_nullptr, int64_nullptr, int64_nullptr, num_elements, 0,
-    64, nullptr);
+    workspace, temp_storage_size, input, sorted_input, input_positions, unique_offsets,
+    num_input_elements, 0, 64, stream);
+
+  mark_run_starts<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
+    sorted_input, run_ids, num_input_elements);
+  cub::DeviceScan::InclusiveSum(
+    workspace, temp_storage_size, run_ids, run_ids, num_input_elements, stream);
+
+  scatter_inverse_indices<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
+    unique_offsets, run_ids, inverse_indices, num_input_elements);
 
-  return temp_size + (4 * num_elements + 1) * sizeof(std::int64_t);
+  cub::DeviceSelect::UniqueByKey(
+    workspace, temp_storage_size, sorted_input, input_positions, unique, unique_offsets,
+    num_unique_d, num_input_elements, stream);
+
+  write_unique_offset_sentinel<<<1, 1, 0, stream>>>(
+    unique_offsets, num_unique_d, num_input_elements);
+  write_unique_counts<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
+    unique_offsets, num_unique_d, unique_counts);
+
+  std::int64_t num_out = 0;
+  cudaMemcpyAsync(&num_out, num_unique_d, sizeof(std::int64_t), cudaMemcpyDeviceToHost, stream);
+  cudaStreamSynchronize(stream);
+  return num_out;
+}
+
+std::size_t get_unique_temp_storage_size(std::size_t num_elements)
+{
+  return query_unique_temp_storage_size(num_elements);
+}
+
+std::size_t get_unique_workspace_size(std::size_t num_elements)
+{
+  const auto temp_size = query_unique_temp_storage_size(num_elements);
+  const auto scratch_offset = align_up(temp_size, alignof(std::int64_t));
+  return scratch_offset + (3 * num_elements + 2U) * sizeof(std::int64_t) +
+         num_elements * sizeof(std::int32_t);
 }

perception/autoware_tensorrt_plugins/src/unique_plugin.cpp

@@ -164,11 +164,12 @@ std::int32_t UniquePlugin::enqueue(
   cudaStream_t stream) noexcept
 {
   std::int64_t num_elements = input_desc[0].dims.d[0];
+  const auto workspace_size = get_unique_workspace_size(static_cast<std::size_t>(num_elements));
 
   std::int64_t num_unique_elements = unique(
     reinterpret_cast<const std::int64_t *>(inputs[0]), reinterpret_cast<std::int64_t *>(outputs[0]),
     reinterpret_cast<std::int64_t *>(outputs[1]), reinterpret_cast<std::int64_t *>(outputs[2]),
-    workspace, num_elements, workspace_size_, stream);
+    workspace, num_elements, workspace_size, stream);
 
   cudaMemcpyAsync(
     reinterpret_cast<std::int64_t *>(outputs[3]), &num_unique_elements, sizeof(std::int64_t),