Skip to content

Commit d790d42

Browse files
authored
Use CuPy for arrays & CUDA Python for low-level CUDA functions (#4958)
Switch to using CuPy for array creation when needed. Can still pass these to Numba or use `__cuda_array_interface__` for low-level operations. Also switch to CUDA Python for low-level CUDA functions calls in Python. Authors: - https://github.com/jakirkham Approvers: - Rick Ratzel (https://github.com/rlratzel) URL: #4958
1 parent 1bdc550 commit d790d42

File tree

2 files changed

+11
-9
lines changed

2 files changed

+11
-9
lines changed

python/cugraph/cugraph/layout/force_atlas2_wrapper.pyx

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2020-2023, NVIDIA CORPORATION.
1+
# Copyright (c) 2020-2025, NVIDIA CORPORATION.
22
# Licensed under the Apache License, Version 2.0 (the "License");
33
# you may not use this file except in compliance with the License.
44
# You may obtain a copy of the License at
@@ -22,7 +22,7 @@ from cugraph.structure.graph_primtypes cimport *
2222
from libcpp cimport bool
2323
from libc.stdint cimport uintptr_t
2424
import cudf
25-
from numba import cuda
25+
import cupy as cp
2626
import numpy as np
2727

2828
cdef extern from "cugraph/legacy/internals.hpp" namespace "cugraph::internals":
@@ -67,7 +67,7 @@ def force_atlas2(input_graph,
6767
cdef GraphCOOView[int,int,double] graph_double
6868

6969
df = cudf.DataFrame()
70-
df['vertex'] = cudf.Series(np.arange(num_verts, dtype=np.int32))
70+
df['vertex'] = cudf.Series(cp.arange(num_verts, dtype=cp.int32))
7171

7272
src = input_graph.edgelist.edgelist_df['src']
7373
dst = input_graph.edgelist.edgelist_df['dst']
@@ -103,12 +103,12 @@ def force_atlas2(input_graph,
103103

104104

105105
# We keep np.float32 as results for both cases
106-
pos = cuda.device_array(
106+
pos = cp.empty(
107107
(num_verts, 2),
108108
order="F",
109-
dtype=np.float32)
109+
dtype=cp.float32)
110110

111-
pos_ptr = pos.device_ctypes_pointer.value
111+
pos_ptr = pos.data.ptr
112112

113113
if input_graph.edgelist.weights \
114114
and input_graph.edgelist.edgelist_df['weights'].dtype == np.float64:

python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import numpy as np
1616
import os
1717

18-
import numba.cuda
18+
from cuda.bindings import runtime as cuda_runtime
1919

2020
from cugraph.gnn import FeatureStore
2121

@@ -28,8 +28,10 @@
2828

2929

3030
def get_cudart_version():
31-
major, minor = numba.cuda.runtime.get_version()
32-
return major * 1000 + minor * 10
31+
status, version = cuda_runtime.getLocalRuntimeVersion()
32+
if status != cuda_runtime.cudaError_t.cudaSuccess:
33+
raise RuntimeError(f"CUDA Error: {status}")
34+
return version
3335

3436

3537
pytestmark = [

0 commit comments

Comments
 (0)