[GraphBolt][CUDA] gb.isin implementation (#6829)

graphbolt/src/cuda/isin.cu

@@ -0,0 +1,38 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *  Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ * @file cuda/isin.cu
+ * @brief IsIn operator implementation on CUDA.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/binary_search.h>
+
+#include <cub/cub.cuh>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+torch::Tensor IsIn(torch::Tensor elements, torch::Tensor test_elements) {
+  auto sorted_test_elements = Sort(test_elements).first;
+  auto allocator = cuda::GetAllocator();
+  auto stream = cuda::GetCurrentStream();
+  const auto exec_policy = thrust::cuda::par_nosync(allocator).on(stream);
+  auto result = torch::empty_like(elements, torch::kBool);
+
+  AT_DISPATCH_INTEGRAL_TYPES(
+      elements.scalar_type(), "IsInOperation", ([&] {
+        thrust::binary_search(
+            exec_policy, sorted_test_elements.data_ptr<scalar_t>(),
+            sorted_test_elements.data_ptr<scalar_t>() +
+                sorted_test_elements.size(0),
+            elements.data_ptr<scalar_t>(),
+            elements.data_ptr<scalar_t>() + elements.size(0),
+            result.data_ptr<bool>());
+      }));
+  return result;
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/isin.cc

@@ -5,8 +5,12 @@
  * @brief Isin op.
  */
 
+#include <graphbolt/cuda_ops.h>
 #include <graphbolt/isin.h>
 
+#include "./macro.h"
+#include "./utils.h"
+
 namespace {
 static constexpr int kSearchGrainSize = 4096;
 }  // namespace
@@ -14,7 +18,7 @@ static constexpr int kSearchGrainSize = 4096;
 namespace graphbolt {
 namespace sampling {
 
-torch::Tensor IsIn(
+torch::Tensor IsInCPU(
     const torch::Tensor& elements, const torch::Tensor& test_elements) {
   torch::Tensor sorted_test_elements;
   std::tie(sorted_test_elements, std::ignore) = test_elements.sort(
@@ -41,5 +45,17 @@ torch::Tensor IsIn(
       }));
   return result;
 }
+
+torch::Tensor IsIn(
+    const torch::Tensor& elements, const torch::Tensor& test_elements) {
+  if (utils::is_accessible_from_gpu(elements) &&
+      utils::is_accessible_from_gpu(test_elements)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "IsInOperation",
+        { return ops::IsIn(elements, test_elements); });
+  } else {
+    return IsInCPU(elements, test_elements);
+  }
+}
 }  // namespace sampling
 }  // namespace graphbolt