From a35a7af7657d5f1ff3e39a9cbfce290645af76a9 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Fri, 1 Aug 2025 11:07:03 +0200 Subject: [PATCH 1/9] Enhance CUDA support by updating adjacency_matrix and propagate functions for COO graphs --- GNNGraphs/ext/GNNGraphsCUDAExt.jl | 37 +++++++++++++++++++++++++++++++ GNNGraphs/src/query.jl | 8 +------ GNNlib/ext/GNNlibCUDAExt.jl | 10 +++++++-- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/GNNGraphs/ext/GNNGraphsCUDAExt.jl b/GNNGraphs/ext/GNNGraphsCUDAExt.jl index af9e9f820..ae7d773ce 100644 --- a/GNNGraphs/ext/GNNGraphsCUDAExt.jl +++ b/GNNGraphs/ext/GNNGraphsCUDAExt.jl @@ -5,13 +5,29 @@ using Random, Statistics, LinearAlgebra using GNNGraphs using GNNGraphs: COO_T, ADJMAT_T, SPARSE_T using SparseArrays +using Graphs const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix} +const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{Nothing, AnyCuArray}} # Query GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1)) +function Graphs.adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out, + weighted = true) + @debug "Using CUDA adjacency_matrix for GNNGraph" + if !g.is_coalesced + # Revisit after + # https://github.com/JuliaGPU/CUDA.jl/issues/1113 + A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesce, construction of sparse matrix is slow + else + A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted) + end + @assert size(A) == (n, n) + return dir == :out ? A : A' +end + # Transform GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz) @@ -35,5 +51,26 @@ function sort_edge_index(u::AnyCuArray, v::AnyCuArray) sort_edge_index(u, v) |> dev end +# Convert + +function GNNGraphs.to_sparse(coo::CUDA_COO_T, T = nothing; dir = :out, num_nodes = nothing, + weighted = true) + s, t, eweight = coo + @debug "Using CUDA to_sparse for COO" + T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T + + if eweight === nothing || !weighted + eweight = fill!(similar(s, T), 1) + end + + num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes + A = CUDA.CUSPARSE.CuSparseMatrixCOO{T,eltype(s)}(s, t, eweight, (num_nodes, num_nodes)) # create sparse matrix in COO format + + num_edges::Int = nnz(A) + if eltype(A) != T + A = T.(A) + end + return A, num_nodes, num_edges +end end #module diff --git a/GNNGraphs/src/query.jl b/GNNGraphs/src/query.jl index e2e5f9bdb..482502b4a 100644 --- a/GNNGraphs/src/query.jl +++ b/GNNGraphs/src/query.jl @@ -231,13 +231,7 @@ If `weighted=true`, the `A` will contain the edge weights if any, otherwise the """ function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType = eltype(g); dir = :out, weighted = true) - if iscuarray(g.graph[1]) - # Revisit after - # https://github.com/JuliaGPU/CUDA.jl/issues/1113 - A, n, m = to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) - else - A, n, m = to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted) - end + A, n, m = to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted) @assert size(A) == (n, n) return dir == :out ? A : A' end diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index 56a6738e9..529db39c0 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -3,7 +3,7 @@ module GNNlibCUDAExt using CUDA using Random, Statistics, LinearAlgebra using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj -using GNNGraphs: GNNGraph, COO_T, SPARSE_T +using GNNGraphs: GNNGraph, COO_T, SPARSE_T, adjacency_matrix ###### PROPAGATE SPECIALIZATIONS #################### @@ -12,7 +12,13 @@ using GNNGraphs: GNNGraph, COO_T, SPARSE_T ## avoid the fast path on gpu until we have better cuda support function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), xi, xj::AnyCuMatrix, e) - propagate((xi, xj, e) -> copy_xj(xi, xj, e), g, +, xi, xj, e) + @debug "Using CUDA propagate for copy_xj" + A = adjacency_matrix(g, eltype(xj); weighted = false) + if g.is_coalesced + return (A' * xj')' # workaround for CUDA issue: https://github.com/JuliaGPU/CUDA.jl/issues/2820 + else + return xj * A + end end ## E_MUL_XJ From 6b3257cb840e4971413d16f0b8de846bc74f80dd Mon Sep 17 00:00:00 2001 From: dferre97 Date: Tue, 12 Aug 2025 09:51:50 +0200 Subject: [PATCH 2/9] Swap edge encoding order in coalesce to fix CUDA.jl issue --- GNNGraphs/src/transform.jl | 3 ++- GNNlib/ext/GNNlibCUDAExt.jl | 7 ++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl index ce8d90b6a..61a11f455 100644 --- a/GNNGraphs/src/transform.jl +++ b/GNNGraphs/src/transform.jl @@ -158,7 +158,8 @@ function Base.coalesce(g::GNNGraph{<:COO_T}; aggr = +) w = get_edge_weight(g) edata = g.edata num_edges = g.num_edges - idxs, idxmax = edge_encoding(s, t, g.num_nodes) + # order by target first and then source as a workaround of CUDA.jl issue: https://github.com/JuliaGPU/CUDA.jl/issues/2820 + idxs, idxmax = edge_encoding(t, s, g.num_nodes) perm = sortperm(idxs) idxs = idxs[perm] diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index 529db39c0..103318772 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -14,11 +14,8 @@ function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), xi, xj::AnyCuMatrix, e) @debug "Using CUDA propagate for copy_xj" A = adjacency_matrix(g, eltype(xj); weighted = false) - if g.is_coalesced - return (A' * xj')' # workaround for CUDA issue: https://github.com/JuliaGPU/CUDA.jl/issues/2820 - else - return xj * A - end + + return xj * A end ## E_MUL_XJ From 4beb3f58170c222f6ae626b2224bf462ca7e66eb Mon Sep 17 00:00:00 2001 From: dferre97 Date: Tue, 12 Aug 2025 13:38:28 +0200 Subject: [PATCH 3/9] Update comments to clarify coalesce behavior --- GNNGraphs/src/gnngraph.jl | 2 +- GNNGraphs/src/transform.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/GNNGraphs/src/gnngraph.jl b/GNNGraphs/src/gnngraph.jl index b90cfc032..e59c8f1e0 100644 --- a/GNNGraphs/src/gnngraph.jl +++ b/GNNGraphs/src/gnngraph.jl @@ -113,7 +113,7 @@ struct GNNGraph{T <: Union{COO_T, ADJMAT_T}} <: AbstractGNNGraph{T} ndata::DataStore edata::DataStore gdata::DataStore - is_coalesced::Bool # only for :coo, true if the graph is coalesced, i.e., indices ordered by row and no multi edges + is_coalesced::Bool # only for :coo, true if the graph is coalesced, i.e., no multi edges and indices ordered by target, then source end # GNNGraph constructor setting the is_coalesced field to false diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl index 61a11f455..1c8cffaa9 100644 --- a/GNNGraphs/src/transform.jl +++ b/GNNGraphs/src/transform.jl @@ -148,7 +148,7 @@ end """ coalesce(g::GNNGraph; aggr=+) -Return a new GNNGraph where all multiple edges between the same pair of nodes are merged (using aggr for edge weights and features), and the edge indices are sorted lexicographically (by source, then target). +Return a new GNNGraph where all multiple edges between the same pair of nodes are merged (using aggr for edge weights and features), and the edge indices are sorted lexicographically (by target, then by source). This method is only applicable to graphs of type `:coo`. `aggr` can take value `+`,`min`, `max` or `mean`. From 992e5bce961731f96784d7e0e2c1f8f08920bc1c Mon Sep 17 00:00:00 2001 From: dferre97 Date: Tue, 26 Aug 2025 15:35:52 +0200 Subject: [PATCH 4/9] Add custom _adjacency_matrix for propagate CUDA COO graphs - Leave public adjacency_matrix interface uniform, always returning a sparse adjacency_matrix - Implement custom _adjacency_matrix for propagate copy_xj for CUDA COO graphs, converting to dense when more efficient --- GNNGraphs/ext/GNNGraphsCUDAExt.jl | 26 +++++++++----------------- GNNlib/ext/GNNlibCUDAExt.jl | 22 ++++++++++++++++++++-- 2 files changed, 29 insertions(+), 19 deletions(-) diff --git a/GNNGraphs/ext/GNNGraphsCUDAExt.jl b/GNNGraphs/ext/GNNGraphsCUDAExt.jl index ae7d773ce..9655f55a7 100644 --- a/GNNGraphs/ext/GNNGraphsCUDAExt.jl +++ b/GNNGraphs/ext/GNNGraphsCUDAExt.jl @@ -14,20 +14,6 @@ const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{ GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1)) -function Graphs.adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out, - weighted = true) - @debug "Using CUDA adjacency_matrix for GNNGraph" - if !g.is_coalesced - # Revisit after - # https://github.com/JuliaGPU/CUDA.jl/issues/1113 - A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesce, construction of sparse matrix is slow - else - A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted) - end - @assert size(A) == (n, n) - return dir == :out ? A : A' -end - # Transform GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz) @@ -54,9 +40,9 @@ end # Convert function GNNGraphs.to_sparse(coo::CUDA_COO_T, T = nothing; dir = :out, num_nodes = nothing, - weighted = true) + weighted = true, is_coalesced = false) s, t, eweight = coo - @debug "Using CUDA to_sparse for COO" + @debug "Using CUDA to_sparse for COO with is_coalesced=$is_coalesced" T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T if eweight === nothing || !weighted @@ -64,8 +50,14 @@ function GNNGraphs.to_sparse(coo::CUDA_COO_T, T = nothing; dir = :out, num_nodes end num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes - A = CUDA.CUSPARSE.CuSparseMatrixCOO{T,eltype(s)}(s, t, eweight, (num_nodes, num_nodes)) # create sparse matrix in COO format + # if coalesced build directly sparse coo matrix + if is_coalesced + A = CUDA.CUSPARSE.CuSparseMatrixCOO{T,eltype(s)}(s, t, eweight, (num_nodes, num_nodes)) + else + A = sparse(s, t, eweight, num_nodes, num_nodes) + end + num_edges::Int = nnz(A) if eltype(A) != T A = T.(A) diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index 103318772..aaaa71e64 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -3,7 +3,9 @@ module GNNlibCUDAExt using CUDA using Random, Statistics, LinearAlgebra using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj -using GNNGraphs: GNNGraph, COO_T, SPARSE_T, adjacency_matrix +using GNNGraphs: GNNGraph, COO_T, SPARSE_T + +const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{Nothing, AnyCuArray}} ###### PROPAGATE SPECIALIZATIONS #################### @@ -13,7 +15,7 @@ using GNNGraphs: GNNGraph, COO_T, SPARSE_T, adjacency_matrix function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), xi, xj::AnyCuMatrix, e) @debug "Using CUDA propagate for copy_xj" - A = adjacency_matrix(g, eltype(xj); weighted = false) + A = _adjacency_matrix(g, eltype(xj); weighted = false) return xj * A end @@ -45,4 +47,20 @@ end # Flux.Zygote.@nograd compute_degree +## CUSTOM ADJACENCY_MATRIX IMPLEMENTATION FOR CUDA COO GRAPHS, returning dense matrix when not coalesced, more efficient + +function _adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out, + weighted = true) + @debug "Using CUDA _adjacency_matrix for COO GNNGraph" + if !g.is_coalesced + # Revisit after + # https://github.com/JuliaGPU/CUDA.jl/issues/1113 + A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesced, construction of sparse matrix is slow + else + A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted, is_coalesced = true) + end + @assert size(A) == (n, n) + return dir == :out ? A : A' +end + end #module From 258d35e5b05fd824a67025031976fe6eae5db811 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Wed, 27 Aug 2025 09:53:17 +0200 Subject: [PATCH 5/9] Fix imports --- GNNlib/ext/GNNlibCUDAExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index aaaa71e64..d5380a729 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -3,7 +3,7 @@ module GNNlibCUDAExt using CUDA using Random, Statistics, LinearAlgebra using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj -using GNNGraphs: GNNGraph, COO_T, SPARSE_T +using GNNGraphs: GNNGraph, COO_T, SPARSE_T, to_dense, to_sparse const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{Nothing, AnyCuArray}} @@ -55,9 +55,9 @@ function _adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); d if !g.is_coalesced # Revisit after # https://github.com/JuliaGPU/CUDA.jl/issues/1113 - A, n, m = GNNGraphs.to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesced, construction of sparse matrix is slow + A, n, m = to_dense(g.graph, T; num_nodes = g.num_nodes, weighted) # if not coalesced, construction of sparse matrix is slow else - A, n, m = GNNGraphs.to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted, is_coalesced = true) + A, n, m = to_sparse(g.graph, T; num_nodes = g.num_nodes, weighted, is_coalesced = true) end @assert size(A) == (n, n) return dir == :out ? A : A' From a2d7a62f6dafa6a0e62ab9d49eacbd675042f918 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Wed, 27 Aug 2025 09:54:42 +0200 Subject: [PATCH 6/9] Update GPU compatibility checks for COO CUDA --- GNNGraphs/test/gnngraph.jl | 7 ++++--- GNNGraphs/test/transform.jl | 6 ++++-- GraphNeuralNetworks/test/layers/conv.jl | 2 +- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/GNNGraphs/test/gnngraph.jl b/GNNGraphs/test/gnngraph.jl index 2b18fe7b7..1d6d27b77 100644 --- a/GNNGraphs/test/gnngraph.jl +++ b/GNNGraphs/test/gnngraph.jl @@ -99,13 +99,14 @@ end mat_gpu = adjacency_matrix(g_gpu) @test mat_gpu isa AbstractMatrix{Int} @test get_device(mat_gpu) isa AbstractGPUDevice - @test Array(mat_gpu) == adj_mat + # Convert to float first because poor Int support in CUSPARSE, throws an error + @test Array(Float32.(mat_gpu)) == Float32.(adj_mat) end end @testset "normalized_laplacian" begin mat = normalized_laplacian(g) - if TEST_GPU && !(dev isa MetalDevice) && GRAPH_T != :sparse + if TEST_GPU && !(dev isa MetalDevice) && GRAPH_T != :sparse && GRAPH_T != :coo mat_gpu = normalized_laplacian(g_gpu) @test mat_gpu isa AbstractMatrix{Float32} @test get_device(mat_gpu)isa AbstractGPUDevice @@ -114,7 +115,7 @@ end end @testset "scaled_laplacian" begin - if TEST_GPU && !(dev isa MetalDevice) && GRAPH_T != :sparse + if TEST_GPU && !(dev isa MetalDevice) && GRAPH_T != :sparse && GRAPH_T != :coo mat = scaled_laplacian(g) mat_gpu = scaled_laplacian(g_gpu) @test mat_gpu isa AbstractMatrix{Float32} diff --git a/GNNGraphs/test/transform.jl b/GNNGraphs/test/transform.jl index 256d851bf..fb7e95bf4 100644 --- a/GNNGraphs/test/transform.jl +++ b/GNNGraphs/test/transform.jl @@ -456,8 +456,10 @@ end s2, t2 = edge_index(g2) w2 = get_edge_weight(g2) - @test s2 == [1, 2, 2, 3, 3, 4, 4] - @test t2 == [2, 1, 3, 2, 4, 3, 4] + # @test s2 == [1, 2, 2, 3, 3, 4, 4] + # @test t2 == [2, 1, 3, 2, 4, 3, 4] + @test s2 == [2, 1, 3, 2, 4, 3, 4] + @test t2 == [1, 2, 2, 3, 3, 4, 4] @test w2 == [1, 1, 2, 2, 3.5, 3.5, 5] @test g2.edata.e == [10.0, 10.0, 20.0, 20.0, 35.0, 35.0, 50.0] end diff --git a/GraphNeuralNetworks/test/layers/conv.jl b/GraphNeuralNetworks/test/layers/conv.jl index 16e9b2fd5..97cc7a355 100644 --- a/GraphNeuralNetworks/test/layers/conv.jl +++ b/GraphNeuralNetworks/test/layers/conv.jl @@ -108,7 +108,7 @@ end if gpu_backend() == "AMDGPU" broken = true - elseif gpu_backend() == "CUDA" && get_graph_type(g) == :sparse + elseif gpu_backend() == "CUDA" && get_graph_type(g) in [:coo, :sparse] broken = true else broken = false From 7274e048b7c5e4e158df979452a10a8e708562c6 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Sun, 31 Aug 2025 10:31:34 +0200 Subject: [PATCH 7/9] Add @non_differentiable annotation to _adjacency_matrix function --- GNNlib/ext/GNNlibCUDAExt.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index d5380a729..cefcf4430 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -4,6 +4,7 @@ using CUDA using Random, Statistics, LinearAlgebra using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj using GNNGraphs: GNNGraph, COO_T, SPARSE_T, to_dense, to_sparse +using ChainRulesCore: @non_differentiable const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{Nothing, AnyCuArray}} @@ -63,4 +64,6 @@ function _adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); d return dir == :out ? A : A' end +@non_differentiable _adjacency_matrix(x...) + end #module From f0f715bb64c5f5843d54e942dee2c99a301f8fe8 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Sun, 31 Aug 2025 14:44:39 +0200 Subject: [PATCH 8/9] Add tests for coalesced COO graphs --- GNNlib/test/test_module.jl | 8 +++++++- GraphNeuralNetworks/test/test_module.jl | 8 +++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl index b6894cdfa..075881af8 100644 --- a/GNNlib/test/test_module.jl +++ b/GNNlib/test/test_module.jl @@ -150,7 +150,7 @@ function test_gradients( return true end -function generate_test_graphs(graph_type) +function generate_test_graphs(graph_type; do_coalesce=false) adj1 = [0 1 0 1 1 0 1 0 0 1 0 1 @@ -168,12 +168,18 @@ function generate_test_graphs(graph_type) g_single_vertex = GNNGraph(adj_single_vertex, ndata = rand(Float32, D_IN, 4); graph_type) + + if graph_type == :coo && do_coalesce + g1 = coalesce(g1) + g_single_vertex = coalesce(g_single_vertex) + end return (g1, g_single_vertex) end GRAPH_TYPES = [:coo, :dense, :sparse] TEST_GRAPHS = [generate_test_graphs(:coo)..., + generate_test_graphs(:coo, do_coalesce=true)..., generate_test_graphs(:dense)..., generate_test_graphs(:sparse)...] diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl index 8f7a0446b..74c25a555 100644 --- a/GraphNeuralNetworks/test/test_module.jl +++ b/GraphNeuralNetworks/test/test_module.jl @@ -157,7 +157,7 @@ function test_gradients( end -function generate_test_graphs(graph_type) +function generate_test_graphs(graph_type; do_coalesce=false) adj1 = [0 1 0 1 1 0 1 0 0 1 0 1 @@ -175,12 +175,18 @@ function generate_test_graphs(graph_type) g_single_vertex = GNNGraph(adj_single_vertex, ndata = rand(Float32, D_IN, 4); graph_type) + + if graph_type == :coo && do_coalesce + g1 = coalesce(g1) + g_single_vertex = coalesce(g_single_vertex) + end return (g1, g_single_vertex) end GRAPH_TYPES = [:coo, :dense, :sparse] TEST_GRAPHS = [generate_test_graphs(:coo)..., + generate_test_graphs(:coo, do_coalesce=true)..., generate_test_graphs(:dense)..., generate_test_graphs(:sparse)...] From 620a2e55939ad1edaab3db2fbe411f9a8bf392e2 Mon Sep 17 00:00:00 2001 From: dferre97 Date: Sun, 31 Aug 2025 16:05:07 +0200 Subject: [PATCH 9/9] Remove debug statements --- GNNGraphs/ext/GNNGraphsCUDAExt.jl | 1 - GNNlib/ext/GNNlibCUDAExt.jl | 2 -- 2 files changed, 3 deletions(-) diff --git a/GNNGraphs/ext/GNNGraphsCUDAExt.jl b/GNNGraphs/ext/GNNGraphsCUDAExt.jl index 9655f55a7..2b38ee739 100644 --- a/GNNGraphs/ext/GNNGraphsCUDAExt.jl +++ b/GNNGraphs/ext/GNNGraphsCUDAExt.jl @@ -42,7 +42,6 @@ end function GNNGraphs.to_sparse(coo::CUDA_COO_T, T = nothing; dir = :out, num_nodes = nothing, weighted = true, is_coalesced = false) s, t, eweight = coo - @debug "Using CUDA to_sparse for COO with is_coalesced=$is_coalesced" T = T === nothing ? (eweight === nothing ? eltype(s) : eltype(eweight)) : T if eweight === nothing || !weighted diff --git a/GNNlib/ext/GNNlibCUDAExt.jl b/GNNlib/ext/GNNlibCUDAExt.jl index cefcf4430..f745d51da 100644 --- a/GNNlib/ext/GNNlibCUDAExt.jl +++ b/GNNlib/ext/GNNlibCUDAExt.jl @@ -15,7 +15,6 @@ const CUDA_COO_T = Tuple{T, T, V} where {T <: AnyCuArray{<:Integer}, V <: Union{ ## avoid the fast path on gpu until we have better cuda support function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+), xi, xj::AnyCuMatrix, e) - @debug "Using CUDA propagate for copy_xj" A = _adjacency_matrix(g, eltype(xj); weighted = false) return xj * A @@ -52,7 +51,6 @@ end function _adjacency_matrix(g::GNNGraph{<:CUDA_COO_T}, T::DataType = eltype(g); dir = :out, weighted = true) - @debug "Using CUDA _adjacency_matrix for COO GNNGraph" if !g.is_coalesced # Revisit after # https://github.com/JuliaGPU/CUDA.jl/issues/1113