perf(autoware_tensorrt_plugins): keep plugin outputs stream-ordered

Author: mojomex

perception/autoware_tensorrt_plugins/include/autoware/tensorrt_plugins/get_indices_pairs_implicit_gemm_plugin.hpp

@@ -63,7 +63,7 @@ class GetIndicesPairsImplicitGemmPlugin : public IPluginV3,
   GetIndicesPairsImplicitGemmPlugin(
     const std::string & name, GetIndicesPairsImplicitGemmParameters const & params);
 
-  ~GetIndicesPairsImplicitGemmPlugin() override = default;
+  ~GetIndicesPairsImplicitGemmPlugin() override;
 
   // IPluginV3 Methods
 
@@ -131,6 +131,7 @@ class GetIndicesPairsImplicitGemmPlugin : public IPluginV3,
 
   std::string layer_name_;
   GetIndicesPairsImplicitGemmParameters params_;
+  std::int32_t * num_act_out_host_{nullptr};
   std::vector<nvinfer1::PluginField> data_to_serialize_;
   nvinfer1::PluginFieldCollection fc_to_serialize_;
 };

perception/autoware_tensorrt_plugins/include/autoware/tensorrt_plugins/get_indices_pairs_plugin.hpp

@@ -61,7 +61,7 @@ class GetIndicesPairsPlugin : public IPluginV3,
 public:
   GetIndicesPairsPlugin(const std::string & name, GetIndicesPairsParameters const & params);
 
-  ~GetIndicesPairsPlugin() override = default;
+  ~GetIndicesPairsPlugin() override;
 
   // IPluginV3 Methods
 
@@ -125,6 +125,7 @@ class GetIndicesPairsPlugin : public IPluginV3,
 
   std::string layer_name_;
   GetIndicesPairsParameters params_;
+  std::int32_t * num_act_out_host_{nullptr};
   std::vector<nvinfer1::PluginField> data_to_serialize_;
   nvinfer1::PluginFieldCollection fc_to_serialize_;
 };

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

@@ -19,10 +19,10 @@
 
 #include <cstdint>
 
-std::int64_t unique(
+cudaError_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 workspace_size, cudaStream_t stream);
+  std::int64_t * unique_counts, std::int64_t * num_unique, 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);

perception/autoware_tensorrt_plugins/src/get_indices_pairs_implicit_gemm_plugin.cpp

@@ -43,6 +43,16 @@ GetIndicesPairsImplicitGemmPlugin::GetIndicesPairsImplicitGemmPlugin(
 : layer_name_{name}, params_{params}
 {
   initFieldsToSerialize();
+  PLUGIN_ASSERT(
+    cudaMallocHost(reinterpret_cast<void **>(&num_act_out_host_), sizeof(std::int32_t)) ==
+    cudaSuccess);
+}
+
+GetIndicesPairsImplicitGemmPlugin::~GetIndicesPairsImplicitGemmPlugin()
+{
+  if (num_act_out_host_ != nullptr) {
+    cudaFreeHost(num_act_out_host_);
+  }
 }
 
 void GetIndicesPairsImplicitGemmPlugin::initFieldsToSerialize()
@@ -437,9 +447,10 @@ std::int32_t GetIndicesPairsImplicitGemmPlugin::enqueue(
 
   std::int32_t num_act_out_real = std::get<1>(pair_res);
   std::int32_t * num_act_out_data = static_cast<std::int32_t *>(outputs[4]);
+  *num_act_out_host_ = num_act_out_real;
 
   cudaError_t const status = cudaMemcpyAsync(
-    num_act_out_data, &num_act_out_real, sizeof(std::int32_t), cudaMemcpyHostToDevice, stream);
+    num_act_out_data, num_act_out_host_, sizeof(std::int32_t), cudaMemcpyHostToDevice, stream);
 
   return status;
 }

perception/autoware_tensorrt_plugins/src/get_indices_pairs_plugin.cpp

@@ -40,6 +40,16 @@ GetIndicesPairsPlugin::GetIndicesPairsPlugin(
 : layer_name_{name}, params_{params}
 {
   initFieldsToSerialize();
+  PLUGIN_ASSERT(
+    cudaMallocHost(reinterpret_cast<void **>(&num_act_out_host_), sizeof(std::int32_t)) ==
+    cudaSuccess);
+}
+
+GetIndicesPairsPlugin::~GetIndicesPairsPlugin()
+{
+  if (num_act_out_host_ != nullptr) {
+    cudaFreeHost(num_act_out_host_);
+  }
 }
 
 void GetIndicesPairsPlugin::initFieldsToSerialize()
@@ -290,11 +300,10 @@ std::int32_t GetIndicesPairsPlugin::enqueue(
   }
 
   std::int32_t * num_act_out_data = static_cast<std::int32_t *>(outputs[3]);
+  *num_act_out_host_ = num_act_out_real;
 
   cudaError_t const status = cudaMemcpyAsync(
-    num_act_out_data, &num_act_out_real, sizeof(std::int32_t), cudaMemcpyHostToDevice, stream);
-
-  cudaStreamSynchronize(stream);
+    num_act_out_data, num_act_out_host_, sizeof(std::int32_t), cudaMemcpyHostToDevice, stream);
 
   return status;
 }

perception/autoware_tensorrt_plugins/src/unique_ops/unique.cu

