append indexsearch

gsampler9 · Jul 31, 2023 · d1f72ad · d1f72ad
1 parent 351b9da
commit d1f72ad
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/src/pybind.cc b/src/pybind.cc
@@ -74,6 +74,7 @@ TORCH_LIBRARY(gs_ops, m) {
   m.def("_CAPI_BatchSplitByOffset", &gs::impl::batch::SplitByOffset);
   m.def("_CAPI_BatchIndptrSplitByOffset",
         &gs::impl::batch::SplitIndptrByOffsetCUDA);
+  m.def("_CAPI_IndexSearch", &IndexSearch);
 }
 
 namespace gs {}
diff --git a/src/tensor_ops.cc b/src/tensor_ops.cc
@@ -27,4 +27,14 @@ std::tuple<torch::Tensor, torch::Tensor> BatchListSampling(
                                                  range);
 }
 
+torch::Tensor IndexSearch(torch::Tensor origin_data, torch::Tensor keys) {
+  torch::Tensor key_buffer, value_buffer;
+
+  std::tie(key_buffer, value_buffer) =
+      impl::IndexHashMapInsertCUDA(origin_data);
+  torch::Tensor result =
+      impl::IndexHashMapSearchCUDA(key_buffer, value_buffer, keys);
+  return result;
+}
+
 }  // namespace gs
diff --git a/src/tensor_ops.h b/src/tensor_ops.h
@@ -26,5 +26,7 @@ std::tuple<torch::Tensor, torch::Tensor> BatchListSamplingProbs(
 std::tuple<torch::Tensor, torch::Tensor> BatchListSampling(int64_t num_picks,
                                                            bool replace,
                                                            torch::Tensor range);
+
+torch::Tensor IndexSearch(torch::Tensor origin_data, torch::Tensor keys);
 }  // namespace gs
 #endif