Skip to content

pyg::subgraph CUDA implementation #42

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rusty1s wants to merge 11 commits into
base: main
Choose a base branch
from
Open

pyg::subgraph CUDA implementation #42

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
07f6a85
Update
rusty1s May 3, 2022
058c0bc
update
rusty1s May 3, 2022
f85ff1e
update
rusty1s May 3, 2022
4aa86b9
initial commit
rusty1s May 3, 2022
4bb8087
update
rusty1s May 3, 2022
ef68b2b
update
rusty1s May 3, 2022
377e261
update
rusty1s May 4, 2022
9b5d9f7
typo
rusty1s May 4, 2022
9b4e6e6
changelog
rusty1s May 4, 2022
bf23274
update
rusty1s May 4, 2022
9113818
Merge branch 'master' into subgraph_cuda
ZenoTan May 5, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 23 additions & 3 deletions pyg_lib/csrc/sampler/cuda/subgraph_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ namespace sampler {

namespace {

#define FULL_MASK 0xffffffff

Comment thread
rusty1s marked this conversation as resolved.
Outdated
template <typename scalar_t>
__global__ void subgraph_deg_kernel_impl(
const scalar_t* __restrict__ rowptr_data,
Expand All @@ -16,7 +18,24 @@ __global__ void subgraph_deg_kernel_impl(
const scalar_t* __restrict__ to_local_node_data,
scalar_t* __restrict__ out_data,
int64_t num_nodes) {
CUDA_1D_KERNEL_LOOP(scalar_t, i, 32 * num_nodes) {}
CUDA_1D_KERNEL_LOOP(scalar_t, thread_idx, 32 * num_nodes) {
scalar_t i = thread_idx >> 5; // thread_idx / 32
scalar_t lane = thread_idx & (32 - 1); // thread_idx % 32

auto v = nodes_data[i];

scalar_t deg = 0;
for (scalar_t j = rowptr_data[v] + lane; j < rowptr_data[v + 1]; j += 32) {
if (to_local_node_data[col_data[j]] >= 0) // contiguous access
deg++;
}

for (scalar_t offset = 16; offset > 0; offset /= 2) // warp-level reduction
deg += __shfl_down_sync(FULL_MASK, deg, offset);

if (lane == 0)
out_data[i] = deg;
}
}

std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph_kernel(
Expand All @@ -32,7 +51,7 @@ std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph_kernel(

// We maintain a O(N) vector to map global node indices to local ones.
// TODO Can we do this without O(N) storage requirement?
const auto to_local_node = nodes.new_empty({rowptr.size(0) - 1});
const auto to_local_node = nodes.new_full({rowptr.size(0) - 1}, -1);
const auto arange = at::arange(nodes.size(0), nodes.options());
to_local_node.index_copy_(/*dim=*/0, nodes, arange);

Expand All @@ -48,6 +67,7 @@ std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph_kernel(
const auto to_local_node_data = to_local_node.data_ptr<scalar_t>();
auto deg_data = deg.data_ptr<scalar_t>();

// Compute induced subgraph degree, parallelize with 32 threads per node:

@yaoyaowd yaoyaowd May 3, 2022

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm actually not sure if it is necessary to parallelize with 32 threads per nodes. Most of the time we are dealing with sparse data and a lot of threads will not go into for loop.

If you are looking for extreme performance, you can bundle to_local_node_data and col_data into one iterator structure and use this function. I haven't seen any better performance than it in the past.
https://nvlabs.github.io/cub/structcub_1_1_device_segmented_reduce.html#a4854a13561cb66d46aa617aab16b8825

@rusty1s rusty1s May 4, 2022

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you have an example of bundling to_local_node_data and col_data into one iterator structure? This looks really interesting.

I am okay with dropping the warp-level parallelism for now, but we will lose the contiguous access to col_data, and probably under-utilize the number of threads available on modern GPUs.

@rusty1s rusty1s May 4, 2022

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On a second look, this doesn't seem possible since col_data refers to edges, while to_local_node_data refers to nodes, while we actually want do the compute across the number of nodes in the induced subgraph.

subgraph_deg_kernel_impl<<<pyg::utils::blocks(32 * nodes.size(0)),
pyg::utils::threads(), 0, stream>>>(
rowptr_data, col_data, nodes_data, to_local_node_data, deg_data,
Expand All @@ -57,7 +77,7 @@ std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>> subgraph_kernel(
at::cumsum_out(tmp, deg, /*dim=*/0);
});

return std::make_tuple(to_local_node, deg, rowptr);
return std::make_tuple(out_rowptr, deg, deg);
}

} // namespace
Expand Down
6 changes: 4 additions & 2 deletions test/csrc/sampler/test_subgraph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,11 @@ TEST(SubgraphTest, BasicAssertions) {
/*col=*/std::get<1>(graph), nodes);

std::cout << std::get<0>(out) << std::endl;
std::cout << std::get<1>(out) << std::endl;
std::cout << std::get<2>(out).value() << std::endl;

/* auto expected_rowptr = at::tensor({0, 1, 3, 5, 6}, options); */
/* EXPECT_TRUE(at::equal(std::get<0>(out), expected_rowptr)); */
auto expected_rowptr = at::tensor({0, 1, 3, 5, 6}, options);
EXPECT_TRUE(at::equal(std::get<0>(out), expected_rowptr));
/* auto expected_col = at::tensor({1, 0, 2, 1, 3, 2}, options); */
/* EXPECT_TRUE(at::equal(std::get<1>(out), expected_col)); */
/* auto expected_edge_id = at::tensor({3, 4, 5, 6, 7, 8}, options); */
Expand Down