Address review re: distance API

ahendriksen · Apr 18, 2023 · 3ea52b8 · 3ea52b8
1 parent 65c1cba
commit 3ea52b8
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 91 deletions.
diff --git a/cpp/include/raft/distance/detail/pairwise_matrix/dispatch-ext.cuh b/cpp/include/raft/distance/detail/pairwise_matrix/dispatch-ext.cuh
@@ -65,6 +65,14 @@ void pairwise_matrix_dispatch(OpT distance_op,
       cudaStream_t stream,                                                             \
       bool is_row_major)
 
+/*
+ * Hierarchy of instantiations:
+ *
+ * This file defines extern template instantiations of the distance kernels. The
+ * instantiation of the public API is handled in raft/distance/distance-ext.cuh.
+ *
+ * After adding an instance here, make sure to also add the instance there.
+ */
 instantiate_raft_distance_detail_pairwise_matrix_dispatch(
   raft::distance::detail::ops::canberra_distance_op, float, float, float, raft::identity_op, int);
 instantiate_raft_distance_detail_pairwise_matrix_dispatch(

diff --git a/cpp/include/raft/distance/distance-ext.cuh b/cpp/include/raft/distance/distance-ext.cuh
@@ -331,6 +331,18 @@ void pairwise_distance(raft::resources const& handle,
 
 #endif  // RAFT_EXPLICIT_INSTANTIATE
 
+/*
+ * Hierarchy of instantiations:
+ *
+ * This file defines the extern template instantiations for the public API of
+ * raft::distance. To improve compile times, the extern template instantiation
+ * of the distance kernels is handled in
+ * distance/detail/pairwise_matrix/dispatch-ext.cuh.
+ *
+ * After adding an instance here, make sure to also add the instance to
+ * dispatch-ext.cuh and the corresponding .cu files.
+ */
+
 #define instantiate_raft_distance_distance(DT, DataT, AccT, OutT, FinalLambda, IdxT)       \
   extern template void raft::distance::distance<DT, DataT, AccT, OutT, FinalLambda, IdxT>( \
     raft::resources const& handle,                                                         \

diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch.cu b/cpp/src/distance/detail/pairwise_matrix/dispatch.cu
diff --git a/cpp/src/distance/distance.cu b/cpp/src/distance/distance.cu
@@ -16,6 +16,15 @@
 
 #include <raft/distance/distance-inl.cuh>
 
+/*
+ * Hierarchy of instantiations:
+ *
+ * This file defines the template instantiations for the public API of
+ * raft::distance. To improve compile times, the compilation of the distance
+ * kernels is handled in distance/detail/pairwise_matrix/dispatch_*.cu.
+ *
+ */
+
 #define instantiate_raft_distance_distance(DT, DataT, AccT, OutT, FinalLambda, IdxT) \
   template void raft::distance::distance<DT, DataT, AccT, OutT, FinalLambda, IdxT>(  \
     raft::resources const& handle,                                                   \