@@ -188,13 +188,14 @@ __global__ void write_unique_counts(
 
 }  // namespace
 
-std::int64_t unique(
+cudaError_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::int64_t * unique_counts, std::int64_t * num_unique, void * workspace,
+  std::size_t num_input_elements, std::size_t unique_workspace_size, cudaStream_t stream)
 {
+  (void)unique_workspace_size;
   if (num_input_elements == 0U) {
-    return 0;
+    return cudaMemsetAsync(num_unique, 0, sizeof(std::int64_t), stream);
   }
 
   const auto temp_storage_size = get_unique_temp_storage_size(num_input_elements);
@@ -204,38 +205,61 @@ std::int64_t unique(
   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);
+  auto * run_ids = reinterpret_cast<std::int32_t *>(unique_offsets + num_input_elements + 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(
+  cudaError_t status = cudaGetLastError();
+  if (status != cudaSuccess) {
+    return status;
+  }
+
+  status = cub::DeviceRadixSort::SortPairs(
     workspace, temp_storage_size, input, sorted_input, input_positions, unique_offsets,
     num_input_elements, 0, 64, stream);
+  if (status != cudaSuccess) {
+    return status;
+  }
 
   mark_run_starts<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
     sorted_input, run_ids, num_input_elements);
-  cub::DeviceScan::InclusiveSum(
+  status = cudaGetLastError();
+  if (status != cudaSuccess) {
+    return status;
+  }
+
+  status = cub::DeviceScan::InclusiveSum(
     workspace, temp_storage_size, run_ids, run_ids, num_input_elements, stream);
+  if (status != cudaSuccess) {
+    return status;
+  }
 
   scatter_inverse_indices<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
     unique_offsets, run_ids, inverse_indices, num_input_elements);
+  status = cudaGetLastError();
+  if (status != cudaSuccess) {
+    return status;
+  }
 
-  cub::DeviceSelect::UniqueByKey(
+  status = cub::DeviceSelect::UniqueByKey(
     workspace, temp_storage_size, sorted_input, input_positions, unique, unique_offsets,
-    num_unique_d, num_input_elements, stream);
+    num_unique, num_input_elements, stream);
+  if (status != cudaSuccess) {
+    return status;
+  }
 
   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);
+    unique_offsets, num_unique, num_input_elements);
+  status = cudaGetLastError();
+  if (status != cudaSuccess) {
+    return status;
+  }
 
-  std::int64_t num_out = 0;
-  cudaMemcpyAsync(&num_out, num_unique_d, sizeof(std::int64_t), cudaMemcpyDeviceToHost, stream);
-  cudaStreamSynchronize(stream);
-  return num_out;
+  write_unique_counts<<<num_blocks, kThreadsPerBlock, 0, stream>>>(
+    unique_offsets, num_unique, unique_counts);
+  return cudaGetLastError();
 }
 
 std::size_t get_unique_temp_storage_size(std::size_t num_elements)
@@ -247,6 +271,6 @@ 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) +
+  return scratch_offset + (3 * num_elements + 1U) * sizeof(std::int64_t) +
          num_elements * sizeof(std::int32_t);
 }

perception/autoware_tensorrt_plugins/src/unique_plugin.cpp

@@ -165,19 +165,10 @@ std::int32_t UniquePlugin::enqueue(
 {
   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(
+  return 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);
-
-  cudaMemcpyAsync(
-    reinterpret_cast<std::int64_t *>(outputs[3]), &num_unique_elements, sizeof(std::int64_t),
-    cudaMemcpyHostToDevice, stream);
-
-  cudaStreamSynchronize(stream);
-
-  return 0;
+    reinterpret_cast<std::int64_t *>(outputs[3]), workspace, num_elements, workspace_size, stream);
 }
 
 std::int32_t UniquePlugin::onShapeChange(

perception/autoware_tensorrt_plugins/test/reference_kernels_test.cpp

@@ -261,15 +261,19 @@ TEST(ReferenceKernelsTest, UniqueMatchesCpuReference)
   DeviceBuffer<std::int64_t> unique_d(input.size());
   DeviceBuffer<std::int64_t> inverse_d(input.size());
   DeviceBuffer<std::int64_t> counts_d(input.size());
+  DeviceBuffer<std::int64_t> num_unique_d(1U);
   DeviceBuffer<std::uint8_t> workspace_d(get_unique_workspace_size(input.size()));
 
   copyToDevice(input_d.get(), input);
 
-  const auto num_unique = unique(
-    input_d.get(), unique_d.get(), inverse_d.get(), counts_d.get(), workspace_d.get(), input.size(),
-    get_unique_workspace_size(input.size()), stream.get());
+  ASSERT_EQ(
+    unique(
+      input_d.get(), unique_d.get(), inverse_d.get(), counts_d.get(), num_unique_d.get(),
+      workspace_d.get(), input.size(), get_unique_workspace_size(input.size()), stream.get()),
+    cudaSuccess);
   ASSERT_EQ(cudaStreamSynchronize(stream.get()), cudaSuccess);
 
+  const auto num_unique = copyToHost(num_unique_d.get(), 1U).front();
   const auto unique_values = copyToHost(unique_d.get(), static_cast<std::size_t>(num_unique));
   const auto inverse_indices = copyToHost(inverse_d.get(), input.size());
   const auto counts = copyToHost(counts_d.get(), static_cast<std::size_t>(num_unique));