Define a selection primtive API

cpp/src/prims/per_v_random_select_and_transform_outgoine_e.cuh

@@ -0,0 +1,170 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+namespace cugraph {
+
+/**
+ * @brief Randomly select and transform the input (tagged-)vertices' outgoing edges with biases.
+ *
+ * @tparam GraphViewType Type of the passed non-owning graph object.
+ * @tparam VertexFrontierBucketType Type of the vertex frontier bucket class which abstracts the
+ * current (tagged-)vertex frontier.
+ * @tparam EdgeSrcValueInputWrapper Type of the wrapper for edge source property values.
+ * @tparam EdgeDstValueInputWrapper Type of the wrapper for edge destination property values.
+ * @tparam EdgeBiasOp Type of the quaternary (or quinary) edge operator to set-up selection bias
+ * values.
+ * @tparam EdgeOp Type of the quaternary (or quinary) edge operator.
+ * @tparam T Type of the selected and transformed edge output values.
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param graph_view Non-owning graph object.
+ * @param frontier VertexFrontierBucketType class object to store the (tagged-)vertex list to sample
+ * outgoing edges.
+ * @param edge_src_value_input Wrapper used to access source input property values (for the edge
+ * sources assigned to this process in multi-GPU). Use either cugraph::edge_src_property_t::view()
+ * (if @p e_op needs to access source property values) or cugraph::edge_src_dummy_property_t::view()
+ * (if @p e_op does not access source property values). Use update_edge_src_property to fill the
+ * wrapper.
+ * @param edge_dst_value_input Wrapper used to access destination input property values (for the
+ * edge destinations assigned to this process in multi-GPU). Use either
+ * cugraph::edge_dst_property_t::view() (if @p e_op needs to access destination property values) or
+ * cugraph::edge_dst_dummy_property_t::view() (if @p e_op does not access destination property
+ * values). Use update_edge_dst_property to fill the wrapper.
+ * @param e_bias_op Quaternary (or quinary) operator takes edge source, edge destination, (optional
+ * edge weight), property values for the source, and property values for the destination and returns
+ * a graph weight type bias value to be used in biased random selection.
+ * @param e_op Quaternary (or quinary) operator takes edge source, edge destination, (optional edge
+ * weight), property values for the source, and property values for the destination and returns a
+ * value to be collected in the output. This function is called only for the selected edges.
+ * @param K Number of outgoing edges to select per (tagged-)vertex.
+ * @param with_replacement A flag to specify whether a single outgoing edge can be selected multiple
+ * times (if @p with_replacement = true) or can be selected only once (if @p with_replacement =
+ * false).
+ * @param invalid_value If @p invalid_value.has_value() is true, this value is used to fill the
+ * output vector for the zero out-degree vertices (if @p with_replacement = true) or the vertices
+ * with their out-degrees smaller than @p K (if @p with_replacement = false). If @p
+ * invalid_value.has_value() is false, fewer than @p K values can be returned for the vertices with
+ * fewer than @p K selected edges. See the return value section for additional details.
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return std::tuple Tuple of an optional offset vector of type
+ * std::optional<rmm::device_uvector<size_t>> and a dataframe buffer storing the output values of
+ * type @p T from the selected edges. If @p invalid_value is std::nullopt, the offset vector is
+ * valid and has the size of @p frontier.size() + 1. If @p invalid_value.has_value() is true,
+ * std::nullopt is returned (the dataframe buffer will store @p frontier.size() * @p K elements).
+ */
+template <typename GraphViewType,
+          typename VertexFrontierBucketType,
+          typename EdgeSrcValueInputWrapper,
+          typename EdgeDstValueInputWrapper,
+          typename EdgeBiasOp,
+          typename EdgeOp,
+          typename T>
+std::tuple<std::optional<rmm::device_uvector<size_t>>,
+           decltype(allocate_dataframe_buffer<T>(size_t{0}, rmm::cuda_stream_view{}))>
+per_v_random_select_and_transform_outgoing_e(raft::handle_t const& handle,
+                                             GraphViewType const& graph_view,
+                                             VertexFrontierBucketType const& frontier,
+                                             EdgeSrcValueInputWrapper edge_src_value_input,
+                                             EdgeDstValueInputWrapper edge_dst_value_input,
+#if 0  // FIXME: This will be necessary to include edge IDs in the output.
+       // Primitives API should be updated to support this in a consistent way.
+                               EdgeValueInputWrapper egde_value_input,
+#endif
+                                             EdgeBiasOp e_bias_op,
+                                             EdgeOp e_op,
+                                             size_t K,
+                                             bool with_replacement,
+                                             std::optional<T> invalid_value,
+                                             bool do_expensive_check = false)
+{
+  static_assert(false, "unimplemented.");
+}
+
+/**
+ * @brief Randomly select and transform the input (tagged-)vertices' outgoing edges.
+ *
+ * This function assumes that every outgoing edge of a given vertex has the same odd to be selected
+ * (uniform neighbor sampling).
+ *
+ * @tparam GraphViewType Type of the passed non-owning graph object.
+ * @tparam VertexFrontierBucketType Type of the vertex frontier bucket class which abstracts the
+ * current (tagged-)vertex frontier.
+ * @tparam EdgeSrcValueInputWrapper Type of the wrapper for edge source property values.
+ * @tparam EdgeDstValueInputWrapper Type of the wrapper for edge destination property values.
+ * @tparam EdgeOp Type of the quaternary (or quinary) edge operator.
+ * @tparam T Type of the selected and transformed edge output values.
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param graph_view Non-owning graph object.
+ * @param frontier VertexFrontierBucketType class object to store the (tagged-)vertex list to sample
+ * outgoing edges.
+ * @param edge_src_value_input Wrapper used to access source input property values (for the edge
+ * sources assigned to this process in multi-GPU). Use either cugraph::edge_src_property_t::view()
+ * (if @p e_op needs to access source property values) or cugraph::edge_src_dummy_property_t::view()
+ * (if @p e_op does not access source property values). Use update_edge_src_property to fill the
+ * wrapper.
+ * @param edge_dst_value_input Wrapper used to access destination input property values (for the
+ * edge destinations assigned to this process in multi-GPU). Use either
+ * cugraph::edge_dst_property_t::view() (if @p e_op needs to access destination property values) or
+ * cugraph::edge_dst_dummy_property_t::view() (if @p e_op does not access destination property
+ * values). Use update_edge_dst_property to fill the wrapper.
+ * @param e_op Quaternary (or quinary) operator takes edge source, edge destination, (optional edge
+ * weight), property values for the source, and property values for the destination and returns a
+ * value to be collected in the output. This function is called only for the selected edges.
+ * @param K Number of outgoing edges to select per (tagged-)vertex.
+ * @param with_replacement A flag to specify whether a single outgoing edge can be selected multiple
+ * times (if @p with_replacement = true) or can be selected only once (if @p with_replacement =
+ * false).
+ * @param invalid_value If @p invalid_value.has_value() is true, this value is used to fill the
+ * output vector for the zero out-degree vertices (if @p with_replacement = true) or the vertices
+ * with their out-degrees smaller than @p K (if @p with_replacement = false). If @p
+ * invalid_value.has_value() is false, fewer than @p K values can be returned for the vertices with
+ * fewer than @p K selected edges. See the return value section for additional details.
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return std::tuple Tuple of an optional offset vector of type
+ * std::optional<rmm::device_uvector<size_t>> and a dataframe buffer storing the output values of
+ * type @p T from the selected edges. If @p invalid_value is std::nullopt, the offset vector is
+ * valid and has the size of @p frontier.size() + 1. If @p invalid_value.has_value() is true,
+ * std::nullopt is returned (the dataframe buffer will store @p frontier.size() * @p K elements).
+ */
+template <typename GraphViewType,
+          typename VertexFrontierBucketType,
+          typename EdgeSrcValueInputWrapper,
+          typename EdgeDstValueInputWrapper,
+          typename EdgeOp,
+          typename T>
+std::tuple<std::optional<rmm::device_uvector<size_t>>,
+           decltype(allocate_dataframe_buffer<T>(size_t{0}, rmm::cuda_stream_view{}))>
+per_v_random_select_and_transform_outgoing_e(raft::handle_t const& handle,
+                                             GraphViewType const& graph_view,
+                                             VertexFrontierBucketType const& frontier,
+                                             EdgeSrcValueInputWrapper edge_src_value_input,
+                                             EdgeDstValueInputWrapper edge_dst_value_input,
+#if 0  // FIXME: This will be necessary to include edge IDs in the output.
+       // Primitives API should be updated to support this in a consistent way.
+                               EdgeValueInputWrapper egde_value_input,
+#endif
+                                             EdgeOp e_op,
+                                             size_t K,
+                                             bool with_replacement,
+                                             std::optional<T> invalid_value,
+                                             bool do_expensive_check = false)
+{
+  static_assert(false, "unimplemented.");
+}
+
+}  // namespace cugraph