From 82552ab903b0cab9f5b673eb38c4bc4eac50eb48 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 24 Jan 2024 16:54:30 -0600 Subject: [PATCH 1/5] nx-cugraph: rename `plc=` to `_plc=` (#4106) As discussed here: https://github.com/rapidsai/cugraph/pull/4093#discussion_r1458012713 Authors: - Erik Welch (https://github.com/eriknw) Approvers: - Rick Ratzel (https://github.com/rlratzel) URL: https://github.com/rapidsai/cugraph/pull/4106 --- python/nx-cugraph/lint.yaml | 6 +-- .../nx_cugraph/algorithms/bipartite/basic.py | 2 +- .../algorithms/centrality/betweenness.py | 4 +- .../algorithms/centrality/eigenvector.py | 2 +- .../nx_cugraph/algorithms/centrality/katz.py | 2 +- .../nx_cugraph/algorithms/cluster.py | 16 ++++++-- .../algorithms/community/louvain.py | 2 +- .../algorithms/components/connected.py | 13 +++---- .../components/strongly_connected.py | 6 +-- .../algorithms/components/weakly_connected.py | 6 +-- .../nx-cugraph/nx_cugraph/algorithms/core.py | 4 +- .../nx-cugraph/nx_cugraph/algorithms/dag.py | 4 +- .../algorithms/link_analysis/hits_alg.py | 2 +- .../algorithms/link_analysis/pagerank_alg.py | 2 +- .../algorithms/shortest_paths/unweighted.py | 4 +- .../traversal/breadth_first_search.py | 38 +++++++++++++++---- .../nx_cugraph/algorithms/tree/recognition.py | 8 ++-- .../nx-cugraph/nx_cugraph/utils/decorators.py | 16 ++++---- 18 files changed, 84 insertions(+), 53 deletions(-) diff --git a/python/nx-cugraph/lint.yaml b/python/nx-cugraph/lint.yaml index 0d4f0b59413..5a4773168b6 100644 --- a/python/nx-cugraph/lint.yaml +++ b/python/nx-cugraph/lint.yaml @@ -50,7 +50,7 @@ repos: - id: black # - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.1.13 + rev: v0.1.14 hooks: - id: ruff args: [--fix-only, --show-fixes] # --unsafe-fixes] @@ -62,7 +62,7 @@ repos: additional_dependencies: &flake8_dependencies # These versions need updated manually - flake8==7.0.0 - - flake8-bugbear==23.12.2 + - flake8-bugbear==24.1.17 - flake8-simplify==0.21.0 - repo: https://github.com/asottile/yesqa rev: v1.5.0 @@ -77,7 +77,7 @@ repos: additional_dependencies: [tomli] files: ^(nx_cugraph|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.1.13 + rev: v0.1.14 hooks: - id: ruff - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py b/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py index d0e9a5c7f1b..46c6b54075b 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/bipartite/basic.py @@ -21,7 +21,7 @@ ] -@networkx_algorithm(plc="triangle_count", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="triangle_count") def is_bipartite(G): G = _to_graph(G) # Counting triangles may not be the fastest way to do this, but it is simple. diff --git a/python/nx-cugraph/nx_cugraph/algorithms/centrality/betweenness.py b/python/nx-cugraph/nx_cugraph/algorithms/centrality/betweenness.py index ba2b3d9c895..f6bb142cded 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/centrality/betweenness.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/centrality/betweenness.py @@ -21,8 +21,8 @@ @networkx_algorithm( is_incomplete=True, # weight not supported is_different=True, # RNG with seed is different - plc="betweenness_centrality", version_added="23.10", + _plc="betweenness_centrality", ) def betweenness_centrality( G, k=None, normalized=True, weight=None, endpoints=False, seed=None @@ -54,8 +54,8 @@ def _(G, k=None, normalized=True, weight=None, endpoints=False, seed=None): @networkx_algorithm( is_incomplete=True, # weight not supported is_different=True, # RNG with seed is different - plc="edge_betweenness_centrality", version_added="23.10", + _plc="edge_betweenness_centrality", ) def edge_betweenness_centrality(G, k=None, normalized=True, weight=None, seed=None): """`weight` parameter is not yet supported, and RNG with seed may be different.""" diff --git a/python/nx-cugraph/nx_cugraph/algorithms/centrality/eigenvector.py b/python/nx-cugraph/nx_cugraph/algorithms/centrality/eigenvector.py index 9e615955a8b..65a8633667a 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/centrality/eigenvector.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/centrality/eigenvector.py @@ -29,8 +29,8 @@ @networkx_algorithm( extra_params=_dtype_param, is_incomplete=True, # nstart not supported - plc="eigenvector_centrality", version_added="23.12", + _plc="eigenvector_centrality", ) def eigenvector_centrality( G, max_iter=100, tol=1.0e-6, nstart=None, weight=None, *, dtype=None diff --git a/python/nx-cugraph/nx_cugraph/algorithms/centrality/katz.py b/python/nx-cugraph/nx_cugraph/algorithms/centrality/katz.py index a2fb950c1aa..4a0684f72ee 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/centrality/katz.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/centrality/katz.py @@ -29,8 +29,8 @@ @networkx_algorithm( extra_params=_dtype_param, is_incomplete=True, # nstart and normalized=False not supported - plc="katz_centrality", version_added="23.12", + _plc="katz_centrality", ) def katz_centrality( G, diff --git a/python/nx-cugraph/nx_cugraph/algorithms/cluster.py b/python/nx-cugraph/nx_cugraph/algorithms/cluster.py index 951c358ff26..a458e6c04db 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/cluster.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/cluster.py @@ -45,7 +45,7 @@ def _triangles(G, nodes, symmetrize=None): @not_implemented_for("directed") -@networkx_algorithm(plc="triangle_count", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="triangle_count") def triangles(G, nodes=None): G = _to_undirected_graph(G) node_ids, triangles, is_single_node = _triangles(G, nodes) @@ -57,9 +57,13 @@ def triangles(G, nodes=None): @not_implemented_for("directed") -@networkx_algorithm(is_incomplete=True, plc="triangle_count", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="triangle_count") def clustering(G, nodes=None, weight=None): """Directed graphs and `weight` parameter are not yet supported.""" + if weight is not None: + raise NotImplementedError( + "Weighted implementation of clustering not currently supported" + ) G = _to_undirected_graph(G) node_ids, triangles, is_single_node = _triangles(G, nodes) if len(G) == 0: @@ -83,9 +87,13 @@ def _(G, nodes=None, weight=None): @not_implemented_for("directed") -@networkx_algorithm(is_incomplete=True, plc="triangle_count", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="triangle_count") def average_clustering(G, nodes=None, weight=None, count_zeros=True): """Directed graphs and `weight` parameter are not yet supported.""" + if weight is not None: + raise NotImplementedError( + "Weighted implementation of average_clustering not currently supported" + ) G = _to_undirected_graph(G) node_ids, triangles, is_single_node = _triangles(G, nodes) if len(G) == 0: @@ -110,7 +118,7 @@ def _(G, nodes=None, weight=None, count_zeros=True): @not_implemented_for("directed") -@networkx_algorithm(is_incomplete=True, plc="triangle_count", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="triangle_count") def transitivity(G): """Directed graphs are not yet supported.""" G = _to_undirected_graph(G) diff --git a/python/nx-cugraph/nx_cugraph/algorithms/community/louvain.py b/python/nx-cugraph/nx_cugraph/algorithms/community/louvain.py index 413ff9ca5e3..f58f1000fc4 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/community/louvain.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/community/louvain.py @@ -36,8 +36,8 @@ }, is_incomplete=True, # seed not supported; self-loops not supported is_different=True, # RNG different - plc="louvain", version_added="23.10", + _plc="louvain", ) def louvain_communities( G, diff --git a/python/nx-cugraph/nx_cugraph/algorithms/components/connected.py b/python/nx-cugraph/nx_cugraph/algorithms/components/connected.py index cdb9f54f6c4..24955e3eac8 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/components/connected.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/components/connected.py @@ -26,7 +26,7 @@ @not_implemented_for("directed") -@networkx_algorithm(plc="weakly_connected_components", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="weakly_connected_components") def number_connected_components(G): G = _to_undirected_graph(G) return _number_connected_components(G) @@ -50,14 +50,11 @@ def _number_connected_components(G, symmetrize=None): @number_connected_components._can_run def _(G): # NetworkX <= 3.2.1 does not check directedness for us - try: - return not G.is_directed() - except Exception: - return False + return not G.is_directed() @not_implemented_for("directed") -@networkx_algorithm(plc="weakly_connected_components", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="weakly_connected_components") def connected_components(G): G = _to_undirected_graph(G) return _connected_components(G) @@ -80,7 +77,7 @@ def _connected_components(G, symmetrize=None): @not_implemented_for("directed") -@networkx_algorithm(plc="weakly_connected_components", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="weakly_connected_components") def is_connected(G): G = _to_undirected_graph(G) return _is_connected(G) @@ -106,7 +103,7 @@ def _is_connected(G, symmetrize=None): @not_implemented_for("directed") -@networkx_algorithm(plc="weakly_connected_components", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="weakly_connected_components") def node_connected_component(G, n): # We could also do plain BFS from n G = _to_undirected_graph(G) diff --git a/python/nx-cugraph/nx_cugraph/algorithms/components/strongly_connected.py b/python/nx-cugraph/nx_cugraph/algorithms/components/strongly_connected.py index 8fdf99ed5ea..d1713129703 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/components/strongly_connected.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/components/strongly_connected.py @@ -51,7 +51,7 @@ def _strongly_connected_components(G): @not_implemented_for("undirected") -@networkx_algorithm(version_added="24.02", plc="strongly_connected_components") +@networkx_algorithm(version_added="24.02", _plc="strongly_connected_components") def strongly_connected_components(G): G = _to_directed_graph(G) if G.src_indices.size == 0: @@ -62,7 +62,7 @@ def strongly_connected_components(G): @not_implemented_for("undirected") -@networkx_algorithm(version_added="24.02", plc="strongly_connected_components") +@networkx_algorithm(version_added="24.02", _plc="strongly_connected_components") def number_strongly_connected_components(G): G = _to_directed_graph(G) if G.src_indices.size == 0: @@ -72,7 +72,7 @@ def number_strongly_connected_components(G): @not_implemented_for("undirected") -@networkx_algorithm(version_added="24.02", plc="strongly_connected_components") +@networkx_algorithm(version_added="24.02", _plc="strongly_connected_components") def is_strongly_connected(G): G = _to_directed_graph(G) if len(G) == 0: diff --git a/python/nx-cugraph/nx_cugraph/algorithms/components/weakly_connected.py b/python/nx-cugraph/nx_cugraph/algorithms/components/weakly_connected.py index 5b797b39118..e42acdd3d84 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/components/weakly_connected.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/components/weakly_connected.py @@ -27,21 +27,21 @@ @not_implemented_for("undirected") -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def weakly_connected_components(G): G = _to_directed_graph(G) return _connected_components(G, symmetrize="union") @not_implemented_for("undirected") -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def number_weakly_connected_components(G): G = _to_directed_graph(G) return _number_connected_components(G, symmetrize="union") @not_implemented_for("undirected") -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def is_weakly_connected(G): G = _to_directed_graph(G) return _is_connected(G, symmetrize="union") diff --git a/python/nx-cugraph/nx_cugraph/algorithms/core.py b/python/nx-cugraph/nx_cugraph/algorithms/core.py index f323cdf6004..71f61abf45b 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/core.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/core.py @@ -28,7 +28,7 @@ @not_implemented_for("directed") @not_implemented_for("multigraph") -@networkx_algorithm(is_incomplete=True, plc="core_number", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="core_number") def core_number(G): """Directed graphs are not yet supported.""" G = _to_undirected_graph(G) @@ -55,7 +55,7 @@ def _(G): @not_implemented_for("directed") @not_implemented_for("multigraph") -@networkx_algorithm(is_incomplete=True, plc="k_truss_subgraph", version_added="23.12") +@networkx_algorithm(is_incomplete=True, version_added="23.12", _plc="k_truss_subgraph") def k_truss(G, k): """ Currently raises `NotImplementedError` for graphs with more than one connected diff --git a/python/nx-cugraph/nx_cugraph/algorithms/dag.py b/python/nx-cugraph/nx_cugraph/algorithms/dag.py index ad5b7594aa1..64be0a58105 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/dag.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/dag.py @@ -45,11 +45,11 @@ def _ancestors_and_descendants(G, source, *, is_ancestors): return G._nodearray_to_set(node_ids[mask]) -@networkx_algorithm(plc="bfs", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="bfs") def descendants(G, source): return _ancestors_and_descendants(G, source, is_ancestors=False) -@networkx_algorithm(plc="bfs", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="bfs") def ancestors(G, source): return _ancestors_and_descendants(G, source, is_ancestors=True) diff --git a/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/hits_alg.py b/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/hits_alg.py index caa01327a56..9e723624a3b 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/hits_alg.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/hits_alg.py @@ -33,8 +33,8 @@ ), **_dtype_param, }, - plc="hits", version_added="23.12", + _plc="hits", ) def hits( G, diff --git a/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/pagerank_alg.py b/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/pagerank_alg.py index d45d019c1b7..55fcc3e520a 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/pagerank_alg.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/link_analysis/pagerank_alg.py @@ -29,8 +29,8 @@ @networkx_algorithm( extra_params=_dtype_param, is_incomplete=True, # dangling not supported - plc={"pagerank", "personalized_pagerank"}, version_added="23.12", + _plc={"pagerank", "personalized_pagerank"}, ) def pagerank( G, diff --git a/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py index b1032a8236b..2012495953e 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py @@ -21,12 +21,12 @@ __all__ = ["single_source_shortest_path_length", "single_target_shortest_path_length"] -@networkx_algorithm(plc="bfs", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="bfs") def single_source_shortest_path_length(G, source, cutoff=None): return _single_shortest_path_length(G, source, cutoff, "Source") -@networkx_algorithm(plc="bfs", version_added="23.12") +@networkx_algorithm(version_added="23.12", _plc="bfs") def single_target_shortest_path_length(G, target, cutoff=None): return _single_shortest_path_length(G, target, cutoff, "Target") diff --git a/python/nx-cugraph/nx_cugraph/algorithms/traversal/breadth_first_search.py b/python/nx-cugraph/nx_cugraph/algorithms/traversal/breadth_first_search.py index aa671bbb7d4..ef1c011363a 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/traversal/breadth_first_search.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/traversal/breadth_first_search.py @@ -57,9 +57,17 @@ def _bfs(G, source, *, depth_limit=None, reverse=False): return distances[mask], predecessors[mask], node_ids[mask] -@networkx_algorithm(is_incomplete=True, plc="bfs", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="bfs") def generic_bfs_edges(G, source, neighbors=None, depth_limit=None, sort_neighbors=None): """`neighbors` and `sort_neighbors` parameters are not yet supported.""" + if neighbors is not None: + raise NotImplementedError( + "neighbors argument in generic_bfs_edges is not currently supported" + ) + if sort_neighbors is not None: + raise NotImplementedError( + "sort_neighbors argument in generic_bfs_edges is not currently supported" + ) return bfs_edges(source, depth_limit=depth_limit) @@ -68,9 +76,13 @@ def _(G, source, neighbors=None, depth_limit=None, sort_neighbors=None): return neighbors is None and sort_neighbors is None -@networkx_algorithm(is_incomplete=True, plc="bfs", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="bfs") def bfs_edges(G, source, reverse=False, depth_limit=None, sort_neighbors=None): """`sort_neighbors` parameter is not yet supported.""" + if sort_neighbors is not None: + raise NotImplementedError( + "sort_neighbors argument in bfs_edges is not currently supported" + ) G = _check_G_and_source(G, source) if depth_limit is not None and depth_limit < 1: return @@ -95,9 +107,13 @@ def _(G, source, reverse=False, depth_limit=None, sort_neighbors=None): return sort_neighbors is None -@networkx_algorithm(is_incomplete=True, plc="bfs", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="bfs") def bfs_tree(G, source, reverse=False, depth_limit=None, sort_neighbors=None): """`sort_neighbors` parameter is not yet supported.""" + if sort_neighbors is not None: + raise NotImplementedError( + "sort_neighbors argument in bfs_tree is not currently supported" + ) G = _check_G_and_source(G, source) if depth_limit is not None and depth_limit < 1: return nxcg.DiGraph.from_coo( @@ -149,9 +165,13 @@ def _(G, source, reverse=False, depth_limit=None, sort_neighbors=None): return sort_neighbors is None -@networkx_algorithm(is_incomplete=True, plc="bfs", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="bfs") def bfs_successors(G, source, depth_limit=None, sort_neighbors=None): """`sort_neighbors` parameter is not yet supported.""" + if sort_neighbors is not None: + raise NotImplementedError( + "sort_neighbors argument in bfs_successors is not currently supported" + ) G = _check_G_and_source(G, source) if depth_limit is not None and depth_limit < 1: yield (source, []) @@ -173,7 +193,7 @@ def _(G, source, depth_limit=None, sort_neighbors=None): return sort_neighbors is None -@networkx_algorithm(plc="bfs", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="bfs") def bfs_layers(G, sources): G = _to_graph(G) if sources in G: @@ -201,9 +221,13 @@ def bfs_layers(G, sources): return (G._nodearray_to_list(groups[key]) for key in range(len(groups))) -@networkx_algorithm(is_incomplete=True, plc="bfs", version_added="24.02") +@networkx_algorithm(is_incomplete=True, version_added="24.02", _plc="bfs") def bfs_predecessors(G, source, depth_limit=None, sort_neighbors=None): """`sort_neighbors` parameter is not yet supported.""" + if sort_neighbors is not None: + raise NotImplementedError( + "sort_neighbors argument in bfs_predecessors is not currently supported" + ) G = _check_G_and_source(G, source) if depth_limit is not None and depth_limit < 1: return @@ -227,7 +251,7 @@ def _(G, source, depth_limit=None, sort_neighbors=None): return sort_neighbors is None -@networkx_algorithm(plc="bfs", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="bfs") def descendants_at_distance(G, source, distance): G = _check_G_and_source(G, source) if distance is None or distance < 0: diff --git a/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py b/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py index 0b82f079d43..74f57b5ea5a 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py @@ -21,20 +21,20 @@ @not_implemented_for("undirected") -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def is_arborescence(G): G = _to_directed_graph(G) return is_tree(G) and int(G._in_degrees_array().max()) <= 1 @not_implemented_for("undirected") -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def is_branching(G): G = _to_directed_graph(G) return is_forest(G) and int(G._in_degrees_array().max()) <= 1 -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def is_forest(G): G = _to_graph(G) if len(G) == 0: @@ -60,7 +60,7 @@ def is_forest(G): return True -@networkx_algorithm(plc="weakly_connected_components", version_added="24.02") +@networkx_algorithm(version_added="24.02", _plc="weakly_connected_components") def is_tree(G): G = _to_graph(G) if len(G) == 0: diff --git a/python/nx-cugraph/nx_cugraph/utils/decorators.py b/python/nx-cugraph/nx_cugraph/utils/decorators.py index d09a9e9617a..011ebfd6ef7 100644 --- a/python/nx-cugraph/nx_cugraph/utils/decorators.py +++ b/python/nx-cugraph/nx_cugraph/utils/decorators.py @@ -59,7 +59,7 @@ def __new__( version_added: str, # Required is_incomplete: bool = False, # See self.extra_doc for details if True is_different: bool = False, # See self.extra_doc for details if True - plc: str | set[str] | None = None, # Hidden from user, may be removed someday + _plc: str | set[str] | None = None, # Hidden from user, may be removed someday ): if func is None: return partial( @@ -67,10 +67,10 @@ def __new__( name=name, extra_params=extra_params, nodes_or_number=nodes_or_number, - plc=plc, version_added=version_added, is_incomplete=is_incomplete, is_different=is_different, + _plc=_plc, ) instance = object.__new__(cls) if nodes_or_number is not None and nx.__version__[:3] > "3.2": @@ -89,12 +89,14 @@ def __new__( f"extra_params must be dict, str, or None; got {type(extra_params)}" ) instance.extra_params = extra_params - if plc is None or isinstance(plc, set): - instance._plc_names = plc - elif isinstance(plc, str): - instance._plc_names = {plc} + if _plc is None or isinstance(_plc, set): + instance._plc_names = _plc + elif isinstance(_plc, str): + instance._plc_names = {_plc} else: - raise TypeError(f"plc argument must be str, set, or None; got {type(plc)}") + raise TypeError( + f"_plc argument must be str, set, or None; got {type(_plc)}" + ) instance.version_added = version_added instance.is_incomplete = is_incomplete instance.is_different = is_different From 3526af4f8776f9e84c664d26b79768f1345aba8b Mon Sep 17 00:00:00 2001 From: Seunghwa Kang <45857425+seunghwak@users.noreply.github.com> Date: Wed, 24 Jan 2024 21:58:17 -0800 Subject: [PATCH 2/5] Implement has_edge() & compute_multiplicity() (#4096) `graph_view.has_edge()` query whether the graph has given (src, dst) pairs as edges. `graph_view.compute_multiplicity()` query the edge multiplicity of given (src, dst) pairs (assumes that multiplicity is 0 if there is no edge between a given pair). This function throws an exception if `graph_view.is_multigraph()` is false (better use `has_edge()` for non-multigraph). In addition to adding the above two functions, this PR includes few code cleanups. * `major_idx_from_major_nocheck()` to `edge_partition_device_view_t` * Move `count_invalid_vertex_pais` from `nbr_intersection.cuh` to `error_check_utils.cuh` * Update `cugraph::test::to_host`, `to_device`, `device_gatherv`, and `device_allgatherv` to support `bool` type (and to handle `std::vector` which stores `bool` values in a packed format) Authors: - Seunghwa Kang (https://github.com/seunghwak) Approvers: - Chuck Hastings (https://github.com/ChuckHastings) - Naim (https://github.com/naimnv) - Joseph Nke (https://github.com/jnke2016) URL: https://github.com/rapidsai/cugraph/pull/4096 --- .../cugraph/edge_partition_device_view.cuh | 22 +- cpp/include/cugraph/graph_view.hpp | 25 +- cpp/src/link_prediction/similarity_impl.cuh | 3 +- cpp/src/prims/detail/nbr_intersection.cuh | 170 +-------- ..._v_pair_transform_dst_nbr_intersection.cuh | 3 +- cpp/src/prims/transform_e.cuh | 43 +-- cpp/src/structure/graph_view_impl.cuh | 345 +++++++++++++++++- cpp/src/utilities/error_check_utils.cuh | 137 +++++++ cpp/tests/CMakeLists.txt | 10 + .../count_self_loops_and_multi_edges_test.cpp | 17 +- ...has_edge_and_compute_multiplicity_test.cpp | 281 ++++++++++++++ ...has_edge_and_compute_multiplicity_test.cpp | 331 +++++++++++++++++ cpp/tests/utilities/device_comm_wrapper.cu | 18 +- cpp/tests/utilities/test_utilities.hpp | 48 +-- 14 files changed, 1217 insertions(+), 236 deletions(-) create mode 100644 cpp/src/utilities/error_check_utils.cuh create mode 100644 cpp/tests/structure/has_edge_and_compute_multiplicity_test.cpp create mode 100644 cpp/tests/structure/mg_has_edge_and_compute_multiplicity_test.cpp diff --git a/cpp/include/cugraph/edge_partition_device_view.cuh b/cpp/include/cugraph/edge_partition_device_view.cuh index 213f9b9497a..d1c2cf3df52 100644 --- a/cpp/include/cugraph/edge_partition_device_view.cuh +++ b/cpp/include/cugraph/edge_partition_device_view.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -298,6 +298,20 @@ class edge_partition_device_view_t major_idx_from_major_nocheck(vertex_t major) const noexcept + { + if (major_hypersparse_first_ && (major >= *major_hypersparse_first_)) { + auto major_hypersparse_idx = + detail::major_hypersparse_idx_from_major_nocheck_impl(*dcs_nzd_vertices_, major); + return major_hypersparse_idx + ? thrust::make_optional((*major_hypersparse_first_ - major_range_first_) + + *major_hypersparse_idx) + : thrust::nullopt; + } else { + return major - major_range_first_; + } + } + __device__ vertex_t major_from_major_idx_nocheck(vertex_t major_idx) const noexcept { if (major_hypersparse_first_) { @@ -339,6 +353,7 @@ class edge_partition_device_view_t{(*dcs_nzd_vertices_).data()} : thrust::nullopt; } + __host__ __device__ thrust::optional dcs_nzd_vertex_count() const { return dcs_nzd_vertices_ @@ -460,6 +475,11 @@ class edge_partition_device_view_t major_idx_from_major_nocheck(vertex_t major) const noexcept + { + return major_offset_from_major_nocheck(major); + } + __device__ vertex_t major_from_major_idx_nocheck(vertex_t major_idx) const noexcept { return major_from_major_offset_nocheck(major_idx); diff --git a/cpp/include/cugraph/graph_view.hpp b/cpp/include/cugraph/graph_view.hpp index 53c66c6483e..93d884a56d9 100644 --- a/cpp/include/cugraph/graph_view.hpp +++ b/cpp/include/cugraph/graph_view.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -631,6 +631,19 @@ class graph_view_t has_edge(raft::handle_t const& handle, + /* (edge_srcs, edge_dsts) should be pre-shuffled */ + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check = false); + + rmm::device_uvector compute_multiplicity( + raft::handle_t const& handle, + /* (edge_srcs, edge_dsts) should be pre-shuffled */ + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check = false); + template std::enable_if_t>> local_sorted_unique_edge_srcs() const @@ -928,6 +941,16 @@ class graph_view_t has_edge(raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check = false); + + rmm::device_uvector compute_multiplicity(raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check = false); + template std::enable_if_t>> local_sorted_unique_edge_srcs() const diff --git a/cpp/src/link_prediction/similarity_impl.cuh b/cpp/src/link_prediction/similarity_impl.cuh index 55e8f5c88d7..7ac294d7719 100644 --- a/cpp/src/link_prediction/similarity_impl.cuh +++ b/cpp/src/link_prediction/similarity_impl.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-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. @@ -18,6 +18,7 @@ #include #include #include +#include #include #include diff --git a/cpp/src/prims/detail/nbr_intersection.cuh b/cpp/src/prims/detail/nbr_intersection.cuh index cefc1836fa6..8261ec747f9 100644 --- a/cpp/src/prims/detail/nbr_intersection.cuh +++ b/cpp/src/prims/detail/nbr_intersection.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-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. @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -63,35 +64,6 @@ namespace cugraph { namespace detail { -// check vertices in the pair are valid and first element of the pair is within the local vertex -// partition range -template -struct is_invalid_input_vertex_pair_t { - vertex_t num_vertices{}; - raft::device_span edge_partition_major_range_firsts{}; - raft::device_span edge_partition_major_range_lasts{}; - vertex_t edge_partition_minor_range_first{}; - vertex_t edge_partition_minor_range_last{}; - - __device__ bool operator()(thrust::tuple pair) const - { - auto major = thrust::get<0>(pair); - auto minor = thrust::get<1>(pair); - if (!is_valid_vertex(num_vertices, major) || !is_valid_vertex(num_vertices, minor)) { - return true; - } - auto it = thrust::upper_bound(thrust::seq, - edge_partition_major_range_lasts.begin(), - edge_partition_major_range_lasts.end(), - major); - if (it == edge_partition_major_range_lasts.end()) { return true; } - auto edge_partition_idx = - static_cast(thrust::distance(edge_partition_major_range_lasts.begin(), it)); - if (major < edge_partition_major_range_firsts[edge_partition_idx]) { return true; } - return (minor < edge_partition_minor_range_first) || (minor >= edge_partition_minor_range_last); - } -}; - // group index determined by major_comm_rank (primary key) and local edge partition index (secondary // key) template @@ -154,24 +126,11 @@ struct update_rx_major_local_degree_t { auto major = rx_majors[rx_group_firsts[major_comm_rank * minor_comm_size + local_edge_partition_idx] + offset_in_local_edge_partition]; - vertex_t major_idx{0}; - edge_t local_degree{0}; - if (multi_gpu && (edge_partition.major_hypersparse_first() && - (major >= *(edge_partition.major_hypersparse_first())))) { - auto major_hypersparse_idx = edge_partition.major_hypersparse_idx_from_major_nocheck(major); - if (major_hypersparse_idx) { - major_idx = - (*(edge_partition.major_hypersparse_first()) - edge_partition.major_range_first()) + - *major_hypersparse_idx; - local_degree = edge_partition.local_degree(major_idx); - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major); - local_degree = edge_partition.local_degree(major_idx); - } + auto major_idx = edge_partition.major_idx_from_major_nocheck(major); + auto local_degree = major_idx ? edge_partition.local_degree(*major_idx) : edge_t{0}; if (edge_partition_e_mask && (local_degree > edge_t{0})) { - auto local_offset = edge_partition.local_offset(major_idx); + auto local_offset = edge_partition.local_offset(*major_idx); local_degree = static_cast( count_set_bits((*edge_partition_e_mask).value_first(), local_offset, local_degree)); } @@ -325,29 +284,11 @@ struct pick_min_degree_t { edge_t local_degree0{0}; vertex_t major0 = thrust::get<0>(pair); if constexpr (std::is_same_v) { - vertex_t major_idx{0}; - if constexpr (multi_gpu) { - if (edge_partition.major_hypersparse_first() && - (major0 >= *(edge_partition.major_hypersparse_first()))) { - auto major_hypersparse_idx = - edge_partition.major_hypersparse_idx_from_major_nocheck(major0); - if (major_hypersparse_idx) { - major_idx = - (*(edge_partition.major_hypersparse_first()) - edge_partition.major_range_first()) + - *major_hypersparse_idx; - local_degree0 = edge_partition.local_degree(major_idx); - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major0); - local_degree0 = edge_partition.local_degree(major_idx); - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major0); - local_degree0 = edge_partition.local_degree(major_idx); - } + auto major_idx = edge_partition.major_idx_from_major_nocheck(major0); + local_degree0 = major_idx ? edge_partition.local_degree(*major_idx) : edge_t{0}; if (edge_partition_e_mask && (local_degree0 > edge_t{0})) { - auto local_offset = edge_partition.local_offset(major_idx); + auto local_offset = edge_partition.local_offset(*major_idx); local_degree0 = count_set_bits((*edge_partition_e_mask).value_first(), local_offset, local_degree0); } @@ -360,29 +301,11 @@ struct pick_min_degree_t { edge_t local_degree1{0}; vertex_t major1 = thrust::get<1>(pair); if constexpr (std::is_same_v) { - vertex_t major_idx{0}; - if constexpr (multi_gpu) { - if (edge_partition.major_hypersparse_first() && - (major1 >= *(edge_partition.major_hypersparse_first()))) { - auto major_hypersparse_idx = - edge_partition.major_hypersparse_idx_from_major_nocheck(major1); - if (major_hypersparse_idx) { - major_idx = - (*(edge_partition.major_hypersparse_first()) - edge_partition.major_range_first()) + - *major_hypersparse_idx; - local_degree1 = edge_partition.local_degree(major_idx); - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major1); - local_degree1 = edge_partition.local_degree(major_idx); - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major1); - local_degree1 = edge_partition.local_degree(major_idx); - } + auto major_idx = edge_partition.major_idx_from_major_nocheck(major1); + local_degree1 = major_idx ? edge_partition.local_degree(*major_idx) : edge_t{0}; if (edge_partition_e_mask && (local_degree1 > edge_t{0})) { - auto local_offset = edge_partition.local_offset(major_idx); + auto local_offset = edge_partition.local_offset(*major_idx); local_degree1 = count_set_bits((*edge_partition_e_mask).value_first(), local_offset, local_degree1); } @@ -699,77 +622,6 @@ struct gatherv_indices_t { } }; -template -size_t count_invalid_vertex_pairs(raft::handle_t const& handle, - GraphViewType const& graph_view, - VertexPairIterator vertex_pair_first, - VertexPairIterator vertex_pair_last) -{ - using vertex_t = typename GraphViewType::vertex_type; - - std::vector h_edge_partition_major_range_firsts( - graph_view.number_of_local_edge_partitions()); - std::vector h_edge_partition_major_range_lasts( - h_edge_partition_major_range_firsts.size()); - vertex_t edge_partition_minor_range_first{}; - vertex_t edge_partition_minor_range_last{}; - if constexpr (GraphViewType::is_multi_gpu) { - for (size_t i = 0; i < graph_view.number_of_local_edge_partitions(); i++) { - if constexpr (GraphViewType::is_storage_transposed) { - h_edge_partition_major_range_firsts[i] = graph_view.local_edge_partition_dst_range_first(i); - h_edge_partition_major_range_lasts[i] = graph_view.local_edge_partition_dst_range_last(i); - } else { - h_edge_partition_major_range_firsts[i] = graph_view.local_edge_partition_src_range_first(i); - h_edge_partition_major_range_lasts[i] = graph_view.local_edge_partition_src_range_last(i); - } - } - if constexpr (GraphViewType::is_storage_transposed) { - edge_partition_minor_range_first = graph_view.local_edge_partition_src_range_first(); - edge_partition_minor_range_last = graph_view.local_edge_partition_src_range_last(); - } else { - edge_partition_minor_range_first = graph_view.local_edge_partition_dst_range_first(); - edge_partition_minor_range_last = graph_view.local_edge_partition_dst_range_last(); - } - } else { - h_edge_partition_major_range_firsts[0] = vertex_t{0}; - h_edge_partition_major_range_lasts[0] = graph_view.number_of_vertices(); - edge_partition_minor_range_first = vertex_t{0}; - edge_partition_minor_range_last = graph_view.number_of_vertices(); - } - rmm::device_uvector d_edge_partition_major_range_firsts( - h_edge_partition_major_range_firsts.size(), handle.get_stream()); - rmm::device_uvector d_edge_partition_major_range_lasts( - h_edge_partition_major_range_lasts.size(), handle.get_stream()); - raft::update_device(d_edge_partition_major_range_firsts.data(), - h_edge_partition_major_range_firsts.data(), - h_edge_partition_major_range_firsts.size(), - handle.get_stream()); - raft::update_device(d_edge_partition_major_range_lasts.data(), - h_edge_partition_major_range_lasts.data(), - h_edge_partition_major_range_lasts.size(), - handle.get_stream()); - - auto num_invalid_pairs = thrust::count_if( - handle.get_thrust_policy(), - vertex_pair_first, - vertex_pair_last, - is_invalid_input_vertex_pair_t{ - graph_view.number_of_vertices(), - raft::device_span(d_edge_partition_major_range_firsts.begin(), - d_edge_partition_major_range_firsts.end()), - raft::device_span(d_edge_partition_major_range_lasts.begin(), - d_edge_partition_major_range_lasts.end()), - edge_partition_minor_range_first, - edge_partition_minor_range_last}); - if constexpr (GraphViewType::is_multi_gpu) { - auto& comm = handle.get_comms(); - num_invalid_pairs = - host_scalar_allreduce(comm, num_invalid_pairs, raft::comms::op_t::SUM, handle.get_stream()); - } - - return num_invalid_pairs; -} - // In multi-GPU, the first element of every vertex pair in [vertex_pair_first, vertex_pair) should // be within the valid edge partition major range assigned to this process and the second element // should be within the valid edge partition minor range assigned to this process. diff --git a/cpp/src/prims/per_v_pair_transform_dst_nbr_intersection.cuh b/cpp/src/prims/per_v_pair_transform_dst_nbr_intersection.cuh index 201c08325d7..469bfcb4e47 100644 --- a/cpp/src/prims/per_v_pair_transform_dst_nbr_intersection.cuh +++ b/cpp/src/prims/per_v_pair_transform_dst_nbr_intersection.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-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. @@ -19,6 +19,7 @@ #include #include #include +#include #include #include diff --git a/cpp/src/prims/transform_e.cuh b/cpp/src/prims/transform_e.cuh index c6623621d24..93a2d040b60 100644 --- a/cpp/src/prims/transform_e.cuh +++ b/cpp/src/prims/transform_e.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023, NVIDIA CORPORATION. + * Copyright (c) 2023-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. @@ -426,28 +426,15 @@ void transform_e(raft::handle_t const& handle, edge_first + edge_partition_offsets[i + 1], [edge_partition, edge_partition_e_mask] __device__(thrust::tuple edge) { - auto major = thrust::get<0>(edge); - auto minor = thrust::get<1>(edge); - vertex_t major_idx{}; - auto major_hypersparse_first = edge_partition.major_hypersparse_first(); - if (major_hypersparse_first) { - if (major < *major_hypersparse_first) { - major_idx = edge_partition.major_offset_from_major_nocheck(major); - } else { - auto major_hypersparse_idx = - edge_partition.major_hypersparse_idx_from_major_nocheck(major); - if (!major_hypersparse_idx) { return true; } - major_idx = - edge_partition.major_offset_from_major_nocheck(*major_hypersparse_first) + - *major_hypersparse_idx; - } - } else { - major_idx = edge_partition.major_offset_from_major_nocheck(major); - } + auto major = thrust::get<0>(edge); + auto minor = thrust::get<1>(edge); + auto major_idx = edge_partition.major_idx_from_major_nocheck(major); + if (!major_idx) { return true; } vertex_t const* indices{nullptr}; edge_t edge_offset{}; edge_t local_degree{}; - thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_idx); + thrust::tie(indices, edge_offset, local_degree) = + edge_partition.local_edges(*major_idx); auto lower_it = thrust::lower_bound(thrust::seq, indices, indices + local_degree, minor); if (*lower_it != minor) { return true; } @@ -494,24 +481,16 @@ void transform_e(raft::handle_t const& handle, auto major = thrust::get<0>(edge); auto minor = thrust::get<1>(edge); - auto major_hypersparse_first = edge_partition.major_hypersparse_first(); - auto major_offset = edge_partition.major_offset_from_major_nocheck(major); - vertex_t major_idx{major_offset}; - - if ((major_hypersparse_first) && (major >= *major_hypersparse_first)) { - auto major_hypersparse_idx = - edge_partition.major_hypersparse_idx_from_major_nocheck(major); - assert(major_hypersparse_idx); - major_idx = edge_partition.major_offset_from_major_nocheck(*major_hypersparse_first) + - *major_hypersparse_idx; - } + auto major_offset = edge_partition.major_offset_from_major_nocheck(major); + auto major_idx = edge_partition.major_idx_from_major_nocheck(major); + assert(major_idx); auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); vertex_t const* indices{nullptr}; edge_t edge_offset{}; edge_t local_degree{}; - thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_idx); + thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(*major_idx); auto lower_it = thrust::lower_bound(thrust::seq, indices, indices + local_degree, minor); auto upper_it = thrust::upper_bound(thrust::seq, lower_it, indices + local_degree, minor); diff --git a/cpp/src/structure/graph_view_impl.cuh b/cpp/src/structure/graph_view_impl.cuh index da0ecc991df..7928c61cf7b 100644 --- a/cpp/src/structure/graph_view_impl.cuh +++ b/cpp/src/structure/graph_view_impl.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -414,6 +415,59 @@ edge_t count_edge_partition_multi_edges( } } +template +std::tuple, std::vector> +compute_edge_indices_and_edge_partition_offsets( + raft::handle_t const& handle, + graph_view_t const& graph_view, + raft::device_span edge_majors, + raft::device_span edge_minors) +{ + auto edge_first = thrust::make_zip_iterator(edge_majors.begin(), edge_minors.begin()); + + rmm::device_uvector edge_indices(edge_majors.size(), handle.get_stream()); + thrust::sequence(handle.get_thrust_policy(), edge_indices.begin(), edge_indices.end(), size_t{0}); + thrust::sort(handle.get_thrust_policy(), + edge_indices.begin(), + edge_indices.end(), + [edge_first] __device__(size_t lhs, size_t rhs) { + return *(edge_first + lhs) < *(edge_first + rhs); + }); + + std::vector h_major_range_lasts(graph_view.number_of_local_edge_partitions()); + for (size_t i = 0; i < h_major_range_lasts.size(); ++i) { + if constexpr (store_transposed) { + h_major_range_lasts[i] = graph_view.local_edge_partition_dst_range_last(i); + } else { + h_major_range_lasts[i] = graph_view.local_edge_partition_src_range_last(i); + } + } + rmm::device_uvector d_major_range_lasts(h_major_range_lasts.size(), handle.get_stream()); + raft::update_device(d_major_range_lasts.data(), + h_major_range_lasts.data(), + h_major_range_lasts.size(), + handle.get_stream()); + rmm::device_uvector d_lower_bounds(d_major_range_lasts.size(), handle.get_stream()); + auto major_first = edge_majors.begin(); + auto sorted_major_first = thrust::make_transform_iterator( + edge_indices.begin(), + cugraph::detail::indirection_t{major_first}); + thrust::lower_bound(handle.get_thrust_policy(), + sorted_major_first, + sorted_major_first + edge_indices.size(), + d_major_range_lasts.begin(), + d_major_range_lasts.end(), + d_lower_bounds.begin()); + std::vector edge_partition_offsets(d_lower_bounds.size() + 1, 0); + raft::update_host(edge_partition_offsets.data() + 1, + d_lower_bounds.data(), + d_lower_bounds.size(), + handle.get_stream()); + handle.sync_stream(); + + return std::make_tuple(std::move(edge_indices), edge_partition_offsets); +} + } // namespace template @@ -751,4 +805,293 @@ edge_t graph_view_tlocal_edge_partition_segment_offsets()); } +template +rmm::device_uvector +graph_view_t>::has_edge( + raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check) +{ + CUGRAPH_EXPECTS( + edge_srcs.size() == edge_dsts.size(), + "Invalid input arguments: edge_srcs.size() does not coincide with edge_dsts.size()."); + + auto edge_first = + thrust::make_zip_iterator(store_transposed ? edge_dsts.begin() : edge_srcs.begin(), + store_transposed ? edge_srcs.begin() : edge_dsts.begin()); + + if (do_expensive_check) { + auto num_invalids = + detail::count_invalid_vertex_pairs(handle, *this, edge_first, edge_first + edge_srcs.size()); + CUGRAPH_EXPECTS(num_invalids == 0, + "Invalid input argument: there are invalid edge (src, dst) pairs."); + } + + auto [edge_indices, edge_partition_offsets] = + compute_edge_indices_and_edge_partition_offsets(handle, + *this, + store_transposed ? edge_dsts : edge_srcs, + store_transposed ? edge_srcs : edge_dsts); + + auto edge_mask_view = this->edge_mask_view(); + + auto sorted_edge_first = thrust::make_transform_iterator( + edge_indices.begin(), cugraph::detail::indirection_t{edge_first}); + rmm::device_uvector ret(edge_srcs.size(), handle.get_stream()); + + for (size_t i = 0; i < this->number_of_local_edge_partitions(); ++i) { + auto edge_partition = + edge_partition_device_view_t(this->local_edge_partition_view(i)); + auto edge_partition_e_mask = + edge_mask_view + ? thrust::make_optional< + detail::edge_partition_edge_property_device_view_t>( + *edge_mask_view, i) + : thrust::nullopt; + thrust::transform(handle.get_thrust_policy(), + sorted_edge_first + edge_partition_offsets[i], + sorted_edge_first + edge_partition_offsets[i + 1], + thrust::make_permutation_iterator( + ret.begin(), edge_indices.begin() + edge_partition_offsets[i]), + [edge_partition, edge_partition_e_mask] __device__(auto e) { + auto major = thrust::get<0>(e); + auto minor = thrust::get<1>(e); + auto major_idx = edge_partition.major_idx_from_major_nocheck(major); + if (major_idx) { + vertex_t const* indices{nullptr}; + edge_t local_edge_offset{}; + edge_t local_degree{}; + thrust::tie(indices, local_edge_offset, local_degree) = + edge_partition.local_edges(*major_idx); + auto it = thrust::lower_bound( + thrust::seq, indices, indices + local_degree, minor); + if ((it != indices + local_degree) && *it == minor) { + if (edge_partition_e_mask) { + return (*edge_partition_e_mask) + .get(local_edge_offset + thrust::distance(indices, it)); + } else { + return true; + } + } else { + return false; + } + } else { + return false; + } + }); + } + + return ret; +} + +template +rmm::device_uvector +graph_view_t>::has_edge( + raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check) +{ + CUGRAPH_EXPECTS( + edge_srcs.size() == edge_dsts.size(), + "Invalid input arguments: edge_srcs.size() does not coincide with edge_dsts.size()."); + + auto edge_first = + thrust::make_zip_iterator(store_transposed ? edge_dsts.begin() : edge_srcs.begin(), + store_transposed ? edge_srcs.begin() : edge_dsts.begin()); + + if (do_expensive_check) { + auto num_invalids = + detail::count_invalid_vertex_pairs(handle, *this, edge_first, edge_first + edge_srcs.size()); + CUGRAPH_EXPECTS(num_invalids == 0, + "Invalid input argument: there are invalid edge (src, dst) pairs."); + } + + auto edge_mask_view = this->edge_mask_view(); + + rmm::device_uvector ret(edge_srcs.size(), handle.get_stream()); + + auto edge_partition = + edge_partition_device_view_t(this->local_edge_partition_view()); + auto edge_partition_e_mask = + edge_mask_view + ? thrust::make_optional< + detail::edge_partition_edge_property_device_view_t>( + *edge_mask_view, 0) + : thrust::nullopt; + thrust::transform( + handle.get_thrust_policy(), + edge_first, + edge_first + edge_srcs.size(), + ret.begin(), + [edge_partition, edge_partition_e_mask] __device__(auto e) { + auto major = thrust::get<0>(e); + auto minor = thrust::get<1>(e); + auto major_offset = edge_partition.major_offset_from_major_nocheck(major); + vertex_t const* indices{nullptr}; + edge_t local_edge_offset{}; + edge_t local_degree{}; + thrust::tie(indices, local_edge_offset, local_degree) = + edge_partition.local_edges(major_offset); + auto it = thrust::lower_bound(thrust::seq, indices, indices + local_degree, minor); + if ((it != indices + local_degree) && *it == minor) { + if (edge_partition_e_mask) { + return (*edge_partition_e_mask).get(local_edge_offset + thrust::distance(indices, it)); + } else { + return true; + } + } else { + return false; + } + }); + + return ret; +} + +template +rmm::device_uvector +graph_view_t>:: + compute_multiplicity(raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check) +{ + CUGRAPH_EXPECTS(this->is_multigraph(), "Use has_edge() instead for non-multigraphs."); + CUGRAPH_EXPECTS( + edge_srcs.size() == edge_dsts.size(), + "Invalid input arguments: edge_srcs.size() does not coincide with edge_dsts.size()."); + + auto edge_first = + thrust::make_zip_iterator(store_transposed ? edge_dsts.begin() : edge_srcs.begin(), + store_transposed ? edge_srcs.begin() : edge_dsts.begin()); + + if (do_expensive_check) { + auto num_invalids = + detail::count_invalid_vertex_pairs(handle, *this, edge_first, edge_first + edge_srcs.size()); + CUGRAPH_EXPECTS(num_invalids == 0, + "Invalid input argument: there are invalid edge (src, dst) pairs."); + } + + auto [edge_indices, edge_partition_offsets] = + compute_edge_indices_and_edge_partition_offsets(handle, + *this, + store_transposed ? edge_dsts : edge_srcs, + store_transposed ? edge_srcs : edge_dsts); + + auto edge_mask_view = this->edge_mask_view(); + + auto sorted_edge_first = thrust::make_transform_iterator( + edge_indices.begin(), cugraph::detail::indirection_t{edge_first}); + rmm::device_uvector ret(edge_srcs.size(), handle.get_stream()); + + for (size_t i = 0; i < this->number_of_local_edge_partitions(); ++i) { + auto edge_partition = + edge_partition_device_view_t(this->local_edge_partition_view(i)); + auto edge_partition_e_mask = + edge_mask_view + ? thrust::make_optional< + detail::edge_partition_edge_property_device_view_t>( + *edge_mask_view, i) + : thrust::nullopt; + thrust::transform( + handle.get_thrust_policy(), + sorted_edge_first + edge_partition_offsets[i], + sorted_edge_first + edge_partition_offsets[i + 1], + thrust::make_permutation_iterator(ret.begin(), + edge_indices.begin() + edge_partition_offsets[i]), + [edge_partition, edge_partition_e_mask] __device__(auto e) { + auto major = thrust::get<0>(e); + auto minor = thrust::get<1>(e); + auto major_idx = edge_partition.major_idx_from_major_nocheck(major); + if (major_idx) { + vertex_t const* indices{nullptr}; + edge_t local_edge_offset{}; + edge_t local_degree{}; + thrust::tie(indices, local_edge_offset, local_degree) = + edge_partition.local_edges(*major_idx); + auto lower_it = thrust::lower_bound(thrust::seq, indices, indices + local_degree, minor); + auto upper_it = thrust::upper_bound(thrust::seq, indices, indices + local_degree, minor); + auto multiplicity = static_cast(thrust::distance(lower_it, upper_it)); + if (edge_partition_e_mask && (multiplicity > 0)) { + multiplicity = static_cast(detail::count_set_bits( + (*edge_partition_e_mask).value_first(), + static_cast(local_edge_offset + thrust::distance(indices, lower_it)), + static_cast(multiplicity))); + } + return multiplicity; + } else { + return edge_t{0}; + } + }); + } + + return ret; +} + +template +rmm::device_uvector +graph_view_t>:: + compute_multiplicity(raft::handle_t const& handle, + raft::device_span edge_srcs, + raft::device_span edge_dsts, + bool do_expensive_check) +{ + CUGRAPH_EXPECTS(this->is_multigraph(), "Use has_edge() instead for non-multigraphs."); + CUGRAPH_EXPECTS( + edge_srcs.size() == edge_dsts.size(), + "Invalid input arguments: edge_srcs.size() does not coincide with edge_dsts.size()."); + + auto edge_first = + thrust::make_zip_iterator(store_transposed ? edge_dsts.begin() : edge_srcs.begin(), + store_transposed ? edge_srcs.begin() : edge_dsts.begin()); + + if (do_expensive_check) { + auto num_invalids = + detail::count_invalid_vertex_pairs(handle, *this, edge_first, edge_first + edge_srcs.size()); + CUGRAPH_EXPECTS(num_invalids == 0, + "Invalid input argument: there are invalid edge (src, dst) pairs."); + } + + auto edge_mask_view = this->edge_mask_view(); + + rmm::device_uvector ret(edge_srcs.size(), handle.get_stream()); + + auto edge_partition = + edge_partition_device_view_t(this->local_edge_partition_view()); + auto edge_partition_e_mask = + edge_mask_view + ? thrust::make_optional< + detail::edge_partition_edge_property_device_view_t>( + *edge_mask_view, 0) + : thrust::nullopt; + thrust::transform( + handle.get_thrust_policy(), + edge_first, + edge_first + edge_srcs.size(), + ret.begin(), + [edge_partition, edge_partition_e_mask] __device__(auto e) { + auto major = thrust::get<0>(e); + auto minor = thrust::get<1>(e); + auto major_offset = edge_partition.major_offset_from_major_nocheck(major); + vertex_t const* indices{nullptr}; + edge_t local_edge_offset{}; + edge_t local_degree{}; + thrust::tie(indices, local_edge_offset, local_degree) = + edge_partition.local_edges(major_offset); + auto lower_it = thrust::lower_bound(thrust::seq, indices, indices + local_degree, minor); + auto upper_it = thrust::upper_bound(thrust::seq, indices, indices + local_degree, minor); + auto multiplicity = static_cast(thrust::distance(lower_it, upper_it)); + if (edge_partition_e_mask && (multiplicity > 0)) { + multiplicity = static_cast(detail::count_set_bits( + (*edge_partition_e_mask).value_first(), + static_cast(local_edge_offset + thrust::distance(indices, lower_it)), + static_cast(multiplicity))); + } + return multiplicity; + }); + + return ret; +} + } // namespace cugraph diff --git a/cpp/src/utilities/error_check_utils.cuh b/cpp/src/utilities/error_check_utils.cuh new file mode 100644 index 00000000000..baaf513d93d --- /dev/null +++ b/cpp/src/utilities/error_check_utils.cuh @@ -0,0 +1,137 @@ +/* + * 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 +#include +#include + +#include +#include +#include +#include + +#include +#include + +#include + +namespace cugraph { +namespace detail { + +// check vertices in the pair are in [0, num_vertices) and belongs to one of the local edge +// partitions. +template +struct is_invalid_input_vertex_pair_t { + vertex_t num_vertices{}; + raft::device_span edge_partition_major_range_firsts{}; + raft::device_span edge_partition_major_range_lasts{}; + vertex_t edge_partition_minor_range_first{}; + vertex_t edge_partition_minor_range_last{}; + + __device__ bool operator()(thrust::tuple pair) const + { + auto major = thrust::get<0>(pair); + auto minor = thrust::get<1>(pair); + if (!is_valid_vertex(num_vertices, major) || !is_valid_vertex(num_vertices, minor)) { + return true; + } + auto it = thrust::upper_bound(thrust::seq, + edge_partition_major_range_lasts.begin(), + edge_partition_major_range_lasts.end(), + major); + if (it == edge_partition_major_range_lasts.end()) { return true; } + auto edge_partition_idx = + static_cast(thrust::distance(edge_partition_major_range_lasts.begin(), it)); + if (major < edge_partition_major_range_firsts[edge_partition_idx]) { return true; } + return (minor < edge_partition_minor_range_first) || (minor >= edge_partition_minor_range_last); + } +}; + +template +size_t count_invalid_vertex_pairs(raft::handle_t const& handle, + GraphViewType const& graph_view, + VertexPairIterator vertex_pair_first, + VertexPairIterator vertex_pair_last) +{ + using vertex_t = typename GraphViewType::vertex_type; + + std::vector h_edge_partition_major_range_firsts( + graph_view.number_of_local_edge_partitions()); + std::vector h_edge_partition_major_range_lasts( + h_edge_partition_major_range_firsts.size()); + vertex_t edge_partition_minor_range_first{}; + vertex_t edge_partition_minor_range_last{}; + if constexpr (GraphViewType::is_multi_gpu) { + for (size_t i = 0; i < graph_view.number_of_local_edge_partitions(); i++) { + if constexpr (GraphViewType::is_storage_transposed) { + h_edge_partition_major_range_firsts[i] = graph_view.local_edge_partition_dst_range_first(i); + h_edge_partition_major_range_lasts[i] = graph_view.local_edge_partition_dst_range_last(i); + } else { + h_edge_partition_major_range_firsts[i] = graph_view.local_edge_partition_src_range_first(i); + h_edge_partition_major_range_lasts[i] = graph_view.local_edge_partition_src_range_last(i); + } + } + if constexpr (GraphViewType::is_storage_transposed) { + edge_partition_minor_range_first = graph_view.local_edge_partition_src_range_first(); + edge_partition_minor_range_last = graph_view.local_edge_partition_src_range_last(); + } else { + edge_partition_minor_range_first = graph_view.local_edge_partition_dst_range_first(); + edge_partition_minor_range_last = graph_view.local_edge_partition_dst_range_last(); + } + } else { + h_edge_partition_major_range_firsts[0] = vertex_t{0}; + h_edge_partition_major_range_lasts[0] = graph_view.number_of_vertices(); + edge_partition_minor_range_first = vertex_t{0}; + edge_partition_minor_range_last = graph_view.number_of_vertices(); + } + rmm::device_uvector d_edge_partition_major_range_firsts( + h_edge_partition_major_range_firsts.size(), handle.get_stream()); + rmm::device_uvector d_edge_partition_major_range_lasts( + h_edge_partition_major_range_lasts.size(), handle.get_stream()); + raft::update_device(d_edge_partition_major_range_firsts.data(), + h_edge_partition_major_range_firsts.data(), + h_edge_partition_major_range_firsts.size(), + handle.get_stream()); + raft::update_device(d_edge_partition_major_range_lasts.data(), + h_edge_partition_major_range_lasts.data(), + h_edge_partition_major_range_lasts.size(), + handle.get_stream()); + + auto num_invalid_pairs = thrust::count_if( + handle.get_thrust_policy(), + vertex_pair_first, + vertex_pair_last, + is_invalid_input_vertex_pair_t{ + graph_view.number_of_vertices(), + raft::device_span(d_edge_partition_major_range_firsts.begin(), + d_edge_partition_major_range_firsts.end()), + raft::device_span(d_edge_partition_major_range_lasts.begin(), + d_edge_partition_major_range_lasts.end()), + edge_partition_minor_range_first, + edge_partition_minor_range_last}); + if constexpr (GraphViewType::is_multi_gpu) { + auto& comm = handle.get_comms(); + num_invalid_pairs = + host_scalar_allreduce(comm, num_invalid_pairs, raft::comms::op_t::SUM, handle.get_stream()); + } + + return num_invalid_pairs; +} + +} // namespace detail +} // namespace cugraph diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d9d2f677abc..3df979fe5c2 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -313,6 +313,11 @@ ConfigureTest(DEGREE_TEST structure/degree_test.cpp) ConfigureTest(COUNT_SELF_LOOPS_AND_MULTI_EDGES_TEST "structure/count_self_loops_and_multi_edges_test.cpp") +################################################################################################### +# - Query edge existence and multiplicity tests --------------------------------------------------- +ConfigureTest(HAS_EDGE_AND_COMPUTE_MULTIPLICITY_TEST + "structure/has_edge_and_compute_multiplicity_test.cpp") + ################################################################################################### # - Coarsening tests ------------------------------------------------------------------------------ ConfigureTest(COARSEN_GRAPH_TEST structure/coarsen_graph_test.cpp) @@ -479,6 +484,11 @@ if(BUILD_CUGRAPH_MG_TESTS) ConfigureTestMG(MG_COUNT_SELF_LOOPS_AND_MULTI_EDGES_TEST "structure/mg_count_self_loops_and_multi_edges_test.cpp") + ############################################################################################### + # - MG Query edge existence and multiplicity tests -------------------------------------------- + ConfigureTestMG(MG_HAS_EDGE_AND_COMPUTE_MULTIPLICITY_TEST + "structure/mg_has_edge_and_compute_multiplicity_test.cpp") + ############################################################################################### # - MG PAGERANK tests ------------------------------------------------------------------------- ConfigureTestMG(MG_PAGERANK_TEST link_analysis/mg_pagerank_test.cpp) diff --git a/cpp/tests/structure/count_self_loops_and_multi_edges_test.cpp b/cpp/tests/structure/count_self_loops_and_multi_edges_test.cpp index 68828d5eee1..b7f1dce2023 100644 --- a/cpp/tests/structure/count_self_loops_and_multi_edges_test.cpp +++ b/cpp/tests/structure/count_self_loops_and_multi_edges_test.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2022, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -208,10 +208,7 @@ INSTANTIATE_TEST_SUITE_P( Tests_CountSelfLoopsAndMultiEdges_File, ::testing::Combine( // enable correctness checks - ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}), + ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{}), ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"), cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx")))); @@ -220,10 +217,7 @@ INSTANTIATE_TEST_SUITE_P( Tests_CountSelfLoopsAndMultiEdges_Rmat, ::testing::Combine( // enable correctness checks - ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}, - CountSelfLoopsAndMultiEdges_Usecase{}), + ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{}), ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false)))); INSTANTIATE_TEST_SUITE_P( @@ -235,10 +229,7 @@ INSTANTIATE_TEST_SUITE_P( Tests_CountSelfLoopsAndMultiEdges_Rmat, ::testing::Combine( // disable correctness checks for large graphs - ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{false}, - CountSelfLoopsAndMultiEdges_Usecase{false}, - CountSelfLoopsAndMultiEdges_Usecase{false}, - CountSelfLoopsAndMultiEdges_Usecase{false}), + ::testing::Values(CountSelfLoopsAndMultiEdges_Usecase{false}), ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false)))); CUGRAPH_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/structure/has_edge_and_compute_multiplicity_test.cpp b/cpp/tests/structure/has_edge_and_compute_multiplicity_test.cpp new file mode 100644 index 00000000000..3ad6953ca03 --- /dev/null +++ b/cpp/tests/structure/has_edge_and_compute_multiplicity_test.cpp @@ -0,0 +1,281 @@ +/* + * 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 governin_from_mtxg permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +struct HasEdgeAndComputeMultiplicity_Usecase { + size_t num_vertex_pairs{}; + bool check_correctness{true}; +}; + +template +class Tests_HasEdgeAndComputeMultiplicity + : public ::testing::TestWithParam< + std::tuple> { + public: + Tests_HasEdgeAndComputeMultiplicity() {} + + static void SetUpTestCase() {} + static void TearDownTestCase() {} + + virtual void SetUp() {} + virtual void TearDown() {} + + template + void run_current_test( + HasEdgeAndComputeMultiplicity_Usecase const& has_edge_and_compute_multiplicity_usecase, + input_usecase_t const& input_usecase) + { + using weight_t = float; + + constexpr bool renumber = true; + + raft::handle_t handle{}; + HighResTimer hr_timer{}; + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.start("Construct graph"); + } + + cugraph::graph_t graph(handle); + std::optional> d_renumber_map_labels{std::nullopt}; + std::tie(graph, std::ignore, d_renumber_map_labels) = + cugraph::test::construct_graph( + handle, input_usecase, false, renumber); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + auto graph_view = graph.view(); + + raft::random::RngState rng_state(0); + rmm::device_uvector edge_srcs( + has_edge_and_compute_multiplicity_usecase.num_vertex_pairs, handle.get_stream()); + rmm::device_uvector edge_dsts(edge_srcs.size(), handle.get_stream()); + cugraph::detail::uniform_random_fill(handle.get_stream(), + edge_srcs.data(), + edge_srcs.size(), + vertex_t{0}, + graph_view.number_of_vertices(), + rng_state); + cugraph::detail::uniform_random_fill(handle.get_stream(), + edge_dsts.data(), + edge_dsts.size(), + vertex_t{0}, + graph_view.number_of_vertices(), + rng_state); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.start("Querying edge existence"); + } + + auto edge_exists = + graph_view.has_edge(handle, + raft::device_span(edge_srcs.data(), edge_srcs.size()), + raft::device_span(edge_dsts.data(), edge_dsts.size())); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.start("Computing multiplicity"); + } + + auto edge_multiplicities = graph_view.compute_multiplicity( + handle, + raft::device_span(edge_srcs.data(), edge_srcs.size()), + raft::device_span(edge_dsts.data(), edge_dsts.size())); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + if (has_edge_and_compute_multiplicity_usecase.check_correctness) { + cugraph::graph_t unrenumbered_graph(handle); + if (renumber) { + std::tie(unrenumbered_graph, std::ignore, std::ignore) = + cugraph::test::construct_graph( + handle, input_usecase, false, false); + } + auto unrenumbered_graph_view = renumber ? unrenumbered_graph.view() : graph_view; + + std::vector h_offsets = cugraph::test::to_host( + handle, unrenumbered_graph_view.local_edge_partition_view().offsets()); + std::vector h_indices = cugraph::test::to_host( + handle, unrenumbered_graph_view.local_edge_partition_view().indices()); + + rmm::device_uvector d_unrenumbered_edge_srcs(edge_srcs.size(), handle.get_stream()); + rmm::device_uvector d_unrenumbered_edge_dsts(edge_dsts.size(), handle.get_stream()); + raft::copy_async( + d_unrenumbered_edge_srcs.data(), edge_srcs.data(), edge_srcs.size(), handle.get_stream()); + raft::copy_async( + d_unrenumbered_edge_dsts.data(), edge_dsts.data(), edge_dsts.size(), handle.get_stream()); + if (renumber) { + cugraph::unrenumber_local_int_vertices(handle, + d_unrenumbered_edge_srcs.data(), + d_unrenumbered_edge_srcs.size(), + (*d_renumber_map_labels).data(), + vertex_t{0}, + graph_view.number_of_vertices()); + cugraph::unrenumber_local_int_vertices(handle, + d_unrenumbered_edge_dsts.data(), + d_unrenumbered_edge_dsts.size(), + (*d_renumber_map_labels).data(), + vertex_t{0}, + graph_view.number_of_vertices()); + } + auto h_unrenumbered_edge_srcs = cugraph::test::to_host(handle, d_unrenumbered_edge_srcs); + auto h_unrenumbered_edge_dsts = cugraph::test::to_host(handle, d_unrenumbered_edge_dsts); + + auto h_cugraph_edge_exists = cugraph::test::to_host(handle, edge_exists); + auto h_cugraph_edge_multiplicities = cugraph::test::to_host(handle, edge_multiplicities); + std::vector h_reference_edge_exists(edge_srcs.size()); + std::vector h_reference_edge_multiplicities(edge_srcs.size()); + for (size_t i = 0; i < edge_srcs.size(); ++i) { + auto src = h_unrenumbered_edge_srcs[i]; + auto dst = h_unrenumbered_edge_dsts[i]; + auto major = store_transposed ? dst : src; + auto minor = store_transposed ? src : dst; + auto lower_it = std::lower_bound( + h_indices.begin() + h_offsets[major], h_indices.begin() + h_offsets[major + 1], minor); + auto upper_it = std::upper_bound( + h_indices.begin() + h_offsets[major], h_indices.begin() + h_offsets[major + 1], minor); + auto multiplicity = static_cast(std::distance(lower_it, upper_it)); + h_reference_edge_exists[i] = multiplicity > 0 ? true : false; + h_reference_edge_multiplicities[i] = multiplicity; + } + + ASSERT_TRUE(std::equal(h_reference_edge_exists.begin(), + h_reference_edge_exists.end(), + h_cugraph_edge_exists.begin())) + << "has_edge() return values do not match with the reference values."; + ASSERT_TRUE(std::equal(h_reference_edge_multiplicities.begin(), + h_reference_edge_multiplicities.end(), + h_cugraph_edge_multiplicities.begin())) + << "compute_multiplicity() return values do not match with the reference values."; + } + } +}; + +using Tests_HasEdgeAndComputeMultiplicity_File = + Tests_HasEdgeAndComputeMultiplicity; +using Tests_HasEdgeAndComputeMultiplicity_Rmat = + Tests_HasEdgeAndComputeMultiplicity; + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_File, CheckInt32Int32FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test(std::get<0>(param), std::get<1>(param)); +} + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int32FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int64FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_Rmat, CheckInt64Int64FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_File, CheckInt32Int32FloatTransposeTrue) +{ + auto param = GetParam(); + run_current_test(std::get<0>(param), std::get<1>(param)); +} + +TEST_P(Tests_HasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int32FloatTransposeTrue) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +INSTANTIATE_TEST_SUITE_P( + file_test, + Tests_HasEdgeAndComputeMultiplicity_File, + ::testing::Combine( + // enable correctness checks + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 128}), + ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"), + cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx")))); + +INSTANTIATE_TEST_SUITE_P( + rmat_small_test, + Tests_HasEdgeAndComputeMultiplicity_Rmat, + ::testing::Combine( + // enable correctness checks + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 128}), + ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false)))); + +INSTANTIATE_TEST_SUITE_P( + rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with + --gtest_filter to select only the rmat_benchmark_test with a specific + vertex & edge type combination) by command line arguments and do not + include more than one Rmat_Usecase that differ only in scale or edge + factor (to avoid running same benchmarks more than once) */ + Tests_HasEdgeAndComputeMultiplicity_Rmat, + ::testing::Combine( + // disable correctness checks for large graphs + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 1024 * 128, false}), + ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false)))); + +CUGRAPH_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/structure/mg_has_edge_and_compute_multiplicity_test.cpp b/cpp/tests/structure/mg_has_edge_and_compute_multiplicity_test.cpp new file mode 100644 index 00000000000..8079de7ebfe --- /dev/null +++ b/cpp/tests/structure/mg_has_edge_and_compute_multiplicity_test.cpp @@ -0,0 +1,331 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include + +struct HasEdgeAndComputeMultiplicity_Usecase { + size_t num_vertex_pairs{}; + bool check_correctness{true}; +}; + +template +class Tests_MGHasEdgeAndComputeMultiplicity + : public ::testing::TestWithParam< + std::tuple> { + public: + Tests_MGHasEdgeAndComputeMultiplicity() {} + + static void SetUpTestCase() { handle_ = cugraph::test::initialize_mg_handle(); } + + static void TearDownTestCase() { handle_.reset(); } + + virtual void SetUp() {} + virtual void TearDown() {} + + // Compare the results of running has_edge & compute_multiplicity on multiple GPUs to that of + // a single-GPU run + template + void run_current_test( + HasEdgeAndComputeMultiplicity_Usecase const& has_edge_and_compute_multiplicity_usecase, + input_usecase_t const& input_usecase) + { + using weight_t = float; + using edge_type_id_t = int32_t; + + HighResTimer hr_timer{}; + + auto const comm_rank = handle_->get_comms().get_rank(); + auto const comm_size = handle_->get_comms().get_size(); + + // 1. create MG graph + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.start("MG Construct graph"); + } + + cugraph::graph_t mg_graph(*handle_); + std::optional> mg_renumber_map{std::nullopt}; + std::tie(mg_graph, std::ignore, mg_renumber_map) = + cugraph::test::construct_graph( + *handle_, input_usecase, false, true); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + auto mg_graph_view = mg_graph.view(); + + // 2. create an edge list to query + + raft::random::RngState rng_state(comm_rank); + size_t num_vertex_pairs_this_gpu = + (has_edge_and_compute_multiplicity_usecase.num_vertex_pairs / comm_size) + + ((comm_rank < has_edge_and_compute_multiplicity_usecase.num_vertex_pairs % comm_size) + ? size_t{1} + : size_t{0}); + rmm::device_uvector d_mg_edge_srcs(num_vertex_pairs_this_gpu, handle_->get_stream()); + rmm::device_uvector d_mg_edge_dsts(d_mg_edge_srcs.size(), handle_->get_stream()); + cugraph::detail::uniform_random_fill(handle_->get_stream(), + d_mg_edge_srcs.data(), + d_mg_edge_srcs.size(), + vertex_t{0}, + mg_graph_view.number_of_vertices(), + rng_state); + cugraph::detail::uniform_random_fill(handle_->get_stream(), + d_mg_edge_dsts.data(), + d_mg_edge_dsts.size(), + vertex_t{0}, + mg_graph_view.number_of_vertices(), + rng_state); + + std::tie(store_transposed ? d_mg_edge_dsts : d_mg_edge_srcs, + store_transposed ? d_mg_edge_srcs : d_mg_edge_dsts, + std::ignore, + std::ignore, + std::ignore) = + cugraph::detail::shuffle_int_vertex_pairs_with_values_to_local_gpu_by_edge_partitioning< + vertex_t, + edge_t, + weight_t, + edge_type_id_t>(*handle_, + std::move(store_transposed ? d_mg_edge_dsts : d_mg_edge_srcs), + std::move(store_transposed ? d_mg_edge_srcs : d_mg_edge_dsts), + std::nullopt, + std::nullopt, + std::nullopt, + mg_graph_view.vertex_partition_range_lasts()); + + // 3. run MG has_edge & compute_multiplicity + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.start("MG Querying edge existence"); + } + + auto d_mg_edge_exists = mg_graph_view.has_edge( + *handle_, + raft::device_span(d_mg_edge_srcs.data(), d_mg_edge_srcs.size()), + raft::device_span(d_mg_edge_dsts.data(), d_mg_edge_dsts.size())); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.start("MG Computing multiplicity"); + } + + auto d_mg_edge_multiplicities = mg_graph_view.compute_multiplicity( + *handle_, + raft::device_span(d_mg_edge_srcs.data(), d_mg_edge_srcs.size()), + raft::device_span(d_mg_edge_dsts.data(), d_mg_edge_dsts.size())); + + if (cugraph::test::g_perf) { + RAFT_CUDA_TRY(cudaDeviceSynchronize()); // for consistent performance measurement + handle_->get_comms().barrier(); + hr_timer.stop(); + hr_timer.display_and_clear(std::cout); + } + + // 4. copmare SG & MG results + + if (has_edge_and_compute_multiplicity_usecase.check_correctness) { + // 4-1. aggregate MG results + + cugraph::unrenumber_int_vertices( + *handle_, + d_mg_edge_srcs.data(), + d_mg_edge_srcs.size(), + (*mg_renumber_map).data(), + mg_graph_view.vertex_partition_range_lasts()); + cugraph::unrenumber_int_vertices( + *handle_, + d_mg_edge_dsts.data(), + d_mg_edge_dsts.size(), + (*mg_renumber_map).data(), + mg_graph_view.vertex_partition_range_lasts()); + + auto d_mg_aggregate_edge_srcs = cugraph::test::device_gatherv( + *handle_, raft::device_span(d_mg_edge_srcs.data(), d_mg_edge_srcs.size())); + auto d_mg_aggregate_edge_dsts = cugraph::test::device_gatherv( + *handle_, raft::device_span(d_mg_edge_dsts.data(), d_mg_edge_dsts.size())); + auto d_mg_aggregate_edge_exists = cugraph::test::device_gatherv( + *handle_, raft::device_span(d_mg_edge_exists.data(), d_mg_edge_exists.size())); + auto d_mg_aggregate_edge_multiplicities = cugraph::test::device_gatherv( + *handle_, + raft::device_span(d_mg_edge_multiplicities.data(), + d_mg_edge_multiplicities.size())); + + cugraph::graph_t sg_graph(*handle_); + std::tie(sg_graph, std::ignore, std::ignore) = cugraph::test::mg_graph_to_sg_graph( + *handle_, + mg_graph_view, + std::optional>{std::nullopt}, + std::make_optional>((*mg_renumber_map).data(), + (*mg_renumber_map).size()), + false); + + if (handle_->get_comms().get_rank() == 0) { + auto sg_graph_view = sg_graph.view(); + + // 4-2. run SG count_self_loops & count_multi_edges + + auto d_sg_edge_exists = sg_graph_view.has_edge( + *handle_, + raft::device_span(d_mg_aggregate_edge_srcs.data(), + d_mg_aggregate_edge_srcs.size()), + raft::device_span(d_mg_aggregate_edge_dsts.data(), + d_mg_aggregate_edge_dsts.size())); + auto d_sg_edge_multiplicities = sg_graph_view.compute_multiplicity( + *handle_, + raft::device_span(d_mg_aggregate_edge_srcs.data(), + d_mg_aggregate_edge_srcs.size()), + raft::device_span(d_mg_aggregate_edge_dsts.data(), + d_mg_aggregate_edge_dsts.size())); + + // 4-3. compare + + auto h_mg_aggregate_edge_exists = + cugraph::test::to_host(*handle_, d_mg_aggregate_edge_exists); + auto h_mg_aggregate_edge_multiplicities = + cugraph::test::to_host(*handle_, d_mg_aggregate_edge_multiplicities); + auto h_sg_edge_exists = cugraph::test::to_host(*handle_, d_sg_edge_exists); + auto h_sg_edge_multiplicities = cugraph::test::to_host(*handle_, d_sg_edge_multiplicities); + + ASSERT_TRUE(std::equal(h_mg_aggregate_edge_exists.begin(), + h_mg_aggregate_edge_exists.end(), + h_sg_edge_exists.begin())); + ASSERT_TRUE(std::equal(h_mg_aggregate_edge_multiplicities.begin(), + h_mg_aggregate_edge_multiplicities.end(), + h_sg_edge_multiplicities.begin())); + } + } + } + + private: + static std::unique_ptr handle_; +}; + +template +std::unique_ptr Tests_MGHasEdgeAndComputeMultiplicity::handle_ = + nullptr; + +using Tests_MGHasEdgeAndComputeMultiplicity_File = + Tests_MGHasEdgeAndComputeMultiplicity; +using Tests_MGHasEdgeAndComputeMultiplicity_Rmat = + Tests_MGHasEdgeAndComputeMultiplicity; + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_File, CheckInt32Int32FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test(std::get<0>(param), std::get<1>(param)); +} + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int32FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int64FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_Rmat, CheckInt64Int64FloatTransposeFalse) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_File, CheckInt32Int32FloatTransposeTrue) +{ + auto param = GetParam(); + run_current_test(std::get<0>(param), std::get<1>(param)); +} + +TEST_P(Tests_MGHasEdgeAndComputeMultiplicity_Rmat, CheckInt32Int32FloatTransposeTrue) +{ + auto param = GetParam(); + run_current_test( + std::get<0>(param), override_Rmat_Usecase_with_cmd_line_arguments(std::get<1>(param))); +} + +INSTANTIATE_TEST_SUITE_P( + file_tests, + Tests_MGHasEdgeAndComputeMultiplicity_File, + ::testing::Combine( + // enable correctness checks + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 128}), + ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"), + cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx")))); + +INSTANTIATE_TEST_SUITE_P( + rmat_small_tests, + Tests_MGHasEdgeAndComputeMultiplicity_Rmat, + ::testing::Combine( + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 128}), + ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false)))); + +INSTANTIATE_TEST_SUITE_P( + rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with + --gtest_filter to select only the rmat_benchmark_test with a specific + vertex & edge type combination) by command line arguments and do not + include more than one Rmat_Usecase that differ only in scale or edge + factor (to avoid running same benchmarks more than once) */ + Tests_MGHasEdgeAndComputeMultiplicity_Rmat, + ::testing::Combine( + ::testing::Values(HasEdgeAndComputeMultiplicity_Usecase{1024 * 1024 * 128, false}), + ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false)))); + +CUGRAPH_MG_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/utilities/device_comm_wrapper.cu b/cpp/tests/utilities/device_comm_wrapper.cu index cfc65b5d741..50727394ad7 100644 --- a/cpp/tests/utilities/device_comm_wrapper.cu +++ b/cpp/tests/utilities/device_comm_wrapper.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2022, NVIDIA CORPORATION. + * Copyright (c) 2021-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. @@ -40,9 +40,10 @@ rmm::device_uvector device_gatherv(raft::handle_t const& handle, rmm::device_uvector gathered_v( is_root ? std::reduce(rx_sizes.begin(), rx_sizes.end()) : size_t{0}, handle.get_stream()); + using comm_datatype_t = std::conditional_t, uint8_t, T>; cugraph::device_gatherv(handle.get_comms(), - d_input.data(), - gathered_v.data(), + reinterpret_cast(d_input.data()), + reinterpret_cast(gathered_v.data()), d_input.size(), rx_sizes, rx_displs, @@ -64,9 +65,10 @@ rmm::device_uvector device_allgatherv(raft::handle_t const& handle, rmm::device_uvector gathered_v(std::reduce(rx_sizes.begin(), rx_sizes.end()), handle.get_stream()); + using comm_datatype_t = std::conditional_t, uint8_t, T>; cugraph::device_allgatherv(handle.get_comms(), - d_input.data(), - gathered_v.data(), + reinterpret_cast(d_input.data()), + reinterpret_cast(gathered_v.data()), rx_sizes, rx_displs, handle.get_stream()); @@ -76,6 +78,9 @@ rmm::device_uvector device_allgatherv(raft::handle_t const& handle, // explicit instantiation +template rmm::device_uvector device_gatherv(raft::handle_t const& handle, + raft::device_span d_input); + template rmm::device_uvector device_gatherv(raft::handle_t const& handle, raft::device_span d_input); @@ -91,6 +96,9 @@ template rmm::device_uvector device_gatherv(raft::handle_t const& handle, template rmm::device_uvector device_gatherv(raft::handle_t const& handle, raft::device_span d_input); +template rmm::device_uvector device_allgatherv(raft::handle_t const& handle, + raft::device_span d_input); + template rmm::device_uvector device_allgatherv(raft::handle_t const& handle, raft::device_span d_input); diff --git a/cpp/tests/utilities/test_utilities.hpp b/cpp/tests/utilities/test_utilities.hpp index 321a0536e02..3fa6ae089d3 100644 --- a/cpp/tests/utilities/test_utilities.hpp +++ b/cpp/tests/utilities/test_utilities.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019-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. @@ -377,18 +377,24 @@ template std::vector to_host(raft::handle_t const& handle, raft::device_span data) { std::vector h_data(data.size()); - raft::update_host(h_data.data(), data.data(), data.size(), handle.get_stream()); - handle.sync_stream(); + if constexpr (std::is_same_v) { // std::vector stores values in a packed format + auto h_tmp = new bool[data.size()]; + raft::update_host(h_tmp, data.data(), data.size(), handle.get_stream()); + handle.sync_stream(); + std::transform( + h_tmp, h_tmp + data.size(), h_data.begin(), [](uint8_t v) { return static_cast(v); }); + delete[] h_tmp; + } else { + raft::update_host(h_data.data(), data.data(), data.size(), handle.get_stream()); + handle.sync_stream(); + } return h_data; } template std::vector to_host(raft::handle_t const& handle, rmm::device_uvector const& data) { - std::vector h_data(data.size()); - raft::update_host(h_data.data(), data.data(), data.size(), handle.get_stream()); - handle.sync_stream(); - return h_data; + return to_host(handle, raft::device_span(data.data(), data.size())); } template @@ -396,11 +402,7 @@ std::optional> to_host(raft::handle_t const& handle, std::optional> data) { std::optional> h_data{std::nullopt}; - if (data) { - h_data = std::vector((*data).size()); - raft::update_host((*h_data).data(), (*data).data(), (*data).size(), handle.get_stream()); - handle.sync_stream(); - } + if (data) { h_data = to_host(handle, *data); } return h_data; } @@ -410,9 +412,7 @@ std::optional> to_host(raft::handle_t const& handle, { std::optional> h_data{std::nullopt}; if (data) { - h_data = std::vector((*data).size()); - raft::update_host((*h_data).data(), (*data).data(), (*data).size(), handle.get_stream()); - handle.sync_stream(); + h_data = to_host(handle, raft::device_span((*data).data(), (*data).size())); } return h_data; } @@ -430,8 +430,16 @@ template rmm::device_uvector to_device(raft::handle_t const& handle, std::vector const& data) { rmm::device_uvector d_data(data.size(), handle.get_stream()); - raft::update_device(d_data.data(), data.data(), data.size(), handle.get_stream()); - handle.sync_stream(); + if constexpr (std::is_same_v) { // std::vector stores values in a packed format + auto h_tmp = new bool[data.size()]; + std::copy(data.begin(), data.end(), h_tmp); + raft::update_device(d_data.data(), h_tmp, h_tmp + data.size(), handle.get_stream()); + handle.sync_stream(); + delete[] h_tmp; + } else { + raft::update_device(d_data.data(), data.data(), data.size(), handle.get_stream()); + handle.sync_stream(); + } return d_data; } @@ -453,11 +461,7 @@ std::optional> to_device(raft::handle_t const& handle, std::optional> const& data) { std::optional> d_data{std::nullopt}; - if (data) { - d_data = rmm::device_uvector(data->size(), handle.get_stream()); - raft::update_host(d_data->data(), data->data(), data->size(), handle.get_stream()); - handle.sync_stream(); - } + if (data) { d_data = to_device(handle, *data); } return d_data; } From 9a261ffd1f5e055e9c8b751c6009f99e9c39d0c1 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 25 Jan 2024 14:05:45 -0600 Subject: [PATCH 3/5] nx-cugraph: add `complement` and `reverse` (#4103) We apparently already had `G.reverse()`, which made that function extra easy :) Authors: - Erik Welch (https://github.com/eriknw) Approvers: - Rick Ratzel (https://github.com/rlratzel) URL: https://github.com/rapidsai/cugraph/pull/4103 --- python/nx-cugraph/_nx_cugraph/__init__.py | 2 + .../nx_cugraph/algorithms/__init__.py | 2 + .../algorithms/operators/__init__.py | 13 +++++ .../nx_cugraph/algorithms/operators/unary.py | 55 +++++++++++++++++++ 4 files changed, 72 insertions(+) create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py index 9bca031a2f0..2f283aa153c 100644 --- a/python/nx-cugraph/_nx_cugraph/__init__.py +++ b/python/nx-cugraph/_nx_cugraph/__init__.py @@ -43,6 +43,7 @@ "chvatal_graph", "circular_ladder_graph", "clustering", + "complement", "complete_bipartite_graph", "complete_graph", "complete_multipartite_graph", @@ -105,6 +106,7 @@ "path_graph", "petersen_graph", "reciprocity", + "reverse", "sedgewick_maze_graph", "single_source_shortest_path_length", "single_target_shortest_path_length", diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py index 08658ad94cb..7aafa85f5b7 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py @@ -17,6 +17,7 @@ community, components, link_analysis, + operators, shortest_paths, traversal, tree, @@ -29,6 +30,7 @@ from .dag import * from .isolate import * from .link_analysis import * +from .operators import * from .reciprocity import * from .shortest_paths import * from .traversal import * diff --git a/python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py new file mode 100644 index 00000000000..32fd45f5726 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py @@ -0,0 +1,13 @@ +# 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. +from .unary import * diff --git a/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py b/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py new file mode 100644 index 00000000000..08abc9f2872 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py @@ -0,0 +1,55 @@ +# 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. +import cupy as cp +import networkx as nx +import numpy as np + +import nx_cugraph as nxcg +from nx_cugraph.convert import _to_graph +from nx_cugraph.utils import index_dtype, networkx_algorithm + +__all__ = ["complement", "reverse"] + + +@networkx_algorithm(version_added="24.02") +def complement(G): + G = _to_graph(G) + N = G._N + # Upcast to int64 so indices don't overflow. + edges_a_b = N * G.src_indices.astype(np.int64) + G.dst_indices + # Now compute flattened indices for all edges except self-loops + # Alt (slower): + # edges_full = np.arange(N * N) + # edges_full = edges_full[(edges_full % (N + 1)).astype(bool)] + edges_full = cp.arange(1, N * (N - 1) + 1) + cp.repeat(cp.arange(N - 1), N) + edges_comp = cp.setdiff1d( + edges_full, + edges_a_b, + assume_unique=not G.is_multigraph(), + ) + src_indices, dst_indices = cp.divmod(edges_comp, N) + return G.__class__.from_coo( + N, + src_indices.astype(index_dtype), + dst_indices.astype(index_dtype), + key_to_id=G.key_to_id, + ) + + +@networkx_algorithm(version_added="24.02") +def reverse(G, copy=True): + if not G.is_directed(): + raise nx.NetworkXError("Cannot reverse an undirected graph.") + if isinstance(G, nx.Graph): + G = nxcg.from_networkx(G, preserve_all_attrs=True) + return G.reverse(copy=copy) From ff76a385c4048ac75d1408c3bcf9d4e3018a88b0 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang <45857425+seunghwak@users.noreply.github.com> Date: Thu, 25 Jan 2024 13:13:58 -0800 Subject: [PATCH 4/5] Update per_v_transform_reduce_incoming|outgoing_e to support edge masking (#4085) per_v_transform_reduce_(incoming|outgoing_e) now supports edge masking. Authors: - Seunghwa Kang (https://github.com/seunghwak) Approvers: - Chuck Hastings (https://github.com/ChuckHastings) - Naim (https://github.com/naimnv) - Joseph Nke (https://github.com/jnke2016) URL: https://github.com/rapidsai/cugraph/pull/4085 --- .../eigenvector_centrality_impl.cuh | 4 +- cpp/src/link_analysis/pagerank_impl.cuh | 4 +- cpp/src/prims/detail/prim_functors.cuh | 60 +++ ...v_transform_reduce_incoming_outgoing_e.cuh | 442 +++++++++++------- ...rm_reduce_v_frontier_outgoing_e_by_dst.cuh | 18 +- .../traversal/od_shortest_distances_impl.cuh | 17 +- ..._v_transform_reduce_incoming_outgoing_e.cu | 27 +- 7 files changed, 381 insertions(+), 191 deletions(-) create mode 100644 cpp/src/prims/detail/prim_functors.cuh diff --git a/cpp/src/centrality/eigenvector_centrality_impl.cuh b/cpp/src/centrality/eigenvector_centrality_impl.cuh index 8d1bea4004d..2129dca6985 100644 --- a/cpp/src/centrality/eigenvector_centrality_impl.cuh +++ b/cpp/src/centrality/eigenvector_centrality_impl.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-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. @@ -117,7 +117,7 @@ rmm::device_uvector eigenvector_centrality( edge_src_centralities.view(), edge_dst_dummy_property_t{}.view(), edge_dummy_property_t{}.view(), - [] __device__(vertex_t, vertex_t, auto src_val, auto, auto) { return src_val * 1.0; }, + [] __device__(vertex_t, vertex_t, auto src_val, auto, auto) { return src_val; }, weight_t{0}, reduce_op::plus{}, centralities.begin()); diff --git a/cpp/src/link_analysis/pagerank_impl.cuh b/cpp/src/link_analysis/pagerank_impl.cuh index 92c70fcff20..9a76ba73f92 100644 --- a/cpp/src/link_analysis/pagerank_impl.cuh +++ b/cpp/src/link_analysis/pagerank_impl.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -288,7 +288,7 @@ centrality_algorithm_metadata_t pagerank( edge_dst_dummy_property_t{}.view(), edge_dummy_property_t{}.view(), [alpha] __device__(vertex_t, vertex_t, auto src_val, auto, auto) { - return src_val * 1.0 * alpha; + return src_val * alpha; }, unvarying_part, reduce_op::plus{}, diff --git a/cpp/src/prims/detail/prim_functors.cuh b/cpp/src/prims/detail/prim_functors.cuh new file mode 100644 index 00000000000..2785ba38dfd --- /dev/null +++ b/cpp/src/prims/detail/prim_functors.cuh @@ -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 + +namespace cugraph { + +namespace detail { + +template +struct call_e_op_t { + edge_partition_device_view_t 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 diff --git a/cpp/src/prims/per_v_transform_reduce_incoming_outgoing_e.cuh b/cpp/src/prims/per_v_transform_reduce_incoming_outgoing_e.cuh index 1a7fc0130c4..24b4f0857b1 100644 --- a/cpp/src/prims/per_v_transform_reduce_incoming_outgoing_e.cuh +++ b/cpp/src/prims/per_v_transform_reduce_incoming_outgoing_e.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -16,6 +16,7 @@ #pragma once #include +#include #include #include #include @@ -63,11 +64,84 @@ namespace detail { int32_t constexpr per_v_transform_reduce_e_kernel_block_size = 512; +template +struct transform_and_atomic_reduce_t { + edge_partition_device_view_t const& edge_partition{}; + result_t identity_element{}; + vertex_t const* indices{nullptr}; + TransformOp const& transform_op{}; + ResultValueOutputIteratorOrWrapper& result_value_output{}; + + __device__ void operator()(edge_t i) const + { + auto e_op_result = transform_op(i); + if (e_op_result != identity_element) { + auto minor = indices[i]; + auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); + if constexpr (multi_gpu) { + reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + } else { + reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + } + } + } +}; + +template +__device__ void update_result_value_output( + edge_partition_device_view_t const& edge_partition, + vertex_t const* indices, + edge_t local_degree, + TransformOp const& transform_op, + result_t init, + ReduceOp const& reduce_op, + size_t output_idx /* relevent only when update_major === true */, + result_t identity_element, + ResultValueOutputIteratorOrWrapper& result_value_output) +{ + if constexpr (update_major) { + *(result_value_output + output_idx) = + thrust::transform_reduce(thrust::seq, + thrust::make_counting_iterator(edge_t{0}), + thrust::make_counting_iterator(local_degree), + transform_op, + init, + reduce_op); + } else { + thrust::for_each( + thrust::seq, + thrust::make_counting_iterator(edge_t{0}), + thrust::make_counting_iterator(local_degree), + transform_and_atomic_reduce_t{ + edge_partition, identity_element, indices, transform_op, result_value_output}); + } +} + template edge_partition_e_mask, ResultValueOutputIteratorOrWrapper result_value_output, EdgeOp e_op, T init /* relevant only if update_major == true */, + T identity_element /* relevant only if update_major == true */, ReduceOp reduce_op) { static_assert(update_major || reduce_op::has_compatible_raft_comms_op_v< @@ -104,6 +180,7 @@ __global__ void per_v_transform_reduce_e_hypersparse( while (idx < static_cast(dcs_nzd_vertex_count)) { auto major = *(edge_partition.major_from_major_hypersparse_idx_nocheck(static_cast(idx))); + auto major_offset = edge_partition.major_offset_from_major_nocheck(major); auto major_idx = major_start_offset + idx; // major_offset != major_idx in the hypersparse region vertex_t const* indices{nullptr}; @@ -111,60 +188,50 @@ __global__ void per_v_transform_reduce_e_hypersparse( edge_t local_degree{}; thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(static_cast(major_idx)); - auto transform_op = [&edge_partition, - &edge_partition_src_value_input, - &edge_partition_dst_value_input, - &edge_partition_e_value_input, - &e_op, - major, - indices, - edge_offset] __device__(auto i) { - auto major_offset = edge_partition.major_offset_from_major_nocheck(major); - 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)); - }; - if constexpr (update_major) { - *(result_value_output + (major - *(edge_partition.major_hypersparse_first()))) = - thrust::transform_reduce(thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - transform_op, - init, - reduce_op); + auto call_e_op = call_e_op_t{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 = + [&edge_partition_e_mask, &call_e_op, identity_element, edge_offset] __device__(auto i) { + if ((*edge_partition_e_mask).get(edge_offset + i)) { + return call_e_op(i); + } else { + return identity_element; + } + }; + + update_result_value_output(edge_partition, + indices, + local_degree, + transform_op, + init, + reduce_op, + major - *(edge_partition).major_hypersparse_first(), + identity_element, + result_value_output); } else { - if constexpr (GraphViewType::is_multi_gpu) { - thrust::for_each( - thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - [&edge_partition, indices, &result_value_output, &transform_op] __device__(auto i) { - auto e_op_result = transform_op(i); - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); - }); - } else { - thrust::for_each( - thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - [&edge_partition, indices, &result_value_output, &transform_op] __device__(auto i) { - auto e_op_result = transform_op(i); - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); - }); - } + update_result_value_output(edge_partition, + indices, + local_degree, + call_e_op, + init, + reduce_op, + major - *(edge_partition).major_hypersparse_first(), + identity_element, + result_value_output); } idx += gridDim.x * blockDim.x; } @@ -175,6 +242,7 @@ template edge_partition_e_mask, ResultValueOutputIteratorOrWrapper result_value_output, EdgeOp e_op, T init /* relevant only if update_major == true */, + T identity_element /* relevant only if update_major == true */, ReduceOp reduce_op) { static_assert(update_major || reduce_op::has_compatible_raft_comms_op_v< @@ -209,71 +279,57 @@ __global__ void per_v_transform_reduce_e_low_degree( auto idx = static_cast(tid); while (idx < static_cast(major_range_last - major_range_first)) { - auto major_offset = major_start_offset + idx; + auto major_offset = static_cast(major_start_offset + idx); + auto major = edge_partition.major_from_major_offset_nocheck(major_offset); vertex_t const* indices{nullptr}; edge_t edge_offset{}; edge_t local_degree{}; thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(static_cast(major_offset)); - auto transform_op = [&edge_partition, - &edge_partition_src_value_input, - &edge_partition_dst_value_input, - &edge_partition_e_value_input, - &e_op, - major_offset, - indices, - edge_offset] __device__(auto i) { - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - auto src = GraphViewType::is_storage_transposed - ? minor - : edge_partition.major_from_major_offset_nocheck(major_offset); - auto dst = GraphViewType::is_storage_transposed - ? edge_partition.major_from_major_offset_nocheck(major_offset) - : minor; - auto src_offset = - GraphViewType::is_storage_transposed ? minor_offset : static_cast(major_offset); - auto dst_offset = - GraphViewType::is_storage_transposed ? static_cast(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)); - }; - if constexpr (update_major) { - *(result_value_output + idx) = - thrust::transform_reduce(thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - transform_op, - init, - reduce_op); + auto call_e_op = call_e_op_t{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 = + [&edge_partition_e_mask, &call_e_op, identity_element, edge_offset] __device__(auto i) { + if ((*edge_partition_e_mask).get(edge_offset + i)) { + return call_e_op(i); + } else { + return identity_element; + } + }; + + update_result_value_output(edge_partition, + indices, + local_degree, + transform_op, + init, + reduce_op, + idx, + identity_element, + result_value_output); } else { - if constexpr (GraphViewType::is_multi_gpu) { - thrust::for_each( - thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - [&edge_partition, indices, &result_value_output, &transform_op] __device__(auto i) { - auto e_op_result = transform_op(i); - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); - }); - } else { - thrust::for_each( - thrust::seq, - thrust::make_counting_iterator(edge_t{0}), - thrust::make_counting_iterator(local_degree), - [&edge_partition, indices, &result_value_output, &transform_op] __device__(auto i) { - auto e_op_result = transform_op(i); - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); - }); - } + update_result_value_output(edge_partition, + indices, + local_degree, + call_e_op, + init, + reduce_op, + idx, + identity_element, + result_value_output); } idx += gridDim.x * blockDim.x; } @@ -284,6 +340,7 @@ template edge_partition_e_mask, ResultValueOutputIteratorOrWrapper result_value_output, EdgeOp e_op, T init /* relevant only if update_major == true */, @@ -327,41 +385,61 @@ __global__ void per_v_transform_reduce_e_mid_degree( raft::warp_size()]; // relevant only if update_major == true while (idx < static_cast(major_range_last - major_range_first)) { - auto major_offset = major_start_offset + idx; + auto major_offset = static_cast(major_start_offset + idx); + auto major = edge_partition.major_from_major_offset_nocheck(major_offset); vertex_t const* indices{nullptr}; edge_t edge_offset{}; edge_t local_degree{}; thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset); + + auto call_e_op = call_e_op_t{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 - for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) { - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - auto src = GraphViewType::is_storage_transposed - ? minor - : edge_partition.major_from_major_offset_nocheck(major_offset); - auto dst = GraphViewType::is_storage_transposed - ? edge_partition.major_from_major_offset_nocheck(major_offset) - : minor; - auto src_offset = - GraphViewType::is_storage_transposed ? minor_offset : static_cast(major_offset); - auto dst_offset = - GraphViewType::is_storage_transposed ? static_cast(major_offset) : minor_offset; - auto e_op_result = 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)); - if constexpr (update_major) { - reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); - } else { - if constexpr (GraphViewType::is_multi_gpu) { - reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + if (edge_partition_e_mask) { + for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) { + if ((*edge_partition_e_mask).get(edge_offset + i)) { + auto e_op_result = call_e_op(i); + if constexpr (update_major) { + reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); + } else { + auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(indices[i]); + if constexpr (GraphViewType::is_multi_gpu) { + reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + } else { + reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + } + } + } + } + } else { + for (edge_t i = lane_id; i < local_degree; i += raft::warp_size()) { + auto e_op_result = call_e_op(i); + if constexpr (update_major) { + reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); } else { - reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(indices[i]); + if constexpr (GraphViewType::is_multi_gpu) { + reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + } else { + reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + } } } } + if constexpr (update_major) { reduced_e_op_result = WarpReduce(temp_storage[threadIdx.x / raft::warp_size()]) .Reduce(reduced_e_op_result, reduce_op); @@ -377,6 +455,7 @@ template edge_partition_e_mask, ResultValueOutputIteratorOrWrapper result_value_output, EdgeOp e_op, T init /* relevant only if update_major == true */, @@ -416,41 +496,61 @@ __global__ void per_v_transform_reduce_e_high_degree( typename BlockReduce::TempStorage temp_storage; // relevant only if update_major == true while (idx < static_cast(major_range_last - major_range_first)) { - auto major_offset = major_start_offset + idx; + auto major_offset = static_cast(major_start_offset + idx); + auto major = edge_partition.major_from_major_offset_nocheck(major_offset); vertex_t const* indices{nullptr}; edge_t edge_offset{}; edge_t local_degree{}; thrust::tie(indices, edge_offset, local_degree) = edge_partition.local_edges(major_offset); + + auto call_e_op = call_e_op_t{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 - for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) { - auto minor = indices[i]; - auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(minor); - auto src = GraphViewType::is_storage_transposed - ? minor - : edge_partition.major_from_major_offset_nocheck(major_offset); - auto dst = GraphViewType::is_storage_transposed - ? edge_partition.major_from_major_offset_nocheck(major_offset) - : minor; - auto src_offset = - GraphViewType::is_storage_transposed ? minor_offset : static_cast(major_offset); - auto dst_offset = - GraphViewType::is_storage_transposed ? static_cast(major_offset) : minor_offset; - auto e_op_result = 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)); - if constexpr (update_major) { - reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); - } else { - if constexpr (GraphViewType::is_multi_gpu) { - reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + if (edge_partition_e_mask) { + for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) { + if ((*edge_partition_e_mask).get(edge_offset + i)) { + auto e_op_result = call_e_op(i); + if constexpr (update_major) { + reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); + } else { + auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(indices[i]); + if constexpr (GraphViewType::is_multi_gpu) { + reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + } else { + reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + } + } + } + } + } else { + for (edge_t i = threadIdx.x; i < local_degree; i += blockDim.x) { + auto e_op_result = call_e_op(i); + if constexpr (update_major) { + reduced_e_op_result = reduce_op(reduced_e_op_result, e_op_result); } else { - reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + auto minor_offset = edge_partition.minor_offset_from_minor_nocheck(indices[i]); + if constexpr (GraphViewType::is_multi_gpu) { + reduce_op::atomic_reduce(result_value_output, minor_offset, e_op_result); + } else { + reduce_op::atomic_reduce(result_value_output + minor_offset, e_op_result); + } } } } + if constexpr (update_major) { reduced_e_op_result = BlockReduce(temp_storage).Reduce(reduced_e_op_result, reduce_op); if (threadIdx.x == 0) { *(result_value_output + idx) = reduced_e_op_result; } @@ -656,10 +756,18 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, if (stream_pool_indices) { handle.sync_stream(); } + auto edge_mask_view = graph_view.edge_mask_view(); + for (size_t i = 0; i < graph_view.number_of_local_edge_partitions(); ++i) { auto edge_partition = edge_partition_device_view_t( graph_view.local_edge_partition_view(i)); + auto edge_partition_e_mask = + edge_mask_view + ? thrust::make_optional< + detail::edge_partition_edge_property_device_view_t>( + *edge_mask_view, i) + : thrust::nullopt; auto major_init = ReduceOp::identity_element; if constexpr (update_major) { @@ -737,9 +845,11 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, edge_partition_src_value_input, edge_partition_dst_value_input, edge_partition_e_value_input, + edge_partition_e_mask, segment_output_buffer, e_op, major_init, + ReduceOp::identity_element, reduce_op); } } @@ -761,9 +871,11 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, edge_partition_src_value_input, edge_partition_dst_value_input, edge_partition_e_value_input, + edge_partition_e_mask, segment_output_buffer, e_op, major_init, + ReduceOp::identity_element, reduce_op); } if ((*segment_offsets)[2] - (*segment_offsets)[1] > 0) { @@ -784,6 +896,7 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, edge_partition_src_value_input, edge_partition_dst_value_input, edge_partition_e_value_input, + edge_partition_e_mask, segment_output_buffer, e_op, major_init, @@ -806,6 +919,7 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, edge_partition_src_value_input, edge_partition_dst_value_input, edge_partition_e_value_input, + edge_partition_e_mask, output_buffer, e_op, major_init, @@ -825,9 +939,11 @@ void per_v_transform_reduce_e(raft::handle_t const& handle, edge_partition_src_value_input, edge_partition_dst_value_input, edge_partition_e_value_input, + edge_partition_e_mask, output_buffer, e_op, major_init, + ReduceOp::identity_element, reduce_op); } } @@ -1056,8 +1172,6 @@ void per_v_transform_reduce_incoming_e(raft::handle_t const& handle, VertexValueOutputIterator vertex_value_output_first, bool do_expensive_check = false) { - CUGRAPH_EXPECTS(!graph_view.has_edge_mask(), "unimplemented."); - if (do_expensive_check) { // currently, nothing to do } @@ -1137,8 +1251,6 @@ void per_v_transform_reduce_outgoing_e(raft::handle_t const& handle, VertexValueOutputIterator vertex_value_output_first, bool do_expensive_check = false) { - CUGRAPH_EXPECTS(!graph_view.has_edge_mask(), "unimplemented."); - if (do_expensive_check) { // currently, nothing to do } diff --git a/cpp/src/prims/transform_reduce_v_frontier_outgoing_e_by_dst.cuh b/cpp/src/prims/transform_reduce_v_frontier_outgoing_e_by_dst.cuh index 745f1a8fd8e..18e722d62cc 100644 --- a/cpp/src/prims/transform_reduce_v_frontier_outgoing_e_by_dst.cuh +++ b/cpp/src/prims/transform_reduce_v_frontier_outgoing_e_by_dst.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-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. @@ -80,7 +80,7 @@ template -struct call_e_op_t { +struct transform_reduce_v_frontier_call_e_op_t { EdgeOp e_op{}; __device__ thrust::optional< @@ -331,13 +331,13 @@ transform_reduce_v_frontier_outgoing_e_by_dst(raft::handle_t const& handle, // 1. fill the buffer - detail::call_e_op_t + detail::transform_reduce_v_frontier_call_e_op_t e_op_wrapper{e_op}; auto [key_buffer, payload_buffer] = diff --git a/cpp/src/traversal/od_shortest_distances_impl.cuh b/cpp/src/traversal/od_shortest_distances_impl.cuh index cc69cb5f67f..58fae83bca0 100644 --- a/cpp/src/traversal/od_shortest_distances_impl.cuh +++ b/cpp/src/traversal/od_shortest_distances_impl.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023, NVIDIA CORPORATION. + * 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. @@ -639,13 +639,14 @@ rmm::device_uvector od_shortest_distances( static_cast(origins.size()), cutoff, invalid_distance}; - detail::call_e_op_t, - weight_t, - vertex_t, - thrust::nullopt_t, - thrust::nullopt_t, - weight_t, - e_op_t> + detail::transform_reduce_v_frontier_call_e_op_t< + thrust::tuple, + weight_t, + vertex_t, + thrust::nullopt_t, + thrust::nullopt_t, + weight_t, + e_op_t> e_op_wrapper{e_op}; auto new_frontier_tagged_vertex_buffer = diff --git a/cpp/tests/prims/mg_per_v_transform_reduce_incoming_outgoing_e.cu b/cpp/tests/prims/mg_per_v_transform_reduce_incoming_outgoing_e.cu index 677d6ce5022..fc8114a4652 100644 --- a/cpp/tests/prims/mg_per_v_transform_reduce_incoming_outgoing_e.cu +++ b/cpp/tests/prims/mg_per_v_transform_reduce_incoming_outgoing_e.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. + * Copyright (c) 2021-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. @@ -150,8 +150,9 @@ struct result_compare { }; struct Prims_Usecase { - bool check_correctness{true}; bool test_weighted{false}; + bool edge_masking{false}; + bool check_correctness{true}; }; template @@ -200,6 +201,13 @@ class Tests_MGPerVTransformReduceIncomingOutgoingE auto mg_graph_view = mg_graph.view(); + std::optional> edge_mask{std::nullopt}; + if (prims_usecase.edge_masking) { + edge_mask = + cugraph::test::generate::edge_property(*handle_, mg_graph_view, 2); + mg_graph_view.attach_edge_mask((*edge_mask).view()); + } + // 2. run MG transform reduce const int hash_bin_count = 5; @@ -674,7 +682,10 @@ INSTANTIATE_TEST_SUITE_P( file_test, Tests_MGPerVTransformReduceIncomingOutgoingE_File, ::testing::Combine( - ::testing::Values(Prims_Usecase{true}), + ::testing::Values(Prims_Usecase{false, false, true}, + Prims_Usecase{false, true, true}, + Prims_Usecase{true, false, true}, + Prims_Usecase{true, true, true}), ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"), cugraph::test::File_Usecase("test/datasets/web-Google.mtx"), cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"), @@ -682,7 +693,10 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P(rmat_small_test, Tests_MGPerVTransformReduceIncomingOutgoingE_Rmat, - ::testing::Combine(::testing::Values(Prims_Usecase{true}), + ::testing::Combine(::testing::Values(Prims_Usecase{false, false, true}, + Prims_Usecase{false, true, true}, + Prims_Usecase{true, false, true}, + Prims_Usecase{true, true, true}), ::testing::Values(cugraph::test::Rmat_Usecase( 10, 16, 0.57, 0.19, 0.19, 0, false, false)))); @@ -694,7 +708,10 @@ INSTANTIATE_TEST_SUITE_P( factor (to avoid running same benchmarks more than once) */ Tests_MGPerVTransformReduceIncomingOutgoingE_Rmat, ::testing::Combine( - ::testing::Values(Prims_Usecase{false}), + ::testing::Values(Prims_Usecase{false, false, false}, + Prims_Usecase{false, true, false}, + Prims_Usecase{true, false, false}, + Prims_Usecase{true, true, false}), ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false)))); CUGRAPH_MG_TEST_PROGRAM_MAIN() From 3ff2abd44f90cd143f1fe09acdd12c1ce848656f Mon Sep 17 00:00:00 2001 From: Tingyu Wang Date: Mon, 29 Jan 2024 10:16:52 -0500 Subject: [PATCH 5/5] Create `cugraph-equivariant` package (#4036) Bring up `cugraph-equivariant` package and add TensorProduct conv layers. Authors: - Tingyu Wang (https://github.com/tingyu66) Approvers: - Jake Awe (https://github.com/AyodeAwe) - https://github.com/DejunL - Maximilian Stadler (https://github.com/stadlmax) - Mario Geiger (https://github.com/mariogeiger) - Rick Ratzel (https://github.com/rlratzel) URL: https://github.com/rapidsai/cugraph/pull/4036 --- .github/workflows/build.yaml | 19 ++ .github/workflows/pr.yaml | 16 ++ .github/workflows/test.yaml | 9 + build.sh | 15 +- ci/build_python.sh | 5 + ci/build_wheel.sh | 3 +- ci/build_wheel_cugraph-equivariant.sh | 6 + ci/test_python.sh | 41 +++ ci/test_wheel_cugraph-equivariant.sh | 33 +++ conda/recipes/cugraph-equivariant/build.sh | 7 + conda/recipes/cugraph-equivariant/meta.yaml | 37 +++ dependencies.yaml | 22 ++ python/cugraph-equivariant/LICENSE | 1 + python/cugraph-equivariant/README.md | 5 + .../cugraph_equivariant/VERSION | 1 + .../cugraph_equivariant/__init__.py | 14 + .../cugraph_equivariant/_version.py | 27 ++ .../cugraph_equivariant/nn/__init__.py | 21 ++ .../nn/tensor_product_conv.py | 259 ++++++++++++++++++ .../cugraph_equivariant/tests/conftest.py | 31 +++ .../cugraph_equivariant/tests/test_scatter.py | 28 ++ .../tests/test_tensor_product_conv.py | 115 ++++++++ .../cugraph_equivariant/utils/__init__.py | 18 ++ .../cugraph_equivariant/utils/scatter.py | 43 +++ python/cugraph-equivariant/pyproject.toml | 64 +++++ python/cugraph-equivariant/setup.py | 20 ++ 26 files changed, 857 insertions(+), 3 deletions(-) create mode 100755 ci/build_wheel_cugraph-equivariant.sh create mode 100755 ci/test_wheel_cugraph-equivariant.sh create mode 100644 conda/recipes/cugraph-equivariant/build.sh create mode 100644 conda/recipes/cugraph-equivariant/meta.yaml create mode 120000 python/cugraph-equivariant/LICENSE create mode 100644 python/cugraph-equivariant/README.md create mode 120000 python/cugraph-equivariant/cugraph_equivariant/VERSION create mode 100644 python/cugraph-equivariant/cugraph_equivariant/__init__.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/_version.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/tests/conftest.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/tests/test_scatter.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/utils/__init__.py create mode 100644 python/cugraph-equivariant/cugraph_equivariant/utils/scatter.py create mode 100644 python/cugraph-equivariant/pyproject.toml create mode 100644 python/cugraph-equivariant/setup.py diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 273a8902eae..243c5f23ec0 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -173,3 +173,22 @@ jobs: sha: ${{ inputs.sha }} date: ${{ inputs.date }} package-name: cugraph-pyg + wheel-build-cugraph-equivariant: + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 + with: + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + sha: ${{ inputs.sha }} + date: ${{ inputs.date }} + script: ci/build_wheel_cugraph-equivariant.sh + wheel-publish-cugraph-equivariant: + needs: wheel-build-cugraph-equivariant + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-24.02 + with: + build_type: ${{ inputs.build_type || 'branch' }} + branch: ${{ inputs.branch }} + sha: ${{ inputs.sha }} + date: ${{ inputs.date }} + package-name: cugraph-equivariant diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 84d22f8e896..1bb2e0ab0a7 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -29,6 +29,8 @@ jobs: - wheel-tests-cugraph-dgl - wheel-build-cugraph-pyg - wheel-tests-cugraph-pyg + - wheel-build-cugraph-equivariant + - wheel-tests-cugraph-equivariant - devcontainer secrets: inherit uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.02 @@ -161,6 +163,20 @@ jobs: build_type: pull-request script: ci/test_wheel_cugraph-pyg.sh matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "11.8.0")) + wheel-build-cugraph-equivariant: + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 + with: + build_type: pull-request + script: ci/build_wheel_cugraph-equivariant.sh + wheel-tests-cugraph-equivariant: + needs: wheel-build-cugraph-equivariant + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 + with: + build_type: pull-request + script: ci/test_wheel_cugraph-equivariant.sh + matrix_filter: map(select(.ARCH == "amd64")) devcontainer: secrets: inherit uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-24.02 diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 773358ede8d..71051bcc529 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -75,3 +75,12 @@ jobs: date: ${{ inputs.date }} sha: ${{ inputs.sha }} script: ci/test_wheel_cugraph-pyg.sh + wheel-tests-cugraph-equivariant: + secrets: inherit + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 + with: + build_type: nightly + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + script: ci/test_wheel_cugraph-equivariant.sh diff --git a/build.sh b/build.sh index 5044b3a55b3..82de45ca9fb 100755 --- a/build.sh +++ b/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Copyright (c) 2019-2023, NVIDIA CORPORATION. +# Copyright (c) 2019-2024, NVIDIA CORPORATION. # cugraph build script @@ -31,6 +31,7 @@ VALIDARGS=" cugraph-service cugraph-pyg cugraph-dgl + cugraph-equivariant nx-cugraph cpp-mgtests cpp-mtmgtests @@ -60,6 +61,7 @@ HELP="$0 [ ...] [ ...] cugraph-service - build the cugraph-service_client and cugraph-service_server Python package cugraph-pyg - build the cugraph-pyg Python package cugraph-dgl - build the cugraph-dgl extensions for DGL + cugraph-equivariant - build the cugraph-equivariant Python package nx-cugraph - build the nx-cugraph Python package cpp-mgtests - build libcugraph and libcugraph_etl MG tests. Builds MPI communicator, adding MPI as a dependency. cpp-mtmgtests - build libcugraph MTMG tests. Adds UCX as a dependency (temporary). @@ -222,7 +224,7 @@ if hasArg uninstall; then # removes the latest one and leaves the others installed. build.sh uninstall # can be run multiple times to remove all of them, but that is not obvious. pip uninstall -y pylibcugraph cugraph cugraph-service-client cugraph-service-server \ - cugraph-dgl cugraph-pyg nx-cugraph + cugraph-dgl cugraph-pyg cugraph-equivariant nx-cugraph fi if hasArg clean; then @@ -359,6 +361,15 @@ if hasArg cugraph-dgl || hasArg all; then fi fi +# Build and install the cugraph-equivariant Python package +if hasArg cugraph-equivariant || hasArg all; then + if hasArg --clean; then + cleanPythonDir ${REPODIR}/python/cugraph-equivariant + else + python ${PYTHON_ARGS_FOR_INSTALL} ${REPODIR}/python/cugraph-equivariant + fi +fi + # Build and install the nx-cugraph Python package if hasArg nx-cugraph || hasArg all; then if hasArg --clean; then diff --git a/ci/build_python.sh b/ci/build_python.sh index a99e5ce63e8..07a4f59396b 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -89,4 +89,9 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then conda/recipes/cugraph-dgl fi +rapids-conda-retry mambabuild \ + --no-test \ + --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \ + conda/recipes/cugraph-equivariant + rapids-upload-conda-to-s3 python diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 828d8948143..30a1c98c106 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -57,7 +57,8 @@ python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check # pure-python packages should not have auditwheel run on them. if [[ ${package_name} == "nx-cugraph" ]] || \ [[ ${package_name} == "cugraph-dgl" ]] || \ - [[ ${package_name} == "cugraph-pyg" ]]; then + [[ ${package_name} == "cugraph-pyg" ]] || \ + [[ ${package_name} == "cugraph-equivariant" ]]; then RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 dist else mkdir -p final_dist diff --git a/ci/build_wheel_cugraph-equivariant.sh b/ci/build_wheel_cugraph-equivariant.sh new file mode 100755 index 00000000000..fcc8e0f774c --- /dev/null +++ b/ci/build_wheel_cugraph-equivariant.sh @@ -0,0 +1,6 @@ +#!/bin/bash +# Copyright (c) 2024, NVIDIA CORPORATION. + +set -euo pipefail + +./ci/build_wheel.sh cugraph-equivariant python/cugraph-equivariant diff --git a/ci/test_python.sh b/ci/test_python.sh index 7eb5a08edc8..5892c37e35b 100755 --- a/ci/test_python.sh +++ b/ci/test_python.sh @@ -247,5 +247,46 @@ else rapids-logger "skipping cugraph_pyg pytest on CUDA != 11.8" fi +# test cugraph-equivariant +if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then + if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then + # Reuse cugraph-dgl's test env for cugraph-equivariant + set +u + conda activate test_cugraph_dgl + set -u + rapids-mamba-retry install \ + --channel "${CPP_CHANNEL}" \ + --channel "${PYTHON_CHANNEL}" \ + --channel pytorch \ + --channel nvidia \ + cugraph-equivariant + pip install e3nn==0.5.1 + + rapids-print-env + + rapids-logger "pytest cugraph-equivariant" + pushd python/cugraph-equivariant/cugraph_equivariant + pytest \ + --cache-clear \ + --junitxml="${RAPIDS_TESTS_DIR}/junit-cugraph-equivariant.xml" \ + --cov-config=../../.coveragerc \ + --cov=cugraph_equivariant \ + --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/cugraph-equivariant-coverage.xml" \ + --cov-report=term \ + . + popd + + # Reactivate the test environment back + set +u + conda deactivate + conda activate test + set -u + else + rapids-logger "skipping cugraph-equivariant pytest on ARM64" + fi +else + rapids-logger "skipping cugraph-equivariant pytest on CUDA!=11.8" +fi + rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/ci/test_wheel_cugraph-equivariant.sh b/ci/test_wheel_cugraph-equivariant.sh new file mode 100755 index 00000000000..f054780b03a --- /dev/null +++ b/ci/test_wheel_cugraph-equivariant.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Copyright (c) 2024, NVIDIA CORPORATION. + +set -eoxu pipefail + +package_name="cugraph-equivariant" +package_dir="python/cugraph-equivariant" + +python_package_name=$(echo ${package_name}|sed 's/-/_/g') + +mkdir -p ./dist +RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})" + +# use 'ls' to expand wildcard before adding `[extra]` requires for pip +RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist +# pip creates wheels using python package names +python -m pip install $(ls ./dist/${python_package_name}*.whl)[test] + + +PKG_CUDA_VER="$(echo ${CUDA_VERSION} | cut -d '.' -f1,2 | tr -d '.')" +PKG_CUDA_VER_MAJOR=${PKG_CUDA_VER:0:2} +if [[ "${PKG_CUDA_VER_MAJOR}" == "12" ]]; then + PYTORCH_CUDA_VER="121" +else + PYTORCH_CUDA_VER=$PKG_CUDA_VER +fi +PYTORCH_URL="https://download.pytorch.org/whl/cu${PYTORCH_CUDA_VER}" + +rapids-logger "Installing PyTorch and e3nn" +rapids-retry python -m pip install torch --index-url ${PYTORCH_URL} +rapids-retry python -m pip install e3nn + +python -m pytest python/cugraph-equivariant/cugraph_equivariant/tests diff --git a/conda/recipes/cugraph-equivariant/build.sh b/conda/recipes/cugraph-equivariant/build.sh new file mode 100644 index 00000000000..f0ff1688b55 --- /dev/null +++ b/conda/recipes/cugraph-equivariant/build.sh @@ -0,0 +1,7 @@ +#!/usr/bin/env bash + +# Copyright (c) 2024, NVIDIA CORPORATION. + +# This assumes the script is executed from the root of the repo directory + +./build.sh cugraph-equivariant diff --git a/conda/recipes/cugraph-equivariant/meta.yaml b/conda/recipes/cugraph-equivariant/meta.yaml new file mode 100644 index 00000000000..a952812f845 --- /dev/null +++ b/conda/recipes/cugraph-equivariant/meta.yaml @@ -0,0 +1,37 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. + +{% set version = environ['RAPIDS_PACKAGE_VERSION'].lstrip('v') + environ.get('VERSION_SUFFIX', '') %} +{% set minor_version = version.split('.')[0] + '.' + version.split('.')[1] %} +{% set py_version = environ['CONDA_PY'] %} +{% set date_string = environ['RAPIDS_DATE_STRING'] %} + +package: + name: cugraph-equivariant + version: {{ version }} + +source: + path: ../../.. + +build: + number: {{ GIT_DESCRIBE_NUMBER }} + build: + number: {{ GIT_DESCRIBE_NUMBER }} + string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }} + +requirements: + host: + - python + run: + - pylibcugraphops ={{ minor_version }} + - python + +tests: + imports: + - cugraph_equivariant + +about: + home: https://rapids.ai/ + dev_url: https://github.com/rapidsai/cugraph + license: Apache-2.0 + license_file: ../../../LICENSE + summary: GPU-accelerated equivariant convolutional layers. diff --git a/dependencies.yaml b/dependencies.yaml index 18ddb6c51dd..e9badf2be9f 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -198,6 +198,28 @@ files: key: test includes: - test_python_common + py_build_cugraph_equivariant: + output: pyproject + pyproject_dir: python/cugraph-equivariant + extras: + table: build-system + includes: + - python_build_wheel + py_run_cugraph_equivariant: + output: pyproject + pyproject_dir: python/cugraph-equivariant + extras: + table: project + includes: + - depends_on_pylibcugraphops + py_test_cugraph_equivariant: + output: pyproject + pyproject_dir: python/cugraph-equivariant + extras: + table: project.optional-dependencies + key: test + includes: + - test_python_common py_build_cugraph_service_client: output: pyproject pyproject_dir: python/cugraph-service/client diff --git a/python/cugraph-equivariant/LICENSE b/python/cugraph-equivariant/LICENSE new file mode 120000 index 00000000000..30cff7403da --- /dev/null +++ b/python/cugraph-equivariant/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/python/cugraph-equivariant/README.md b/python/cugraph-equivariant/README.md new file mode 100644 index 00000000000..d5de8852709 --- /dev/null +++ b/python/cugraph-equivariant/README.md @@ -0,0 +1,5 @@ +# cugraph-equivariant + +## Description + +cugraph-equivariant library provides fast symmetry-preserving (equivariant) operations and convolutional layers, to accelerate the equivariant neural networks in drug discovery and other domains. diff --git a/python/cugraph-equivariant/cugraph_equivariant/VERSION b/python/cugraph-equivariant/cugraph_equivariant/VERSION new file mode 120000 index 00000000000..d62dc733efd --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/VERSION @@ -0,0 +1 @@ +../../../VERSION \ No newline at end of file diff --git a/python/cugraph-equivariant/cugraph_equivariant/__init__.py b/python/cugraph-equivariant/cugraph_equivariant/__init__.py new file mode 100644 index 00000000000..20507bd9329 --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/__init__.py @@ -0,0 +1,14 @@ +# 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. + +from cugraph_equivariant._version import __git_commit__, __version__ diff --git a/python/cugraph-equivariant/cugraph_equivariant/_version.py b/python/cugraph-equivariant/cugraph_equivariant/_version.py new file mode 100644 index 00000000000..31a707bb17e --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/_version.py @@ -0,0 +1,27 @@ +# 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. + +import importlib.resources + +# Read VERSION file from the module that is symlinked to VERSION file +# in the root of the repo at build time or copied to the module at +# installation. VERSION is a separate file that allows CI build-time scripts +# to update version info (including commit hashes) without modifying +# source files. +__version__ = ( + importlib.resources.files("cugraph_equivariant") + .joinpath("VERSION") + .read_text() + .strip() +) +__git_commit__ = "" diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py b/python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py new file mode 100644 index 00000000000..8f4d8de0042 --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py @@ -0,0 +1,21 @@ +# 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. + +from .tensor_product_conv import FullyConnectedTensorProductConv + +DiffDockTensorProductConv = FullyConnectedTensorProductConv + +__all__ = [ + "FullyConnectedTensorProductConv", + "DiffDockTensorProductConv", +] diff --git a/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py new file mode 100644 index 00000000000..5120a23180d --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py @@ -0,0 +1,259 @@ +# 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. + +from typing import Optional, Sequence, Union + +import torch +from torch import nn +from e3nn import o3 +from e3nn.nn import BatchNorm + +from cugraph_equivariant.utils import scatter_reduce + +from pylibcugraphops.pytorch.operators import FusedFullyConnectedTensorProduct + + +class FullyConnectedTensorProductConv(nn.Module): + r"""Message passing layer for tensor products in DiffDock-like architectures. + The left operand of tensor product is the spherical harmonic representation + of edge vector; the right operand consists of node features in irreps. + + .. math:: + \sum_{b \in \mathcal{N}_a} Y\left(\hat{r}_{a b}\right) + \otimes_{\psi_{a b}} \mathbf{h}_b + + where the path weights :math:`\psi_{a b}` can be constructed from edge + embeddings and scalar features using an MLP: + + .. math:: + \psi_{a b} = \operatorname{MLP} + \left(e_{a b}, \mathbf{h}_a^0, \mathbf{h}_b^0\right) + + Users have the option to either directly input the weights or provide the + MLP parameters and scalar features from edges and nodes. + + Parameters + ---------- + in_irreps : e3nn.o3.Irreps + Irreps for the input node features. + + sh_irreps : e3nn.o3.Irreps + Irreps for the spherical harmonic representations of edge vectors. + + out_irreps : e3nn.o3.Irreps + Irreps for the output. + + batch_norm : bool, optional (default=True) + If true, batch normalization is applied. + + mlp_channels : sequence of ints, optional (default=None) + A sequence of integers defining number of neurons in each layer in MLP + before the output layer. If `None`, no MLP will be added. The input layer + contains edge embeddings and node scalar features. + + mlp_activation : nn.Module or sequence of nn.Module, optional (default=nn.GELU()) + A sequence of functions to be applied in between linear layers in MLP, + e.g., `nn.Sequential(nn.ReLU(), nn.Dropout(0.4))`. + + e3nn_compat_mode: bool, optional (default=False) + cugraph-ops and e3nn use different memory layout for Irreps-tensors. + The last (fastest moving) dimension is num_channels for cugraph-ops and + ir.dim for e3nn. When enabled, the input and output of this layer will + follow e3nn's memory layout. + + Examples + -------- + >>> # Case 1: MLP with the input layer having 6 channels and 2 hidden layers + >>> # having 16 channels. edge_emb.size(1) must match the size of + >>> # the input layer: 6 + >>> + >>> conv1 = FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps, + >>> mlp_channels=[6, 16, 16], mlp_activation=nn.ReLU()).cuda() + >>> out = conv1(src_features, edge_sh, edge_emb, graph) + >>> + >>> # Case 2: Same as case 1 but with the scalar features from edges, sources + >>> # and destinations passed in separately. + >>> + >>> conv2 = FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps, + >>> mlp_channels=[6, 16, 16], mlp_activation=nn.ReLU()).cuda() + >>> out = conv3(src_features, edge_sh, edge_scalars, graph, + >>> src_scalars=src_scalars, dst_scalars=dst_scalars) + >>> + >>> # Case 3: No MLP, edge_emb will be directly used as the tensor product weights + >>> + >>> conv3 = FullyConnectedTensorProductConv(in_irreps, sh_irreps, out_irreps, + >>> mlp_channels=None).cuda() + >>> out = conv2(src_features, edge_sh, edge_emb, graph) + + """ + + def __init__( + self, + in_irreps: o3.Irreps, + sh_irreps: o3.Irreps, + out_irreps: o3.Irreps, + batch_norm: bool = True, + mlp_channels: Optional[Sequence[int]] = None, + mlp_activation: Union[nn.Module, Sequence[nn.Module]] = nn.GELU(), + e3nn_compat_mode: bool = False, + ): + super().__init__() + self.in_irreps = in_irreps + self.out_irreps = out_irreps + self.sh_irreps = sh_irreps + self.e3nn_compat_mode = e3nn_compat_mode + + self.tp = FusedFullyConnectedTensorProduct( + in_irreps, sh_irreps, out_irreps, e3nn_compat_mode=e3nn_compat_mode + ) + + self.batch_norm = BatchNorm(out_irreps) if batch_norm else None + + if mlp_activation is None: + mlp_activation = [] + elif hasattr(mlp_activation, "__len__") and hasattr( + mlp_activation, "__getitem__" + ): + mlp_activation = list(mlp_activation) + else: + mlp_activation = [mlp_activation] + + if mlp_channels is not None: + dims = list(mlp_channels) + [self.tp.weight_numel] + mlp = [] + for i in range(len(dims) - 1): + mlp.append(nn.Linear(dims[i], dims[i + 1])) + if i != len(dims) - 2: + mlp.extend(mlp_activation) + self.mlp = nn.Sequential(*mlp) + else: + self.mlp = None + + def forward( + self, + src_features: torch.Tensor, + edge_sh: torch.Tensor, + edge_emb: torch.Tensor, + graph: tuple[torch.Tensor, tuple[int, int]], + src_scalars: Optional[torch.Tensor] = None, + dst_scalars: Optional[torch.Tensor] = None, + reduce: str = "mean", + edge_envelope: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass. + + Parameters + ---------- + src_features : torch.Tensor + Source node features. + Shape: (num_src_nodes, in_irreps.dim) + + edge_sh : torch.Tensor + The spherical harmonic representations of the edge vectors. + Shape: (num_edges, sh_irreps.dim) + + edge_emb: torch.Tensor + Edge embeddings that are fed into MLPs to generate tensor product weights. + Shape: (num_edges, dim), where `dim` should be: + - `tp.weight_numel` when the layer does not contain MLPs. + - num_edge_scalars, with the sum of num_[edge/src/dst]_scalars being + mlp_channels[0] + + graph : tuple + A tuple that stores the graph information, with the first element being + the adjacency matrix in COO, and the second element being its shape: + (num_src_nodes, num_dst_nodes). + + src_scalars: torch.Tensor, optional + Scalar features of source nodes. + Shape: (num_src_nodes, num_src_scalars) + + dst_scalars: torch.Tensor, optional + Scalar features of destination nodes. + Shape: (num_dst_nodes, num_dst_scalars) + + reduce : str, optional (default="mean") + Reduction operator. Choose between "mean" and "sum". + + edge_envelope: torch.Tensor, optional + Typically used as attenuation factors to fade out messages coming + from nodes close to the cutoff distance used to create the graph. + This is important to make the model smooth to the changes in node's + coordinates. + Shape: (num_edges,) + + Returns + ------- + torch.Tensor + Output node features. + Shape: (num_dst_nodes, out_irreps.dim) + """ + edge_emb_size = edge_emb.size(-1) + src_scalars_size = 0 if src_scalars is None else src_scalars.size(-1) + dst_scalars_size = 0 if dst_scalars is None else dst_scalars.size(-1) + + if self.mlp is None: + if self.tp.weight_numel != edge_emb_size: + raise RuntimeError( + f"When MLP is not present, edge_emb's last dimension must " + f"equal tp.weight_numel (but got {edge_emb_size} and " + f"{self.tp.weight_numel})" + ) + else: + total_size = edge_emb_size + src_scalars_size + dst_scalars_size + if self.mlp[0].in_features != total_size: + raise RuntimeError( + f"The size of MLP's input layer ({self.mlp[0].in_features}) " + f"does not match the total number of scalar features from " + f"edge_emb, src_scalars and dst_scalars ({total_size})" + ) + + if reduce not in ["mean", "sum"]: + raise RuntimeError( + f"reduce argument must be either 'mean' or 'sum', got {reduce}." + ) + + (src, dst), (num_src_nodes, num_dst_nodes) = graph + + if self.mlp is not None: + if src_scalars is None and dst_scalars is None: + tp_weights = self.mlp(edge_emb) + else: + w_edge, w_src, w_dst = torch.split( + self.mlp[0].weight, + (edge_emb_size, src_scalars_size, dst_scalars_size), + dim=-1, + ) + tp_weights = edge_emb @ w_edge.T + self.mlp[0].bias + + if src_scalars is not None: + tp_weights += (src_scalars @ w_src.T)[src] + + if dst_scalars is not None: + tp_weights += (dst_scalars @ w_dst.T)[dst] + + tp_weights = self.mlp[1:](tp_weights) + else: + tp_weights = edge_emb + + out = self.tp(src_features[src], edge_sh, tp_weights) + + if edge_envelope is not None: + out = out * edge_envelope.view(-1, 1) + + out = scatter_reduce(out, dst, dim=0, dim_size=num_dst_nodes, reduce=reduce) + + if self.batch_norm: + out = self.batch_norm(out) + + return out diff --git a/python/cugraph-equivariant/cugraph_equivariant/tests/conftest.py b/python/cugraph-equivariant/cugraph_equivariant/tests/conftest.py new file mode 100644 index 00000000000..c7c6bad07db --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/tests/conftest.py @@ -0,0 +1,31 @@ +# 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. + +import pytest +import torch + + +@pytest.fixture +def example_scatter_data(): + src_feat = torch.Tensor([3, 1, 0, 1, 1, 2]) + dst_indices = torch.Tensor([0, 1, 2, 2, 3, 1]) + + results = { + "sum": torch.Tensor([3.0, 3.0, 1.0, 1.0]), + "mean": torch.Tensor([3.0, 1.5, 0.5, 1.0]), + "prod": torch.Tensor([3.0, 2.0, 0.0, 1.0]), + "amax": torch.Tensor([3.0, 2.0, 1.0, 1.0]), + "amin": torch.Tensor([3.0, 1.0, 0.0, 1.0]), + } + + return src_feat, dst_indices, results diff --git a/python/cugraph-equivariant/cugraph_equivariant/tests/test_scatter.py b/python/cugraph-equivariant/cugraph_equivariant/tests/test_scatter.py new file mode 100644 index 00000000000..ff8048468ee --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/tests/test_scatter.py @@ -0,0 +1,28 @@ +# 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. + +import pytest +import torch +from cugraph_equivariant.utils import scatter_reduce + + +@pytest.mark.parametrize("reduce", ["sum", "mean", "prod", "amax", "amin"]) +def test_scatter_reduce(example_scatter_data, reduce): + device = torch.device("cuda:0") + src, index, out_true = example_scatter_data + src = src.to(device) + index = index.to(device) + + out = scatter_reduce(src, index, dim=0, dim_size=None, reduce=reduce) + + assert torch.allclose(out.cpu(), out_true[reduce]) diff --git a/python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py b/python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py new file mode 100644 index 00000000000..a2a13b32cd2 --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py @@ -0,0 +1,115 @@ +# 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. + +import pytest + +import torch +from torch import nn +from e3nn import o3 +from cugraph_equivariant.nn import FullyConnectedTensorProductConv + +device = torch.device("cuda:0") + + +@pytest.mark.parametrize("e3nn_compat_mode", [True, False]) +@pytest.mark.parametrize("batch_norm", [True, False]) +@pytest.mark.parametrize( + "mlp_channels, mlp_activation, scalar_sizes", + [ + [(30, 8, 8), nn.Sequential(nn.Dropout(0.3), nn.ReLU()), (15, 15, 0)], + [(7,), nn.GELU(), (2, 3, 2)], + [None, None, None], + ], +) +def test_tensor_product_conv_equivariance( + mlp_channels, mlp_activation, scalar_sizes, batch_norm, e3nn_compat_mode +): + torch.manual_seed(12345) + + in_irreps = o3.Irreps("10x0e + 10x1e") + out_irreps = o3.Irreps("20x0e + 10x1e") + sh_irreps = o3.Irreps.spherical_harmonics(lmax=2) + + tp_conv = FullyConnectedTensorProductConv( + in_irreps=in_irreps, + sh_irreps=sh_irreps, + out_irreps=out_irreps, + mlp_channels=mlp_channels, + mlp_activation=mlp_activation, + batch_norm=batch_norm, + e3nn_compat_mode=e3nn_compat_mode, + ).to(device) + + num_src_nodes, num_dst_nodes = 9, 7 + num_edges = 40 + src = torch.randint(num_src_nodes, (num_edges,), device=device) + dst = torch.randint(num_dst_nodes, (num_edges,), device=device) + edge_index = torch.vstack((src, dst)) + + src_pos = torch.randn(num_src_nodes, 3, device=device) + dst_pos = torch.randn(num_dst_nodes, 3, device=device) + edge_vec = dst_pos[dst] - src_pos[src] + edge_sh = o3.spherical_harmonics( + tp_conv.sh_irreps, edge_vec, normalize=True, normalization="component" + ).to(device) + src_features = torch.randn(num_src_nodes, in_irreps.dim, device=device) + + rot = o3.rand_matrix() + D_in = tp_conv.in_irreps.D_from_matrix(rot).to(device) + D_sh = tp_conv.sh_irreps.D_from_matrix(rot).to(device) + D_out = tp_conv.out_irreps.D_from_matrix(rot).to(device) + + if mlp_channels is None: + edge_emb = torch.randn(num_edges, tp_conv.tp.weight_numel, device=device) + src_scalars = dst_scalars = None + else: + if scalar_sizes: + edge_emb = torch.randn(num_edges, scalar_sizes[0], device=device) + src_scalars = ( + None + if scalar_sizes[1] == 0 + else torch.randn(num_src_nodes, scalar_sizes[1], device=device) + ) + dst_scalars = ( + None + if scalar_sizes[2] == 0 + else torch.randn(num_dst_nodes, scalar_sizes[2], device=device) + ) + else: + edge_emb = torch.randn(num_edges, tp_conv.mlp[0].in_features, device=device) + src_scalars = dst_scalars = None + + # rotate before + out_before = tp_conv( + src_features=src_features @ D_in.T, + edge_sh=edge_sh @ D_sh.T, + edge_emb=edge_emb, + graph=(edge_index, (num_src_nodes, num_dst_nodes)), + src_scalars=src_scalars, + dst_scalars=dst_scalars, + ) + + # rotate after + out_after = ( + tp_conv( + src_features=src_features, + edge_sh=edge_sh, + edge_emb=edge_emb, + graph=(edge_index, (num_src_nodes, num_dst_nodes)), + src_scalars=src_scalars, + dst_scalars=dst_scalars, + ) + @ D_out.T + ) + + torch.allclose(out_before, out_after, rtol=1e-4, atol=1e-4) diff --git a/python/cugraph-equivariant/cugraph_equivariant/utils/__init__.py b/python/cugraph-equivariant/cugraph_equivariant/utils/__init__.py new file mode 100644 index 00000000000..b4acfe8d090 --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/utils/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from .scatter import scatter_reduce + +__all__ = [ + "scatter_reduce", +] diff --git a/python/cugraph-equivariant/cugraph_equivariant/utils/scatter.py b/python/cugraph-equivariant/cugraph_equivariant/utils/scatter.py new file mode 100644 index 00000000000..45cc541fc7b --- /dev/null +++ b/python/cugraph-equivariant/cugraph_equivariant/utils/scatter.py @@ -0,0 +1,43 @@ +# 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. + +from typing import Optional + +import torch + + +def broadcast(src: torch.Tensor, ref: torch.Tensor, dim: int) -> torch.Tensor: + size = ((1,) * dim) + (-1,) + ((1,) * (ref.dim() - dim - 1)) + return src.view(size).expand_as(ref) + + +def scatter_reduce( + src: torch.Tensor, + index: torch.Tensor, + dim: int = 0, + dim_size: Optional[int] = None, # value of out.size(dim) + reduce: str = "sum", # "sum", "prod", "mean", "amax", "amin" +): + # scatter() expects index to be int64 + index = broadcast(index, src, dim).to(torch.int64) + + size = list(src.size()) + + if dim_size is not None: + assert dim_size >= int(index.max()) + 1 + size[dim] = dim_size + else: + size[dim] = int(index.max()) + 1 + + out = torch.zeros(size, dtype=src.dtype, device=src.device) + return out.scatter_reduce_(dim, index, src, reduce, include_self=False) diff --git a/python/cugraph-equivariant/pyproject.toml b/python/cugraph-equivariant/pyproject.toml new file mode 100644 index 00000000000..f261b0e3535 --- /dev/null +++ b/python/cugraph-equivariant/pyproject.toml @@ -0,0 +1,64 @@ +# Copyright (c) 2023, 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. + +[build-system] +requires = [ + "setuptools>=61.0.0", + "wheel", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[project] +name = "cugraph-equivariant" +dynamic = ["version"] +description = "Fast GPU-based equivariant operations and convolutional layers." +readme = { file = "README.md", content-type = "text/markdown" } +authors = [ + { name = "NVIDIA Corporation" }, +] +license = { text = "Apache 2.0" } +requires-python = ">=3.9" +classifiers = [ + "Intended Audience :: Developers", + "Programming Language :: Python", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", +] +dependencies = [ + "pylibcugraphops==24.2.*", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[project.urls] +Homepage = "https://github.com/rapidsai/cugraph" +Documentation = "https://docs.rapids.ai/api/cugraph/stable/api_docs/cugraph-ops/" + +[project.optional-dependencies] +test = [ + "pandas", + "pytest", + "pytest-benchmark", + "pytest-cov", + "pytest-xdist", + "scipy", +] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. + +[tool.setuptools] +license-files = ["LICENSE"] + +[tool.setuptools.dynamic] +version = {file = "cugraph_equivariant/VERSION"} + +[tool.setuptools.packages.find] +include = [ + "cugraph_equivariant*", + "cugraph_equivariant.*", +] diff --git a/python/cugraph-equivariant/setup.py b/python/cugraph-equivariant/setup.py new file mode 100644 index 00000000000..acd0df3f717 --- /dev/null +++ b/python/cugraph-equivariant/setup.py @@ -0,0 +1,20 @@ +# 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. + +from setuptools import find_packages, setup + +if __name__ == "__main__": + packages = find_packages(include=["cugraph_equivariant*"]) + setup( + package_data={key: ["VERSION"] for key in packages}, + )