프로그램/Python

cupy, numba

naudhizb 2023. 3. 28. 22:43
반응형
import cupy as cp
from cupy.cuda import StreamingPriorityQueue
from numba import cuda

# 그래프 정점의 개수
n = 6

# 인접 리스트로 그래프 생성
graph = [[] for _ in range(n)]
graph[0].append(1)
graph[0].append(2)
graph[1].append(3)
graph[2].append(3)
graph[2].append(4)
graph[3].append(5)
graph[4].append(5)

# DFS 함수
@cuda.jit(device=True)
def dfs_kernel(graph, visited, v, visited_matrix):
    visited[v] = True
    visited_matrix[v] = True
    for i in graph[v]:
        if not visited[i]:
            dfs_kernel(graph, visited, i, visited_matrix)

@cuda.jit
def dfs(graph, visited, start):
    visited_matrix = cuda.local.array((n, n), dtype=cp.bool_)
    visited_matrix[:, :] = False
    dfs_kernel(graph, visited, start, visited_matrix)
    visited[:] = visited_matrix[start, :]

# BFS 함수
@cuda.jit(device=True)
def bfs_kernel(graph, visited, q, visited_matrix):
    while not q.empty():
        v = q.get()
        visited_matrix[v] = True
        for i in graph[v]:
            if not visited[i]:
                q.put(i)
                visited[i] = True

@cuda.jit
def bfs(graph, visited, start):
    q = StreamingPriorityQueue()
    visited_matrix = cuda.local.array((n, n), dtype=cp.bool_)
    visited_matrix[:, :] = False
    q.put((0, start))
    visited[start] = True
    visited_matrix[start, :] = True
    bfs_kernel(graph, visited, q, visited_matrix)
    visited[:] = visited_matrix[start, :]

# 방문 여부를 나타내는 배열
visited_dfs = cp.zeros(n, dtype=cp.bool_)
visited_bfs = cp.zeros(n, dtype=cp.bool_)

# DFS 실행
threadsperblock = 32
blockspergrid = (n + (threadsperblock - 1)) // threadsperblock
dfs[blockspergrid, threadsperblock](cp.asarray(graph), visited_dfs, 0)
print("DFS visited matrix:\n", visited_dfs)

# BFS 실행
bfs[blockspergrid, threadsperblock](cp.asarray(graph), visited_bfs, 0)
print("BFS visited matrix:\n", visited_bfs)

 

numba에서 device=True 데코레이터는 CUDA 커널 함수 내부에서 호출될 수 있는 다른 CUDA 커널 함수를 정의하는 데 사용됩니다. 이 데코레이터를 함수에 추가하면 해당 함수는 CUDA 장치에서 실행 가능한 함수가 되며, 해당 함수는 다른 CUDA 커널 함수에서 호출될 수 있습니다.

즉, device=True 데코레이터는 CUDA 장치에서 실행 가능한 코드를 생성하는 데 사용됩니다. 따라서, numba에서 CUDA 커널 함수를 정의할 때에는 device=True 데코레이터를 사용하여 해당 함수가 CUDA 장치에서 실행 가능하도록 만들어야 합니다.

또한, device=True 데코레이터를 사용하면 해당 함수는 호스트 코드에서 직접 호출할 수 없습니다. 대신, CUDA 커널 함수 내부에서 호출할 수 있습니다.

반응형