Skip to content

Commit

Permalink
code refactor for re-use for other primitives
Browse files Browse the repository at this point in the history
  • Loading branch information
@seunghwak
seunghwak committed Jan 24, 2024
1 parent 304b5c0 commit b2308d6
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 86 deletions.
60 changes: 60 additions & 0 deletions cpp/src/prims/detail/prim_functors.cuh
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
/*
* Copyright (c) 2024, 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

#include <cugraph/edge_partition_device_view.cuh>

namespace cugraph {

namespace detail {

template <typename GraphViewType,
typename EdgePartitionSrcValueInputWrapper,
typename EdgePartitionDstValueInputWrapper,
typename EdgePartitionEdgeValueInputWrapper,
typename EdgeOp>
struct call_e_op_t {
edge_partition_device_view_t<typename GraphViewType::vertex_type,
typename GraphViewType::edge_type,
GraphViewType::is_multi_gpu> const& edge_partition{};
EdgePartitionSrcValueInputWrapper const& edge_partition_src_value_input{};
EdgePartitionDstValueInputWrapper const& edge_partition_dst_value_input{};
EdgePartitionEdgeValueInputWrapper const& edge_partition_e_value_input{};
EdgeOp const& e_op{};
typename GraphViewType::vertex_type major{};
typename GraphViewType::vertex_type major_offset{};
typename GraphViewType::vertex_type const* indices{nullptr};
typename GraphViewType::edge_type edge_offset{};

__device__ auto operator()(typename GraphViewType::edge_type i) const
{
auto minor = indices[i];
auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor);
auto src = GraphViewType::is_storage_transposed ? minor : major;
auto dst = GraphViewType::is_storage_transposed ? major : minor;
auto src_offset = GraphViewType::is_storage_transposed ? minor_offset : major_offset;
auto dst_offset = GraphViewType::is_storage_transposed ? major_offset : minor_offset;
return e_op(src,
dst,
edge_partition_src_value_input.get(src_offset),
edge_partition_dst_value_input.get(dst_offset),
edge_partition_e_value_input.get(edge_offset + i));
}
};

} // namespace detail

} // namespace cugraph
139 changes: 53 additions & 86 deletions cpp/src/prims/per_v_transform_reduce_incoming_outgoing_e.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#pragma once

#include <detail/graph_partition_utils.cuh>
#include <prims/detail/prim_functors.cuh>
#include <prims/fill_edge_src_dst_property.cuh>
#include <prims/property_op_utils.cuh>
#include <prims/reduce_op.cuh>
Expand Down Expand Up @@ -63,40 +64,6 @@ namespace detail {

int32_t constexpr per_v_transform_reduce_e_kernel_block_size = 512;

template <typename GraphViewType,
typename EdgePartitionSrcValueInputWrapper,
typename EdgePartitionDstValueInputWrapper,
typename EdgePartitionEdgeValueInputWrapper,
typename EdgeOp>
struct per_v_transform_reduce_call_e_op_t {
edge_partition_device_view_t<typename GraphViewType::vertex_type,
typename GraphViewType::edge_type,
GraphViewType::is_multi_gpu> const& edge_partition{};
EdgePartitionSrcValueInputWrapper const& edge_partition_src_value_input{};
EdgePartitionDstValueInputWrapper const& edge_partition_dst_value_input{};
EdgePartitionEdgeValueInputWrapper const& edge_partition_e_value_input{};
EdgeOp const& e_op{};
typename GraphViewType::vertex_type major{};
typename GraphViewType::vertex_type major_offset{};
typename GraphViewType::vertex_type const* indices{nullptr};
typename GraphViewType::edge_type edge_offset{};

__device__ auto operator()(typename GraphViewType::edge_type i) const
{
auto minor = indices[i];
auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor);
auto src = GraphViewType::is_storage_transposed ? minor : major;
auto dst = GraphViewType::is_storage_transposed ? major : minor;
auto src_offset = GraphViewType::is_storage_transposed ? minor_offset : major_offset;
auto dst_offset = GraphViewType::is_storage_transposed ? major_offset : minor_offset;
return e_op(src,
dst,
edge_partition_src_value_input.get(src_offset),
edge_partition_dst_value_input.get(dst_offset),
edge_partition_e_value_input.get(edge_offset + i));
}
};

template <typename vertex_t,
typename edge_t,
bool multi_gpu,
Expand Down Expand Up @@ -222,19 +189,19 @@ __global__ void per_v_transform_reduce_e_hypersparse(
thrust::tie(indices, edge_offset, local_degree) =
edge_partition.local_edges(static_cast<vertex_t>(major_idx));

auto call_e_op = per_v_transform_reduce_call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};
auto call_e_op = call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};

if (edge_partition_e_mask) {
auto transform_op =
Expand Down Expand Up @@ -322,19 +289,19 @@ __global__ void per_v_transform_reduce_e_low_degree(
thrust::tie(indices, edge_offset, local_degree) =
edge_partition.local_edges(static_cast<vertex_t>(major_offset));

auto call_e_op = per_v_transform_reduce_call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};
auto call_e_op = call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};

if (edge_partition_e_mask) {
auto transform_op =
Expand Down Expand Up @@ -429,19 +396,19 @@ __global__ void per_v_transform_reduce_e_mid_degree(
edge_t local_degree{};
thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset);

auto call_e_op = per_v_transform_reduce_call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};
auto call_e_op = call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};

[[maybe_unused]] auto reduced_e_op_result =
lane_id == 0 ? init : identity_element; // relevant only if update_major == true
Expand Down Expand Up @@ -540,19 +507,19 @@ __global__ void per_v_transform_reduce_e_high_degree(
edge_t local_degree{};
thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset);

auto call_e_op = per_v_transform_reduce_call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};
auto call_e_op = call_e_op_t<GraphViewType,
EdgePartitionSrcValueInputWrapper,
EdgePartitionDstValueInputWrapper,
EdgePartitionEdgeValueInputWrapper,
EdgeOp>{edge_partition,
edge_partition_src_value_input,
edge_partition_dst_value_input,
edge_partition_e_value_input,
e_op,
major,
major_offset,
indices,
edge_offset};

[[maybe_unused]] auto reduced_e_op_result =
threadIdx.x == 0 ? init : identity_element; // relevant only if update_major == true
Expand Down

0 comments on commit b2308d6

Please sign in to comment.