Remove usage of TRHUTS host&device vector

- as there is no guarantee that THRUST host and device vectors can be used in host-only programs to remove their usage from DALI Signed-off-by: Janusz Lisiecki <[email protected]>
NVIDIA · Apr 22, 2024 · b2bc44b · b2bc44b
1 parent 8801462
commit b2bc44b
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 49 deletions.
diff --git a/dali/c_api/operator_trace_test.cc b/dali/c_api/operator_trace_test.cc
@@ -13,12 +13,13 @@
 // limitations under the License.
 
 #include <gtest/gtest.h>
-#include <thrust/host_vector.h>
-#include <thrust/device_vector.h>
 #include <memory>
 #include "dali/c_api.h"
 #include "dali/pipeline/pipeline.h"
 #include "dali/test/dali_test_config.h"
+#include "dali/test/test_tensors.h"
+#include "dali/test/tensor_test_utils.h"
+
 
 namespace dali::test {
 
@@ -206,18 +207,12 @@ class OperatorTraceTestExternalInput : public OperatorTraceTest {
 namespace {
 
 template<typename T>
-thrust::host_vector<T> random_vector_cpu(std::mt19937 &mt, size_t size) {
-  std::vector<T> ret(size);
+kernels::TestTensorList<T> random_vector(std::mt19937 &mt, size_t size) {
+  kernels::TestTensorList<T> ret;
+  ret.reshape(uniform_list_shape(1, {size}));
   std::uniform_int_distribution<T> dist{0, 255};
   auto gen = [&]() { return dist(mt); };
-  std::generate(ret.begin(), ret.end(), gen);
-  return ret;
-}
-
-
-template<typename T>
-thrust::device_vector<T> random_vector_gpu(std::mt19937 &mt, size_t size) {
-  thrust::device_vector<T> ret = random_vector_cpu<T>(mt, size);
+  Fill(ret.cpu(), gen);
   return ret;
 }
 
@@ -238,9 +233,9 @@ TEST_P(OperatorTraceTestExternalInput, OperatorTraceTestExternalInput) {
     ASSERT_GE(feed_count_cpu, 1);
     for (int i = 0; i < feed_count_cpu; i++) {
       size_t sample_size = 42;
-      auto in_data = random_vector_cpu<uint8_t>(rng, sample_size * batch_size_);
+      auto in_data = random_vector<uint8_t>(rng, sample_size * batch_size_);
       std::vector<int64_t> shapes(batch_size_, sample_size);
-      daliSetExternalInput(&h, "OP_TRACE_IN_CPU", device_type_t::CPU, in_data.data(),
+      daliSetExternalInput(&h, "OP_TRACE_IN_CPU", device_type_t::CPU, in_data.cpu().tensor_data(0),
                            dali_data_type_t::DALI_UINT8, shapes.data(), 1, nullptr,
                            DALI_ext_default);
     }
@@ -250,10 +245,10 @@ TEST_P(OperatorTraceTestExternalInput, OperatorTraceTestExternalInput) {
     ASSERT_GE(feed_count_gpu, 1);
     for (int i = 0; i < feed_count_gpu; i++) {
       int sample_size = 42;
-      auto in_data = random_vector_gpu<uint8_t>(rng, sample_size * batch_size_);
+      auto in_data = random_vector<uint8_t>(rng, sample_size * batch_size_);
       std::vector<int64_t> shapes(batch_size_, sample_size);
       daliSetExternalInput(&h, "OP_TRACE_IN_GPU", device_type_t::GPU,
-                           thrust::raw_pointer_cast(in_data.data()), dali_data_type_t::DALI_UINT8,
+                           in_data.gpu().tensor_data(0), dali_data_type_t::DALI_UINT8,
                            shapes.data(), 1, nullptr, DALI_ext_default);
     }
 

diff --git a/dali/operators/input/input_operator_test.cu b/dali/operators/input/input_operator_test.cu
@@ -13,15 +13,15 @@
 // limitations under the License.
 
 #include <gtest/gtest.h>
-#include <thrust/host_vector.h>
-#include <thrust/device_vector.h>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 #include "dali/c_api.h"
 #include "dali/pipeline/pipeline.h"
 #include "dali/test/dali_test_config.h"
+#include "dali/test/test_tensors.h"
+#include "dali/test/tensor_test_utils.h"
 
 namespace dali::test {
 
@@ -55,16 +55,6 @@ InputOperatorMixedTestParam input_operator_test_params_pipelined_executor_separa
         {3, 2, true, true,  false},
 };
 
-
-template<typename T>
-thrust::host_vector<T> random_vector_cpu(std::mt19937 &mt, size_t size) {
-  thrust::host_vector<T> cpu(size);
-  std::uniform_int_distribution<T> dist{0, 255};
-  auto gen = [&]() { return dist(mt); };
-  thrust::generate(cpu.begin(), cpu.end(), gen);
-  return cpu;
-}
-
 }  // namespace
 
 
@@ -132,13 +122,19 @@ TEST_P(InputOperatorMixedTest, InputOperatorMixedTest) {
   for (int iteration = 0; iteration < n_iterations_; iteration++) {
     int prefetch_depth = daliInputFeedCount(&h, operator_name_.c_str());
     size_t sample_size = 42;
-    thrust::host_vector<int32_t> in_data(sample_size * batch_size_, 2137);
-    thrust::device_vector<int32_t> ref_data = in_data;
+    kernels::TestTensorList<int32_t> in_data;
+    in_data.reshape(uniform_list_shape(batch_size_, {sample_size}));
+    ConstantFill(in_data.cpu(), 0xDEADBEEF);
+    kernels::TestTensorList<int32_t> ref_data;
+    ref_data.reshape(uniform_list_shape(batch_size_, {sample_size}));
+    memcpy(ref_data.cpu().tensor_data(0), in_data.cpu().tensor_data(0),
+           batch_size_ * sample_size * sizeof(in_data.cpu().tensor_data(0)[0]));
 
     // Feed CPU input data.
     for (int i = 0; i < prefetch_depth; i++) {
       std::vector<int64_t> shapes(batch_size_, sample_size);
-      daliSetExternalInput(&h, operator_name_.c_str(), device_type_t::CPU, in_data.data(),
+      daliSetExternalInput(&h, operator_name_.c_str(), device_type_t::CPU,
+                           in_data.cpu().tensor_data(0),
                            dali_data_type_t::DALI_INT32, shapes.data(), 1, nullptr,
                            DALI_ext_force_copy);
     }
@@ -148,11 +144,15 @@ TEST_P(InputOperatorMixedTest, InputOperatorMixedTest) {
     for (int i = 0; i < num_output_batches; i++) {
       daliShareOutput(&h);
       auto sz = daliNumElements(&h, 0);
-      thrust::device_vector<int32_t> out_data(sz);
-      daliOutputCopy(&h, thrust::raw_pointer_cast(out_data.data()), 0, device_type_t::GPU, 0,
+      kernels::TestTensorList<int32_t> out_data;
+      out_data.reshape(uniform_list_shape(batch_size_, {sample_size}));
+      EXPECT_EQ(sz, out_data.gpu().shape.num_elements());
+      daliOutputCopy(&h, out_data.gpu().tensor_data(0), 0, device_type_t::GPU, 0,
                      DALI_ext_force_sync);
-
-      EXPECT_EQ(out_data, ref_data);
+      auto tv_out_data = out_data.cpu();
+      auto tv_ref_data = ref_data.cpu();
+      cudaDeviceSynchronize();
+      Check(tv_out_data, tv_ref_data);
 
       daliOutputRelease(&h);
     }

diff --git a/dali/operators/input/video_input.h b/dali/operators/input/video_input.h
@@ -15,8 +15,6 @@
 #ifndef DALI_OPERATORS_INPUT_VIDEO_INPUT_H_
 #define DALI_OPERATORS_INPUT_VIDEO_INPUT_H_
 
-#include <thrust/host_vector.h>
-#include <thrust/device_vector.h>
 #include <deque>
 #include <string>
 #include <utility>
@@ -85,10 +83,6 @@ template<typename OutBackend, typename PadType>
 struct PadFrameCreator {
   static constexpr bool is_cpu = std::is_same_v<OutBackend, CPUBackend>;
 
-  template<typename T>
-  using container_type =
-          std::conditional_t<is_cpu, thrust::host_vector<T>, thrust::device_vector<T>>;
-
   PadFrameCreator() = default;
 
 

diff --git a/dali/pipeline/operator/builtin/external_source_test.cc b/dali/pipeline/operator/builtin/external_source_test.cc
@@ -13,8 +13,6 @@
 // limitations under the License.
 
 #include <gtest/gtest.h>
-#include <thrust/host_vector.h>
-#include <thrust/device_vector.h>
 #include <fstream>
 #include <memory>
 #include <tuple>
@@ -25,7 +23,8 @@
 #include "dali/pipeline/operator/builtin/external_source.h"
 #include "dali/test/dali_test_config.h"
 #include "dali/c_api.h"
-
+#include "dali/test/test_tensors.h"
+#include "dali/test/tensor_test_utils.h"
 namespace dali {
 
 
@@ -61,17 +60,17 @@ class ExternalSourceBasicTest : public ::testing::Test {
     daliSetExternalInputBatchSize(h, input_name_.c_str(), batch_size_);
     int bytes_per_sample = 10;
     std::vector<int64_t> shapes(batch_size_, bytes_per_sample);
-    thrust::host_vector<uint8_t> data_cpu(batch_size_ * bytes_per_sample);
+    kernels::TestTensorList<uint8_t> data;
+    data.reshape(uniform_list_shape(batch_size_, {bytes_per_sample}));
     assert(bytes_per_sample * batch_size_ < 255);
-    std::iota(data_cpu.begin(), data_cpu.end(), 1);
+    SequentialFill(data.cpu(), 1);
     if constexpr (is_cpu) {
-      daliSetExternalInput(h, input_name_.c_str(), device_type_t::CPU, data_cpu.data(),
+      daliSetExternalInput(h, input_name_.c_str(), device_type_t::CPU, data.cpu().tensor_data(0),
                            dali_data_type_t::DALI_UINT8, shapes.data(), 1, nullptr,
                            DALI_ext_force_copy);
     } else {
-      thrust::device_vector<uint8_t> data_gpu = data_cpu;
       daliSetExternalInput(h, input_name_.c_str(), device_type_t::GPU,
-                           thrust::raw_pointer_cast(data_gpu.data()),
+                           data.gpu().tensor_data(0),
                            dali_data_type_t::DALI_UINT8, shapes.data(), 1, nullptr,
                            DALI_ext_force_copy);
     }