프로그램/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 커널 함수 내부에서 호출할 수 있습니다.
반응형