pub_soft/dgl
3
0
Fork 0
mirror of https://github.com/dmlc/dgl.git synced 2026-06-03 19:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

[GraphBolt][CUDA] Remove unused num_bits argument from UniqueAndCompact (#7764)

Browse Source
This commit is contained in:
Muhammed Fatih BALIN
2024-09-02 17:50:08 -04:00
committed by GitHub
parent 6cd7fe99c9
commit 9a86a665ef
2 changed files with 7 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
graphbolt/include/graphbolt/cuda_ops.h
View File

@@ -299,7 +299,7 @@ torch::Tensor IndptrEdgeIdsImpl(
*/
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> UniqueAndCompact(
const torch::Tensor src_ids, const torch::Tensor dst_ids,
const torch::Tensor unique_dst_ids, int num_bits = 0);
const torch::Tensor unique_dst_ids);
/**
* @brief Batched version of UniqueAndCompact. The ith element of the return
@@ -310,7 +310,7 @@ std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
UniqueAndCompactBatched(
const std::vector<torch::Tensor>& src_ids,
const std::vector<torch::Tensor>& dst_ids,
const std::vector<torch::Tensor>& unique_dst_ids, int num_bits = 0);
const std::vector<torch::Tensor>& unique_dst_ids);
} // namespace ops
} // namespace graphbolt

12
graphbolt/src/cuda/unique_and_compact_impl.cu
View File

@@ -61,7 +61,7 @@ std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
UniqueAndCompactBatchedSortBased(
const std::vector<torch::Tensor>& src_ids,
const std::vector<torch::Tensor>& dst_ids,
const std::vector<torch::Tensor>& unique_dst_ids, int num_bits) {
const std::vector<torch::Tensor>& unique_dst_ids, int num_bits = 0) {
auto allocator = cuda::GetAllocator();
auto stream = cuda::GetCurrentStream();
auto scalar_type = src_ids.at(0).scalar_type();
@@ -276,7 +276,7 @@ std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
UniqueAndCompactBatched(
const std::vector<torch::Tensor>& src_ids,
const std::vector<torch::Tensor>& dst_ids,
const std::vector<torch::Tensor>& unique_dst_ids, int num_bits) {
const std::vector<torch::Tensor>& unique_dst_ids) {
if (cuda::compute_capability() >= 70) {
// Utilizes a hash table based implementation, the mapped id of a vertex
// will be monotonically increasing as the first occurrence index of it in
@@ -287,15 +287,13 @@ UniqueAndCompactBatched(
// Utilizes a sort based algorithm, the mapped id of a vertex part of the
// src_ids but not part of the unique_dst_ids will be monotonically increasing
// as the actual vertex id increases. Thus, it is deterministic.
return UniqueAndCompactBatchedSortBased(
src_ids, dst_ids, unique_dst_ids, num_bits);
return UniqueAndCompactBatchedSortBased(src_ids, dst_ids, unique_dst_ids);
}
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> UniqueAndCompact(
const torch::Tensor src_ids, const torch::Tensor dst_ids,
const torch::Tensor unique_dst_ids, int num_bits) {
return UniqueAndCompactBatched(
{src_ids}, {dst_ids}, {unique_dst_ids}, num_bits)[0];
const torch::Tensor unique_dst_ids) {
return UniqueAndCompactBatched({src_ids}, {dst_ids}, {unique_dst_ids})[0];
}
} // namespace ops