Skip to content

Commit 3802405

Browse files
oleksandr-pavlykoleksandr-pavlyk
committed
Consolidated example for strided_memory_view_cpu/gpu to single run functions
1 parent 5775dda commit 3802405

2 files changed

Lines changed: 174 additions & 98 deletions

File tree

cuda_core/examples/strided_memory_view.py renamed to cuda_core/examples/strided_memory_view_cpu.py

Lines changed: 45 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,8 @@
2626
except ImportError:
2727
print("cffi is not installed, the CPU example will be skipped", file=sys.stderr)
2828
FFI = None
29-
try:
30-
import cupy as cp
31-
except ImportError:
32-
print("cupy is not installed, the GPU example will be skipped", file=sys.stderr)
33-
cp = None
3429
import numpy as np
3530

36-
from cuda.core.experimental import Device, LaunchConfig, Program, ProgramOptions, launch
3731
from cuda.core.experimental.utils import StridedMemoryView, args_viewable_as_strided_memory
3832

3933
# ################################################################################
@@ -54,8 +48,43 @@
5448
func_name = "inplace_plus_arange_N"
5549
func_sig = f"void {func_name}(int* data, size_t N)"
5650

57-
# Here is a concrete (very naive!) implementation on CPU:
58-
if FFI:
51+
52+
# Now we are prepared to run the code from the user's perspective!
53+
#
54+
# ################################################################################
55+
56+
57+
# Below, as a user we want to perform the said in-place operation on a CPU
58+
# or GPU, by calling the corresponding function implemented "elsewhere"
59+
# (in the body of run function).
60+
61+
62+
# We assume the 0-th argument supports either DLPack or CUDA Array Interface (both
63+
# of which are supported by StridedMemoryView).
64+
@args_viewable_as_strided_memory((0,))
65+
def my_func(arr):
66+
global cpu_func
67+
global cpu_prog
68+
# Create a memory view over arr (assumed to be a 1D array of int32). The stream
69+
# ordering is taken care of, so that arr can be safely accessed on our work
70+
# stream (ordered after a data stream on which arr is potentially prepared).
71+
view = arr.view(-1)
72+
assert isinstance(view, StridedMemoryView)
73+
assert len(view.shape) == 1
74+
assert view.dtype == np.int32
75+
assert not view.is_device_accessible
76+
77+
size = view.shape[0]
78+
# DLPack also supports host arrays. We want to know if the array data is
79+
# accessible from the GPU, and dispatch to the right routine accordingly.
80+
cpu_func(cpu_prog.cast("int*", view.ptr), size)
81+
82+
83+
def run():
84+
global my_func
85+
if not FFI:
86+
return
87+
# Here is a concrete (very naive!) implementation on CPU:
5988
cpu_code = string.Template(r"""
6089
extern "C"
6190
$func_sig {
@@ -76,113 +105,31 @@
76105
extra_compile_args=["-std=c++11"],
77106
)
78107
temp_dir = tempfile.mkdtemp()
108+
saved_sys_path = sys.path.copy()
79109
try:
80110
cpu_prog.compile(tmpdir=temp_dir)
81-
finally:
82-
shutil.rmtree(temp_dir)
83-
saved_sys_path = sys.path
84-
try:
111+
85112
sys.path.append(temp_dir)
86113
cpu_func = getattr(importlib.import_module("_cpu_obj.lib"), func_name)
87-
finally:
88-
sys.path = saved_sys_path
89-
shutil.rmtree(temp_dir)
90-
91-
# Here is a concrete (again, very naive!) implementation on GPU:
92-
if cp:
93-
gpu_code = string.Template(r"""
94-
extern "C"
95-
__global__ $func_sig {
96-
const size_t tid = threadIdx.x + blockIdx.x * blockDim.x;
97-
const size_t stride_size = gridDim.x * blockDim.x;
98-
for (size_t i = tid; i < N; i += stride_size) {
99-
data[i] += i;
100-
}
101-
}
102-
""").substitute(func_sig=func_sig)
103-
104-
# To know the GPU's compute capability, we need to identify which GPU to use.
105-
dev = Device(0)
106-
dev.set_current()
107-
arch = "".join(f"{i}" for i in dev.compute_capability)
108-
gpu_prog = Program(gpu_code, code_type="c++", options=ProgramOptions(arch=f"sm_{arch}", std="c++11"))
109-
mod = gpu_prog.compile(target_type="cubin")
110-
gpu_ker = mod.get_kernel(func_name)
111-
112-
# Now we are prepared to run the code from the user's perspective!
113-
#
114-
# ################################################################################
115-
116114

117-
# Below, as a user we want to perform the said in-place operation on either CPU
118-
# or GPU, by calling the corresponding function implemented "elsewhere" (done above).
119-
120-
121-
# We assume the 0-th argument supports either DLPack or CUDA Array Interface (both
122-
# of which are supported by StridedMemoryView).
123-
@args_viewable_as_strided_memory((0,))
124-
def my_func(arr, work_stream):
125-
# Create a memory view over arr (assumed to be a 1D array of int32). The stream
126-
# ordering is taken care of, so that arr can be safely accessed on our work
127-
# stream (ordered after a data stream on which arr is potentially prepared).
128-
view = arr.view(work_stream.handle if work_stream else -1)
129-
assert isinstance(view, StridedMemoryView)
130-
assert len(view.shape) == 1
131-
assert view.dtype == np.int32
132-
133-
size = view.shape[0]
134-
# DLPack also supports host arrays. We want to know if the array data is
135-
# accessible from the GPU, and dispatch to the right routine accordingly.
136-
if view.is_device_accessible:
137-
block = 256
138-
grid = (size + block - 1) // block
139-
config = LaunchConfig(grid=grid, block=block)
140-
launch(work_stream, config, gpu_ker, view.ptr, np.uint64(size))
141-
# Here we're being conservative and synchronize over our work stream,
142-
# assuming we do not know the data stream; if we know then we could
143-
# just order the data stream after the work stream here, e.g.
144-
#
145-
# data_stream.wait(work_stream)
146-
#
147-
# without an expensive synchronization (with respect to the host).
148-
work_stream.sync()
149-
else:
150-
cpu_func(cpu_prog.cast("int*", view.ptr), size)
151-
152-
153-
# This takes the GPU path
154-
if cp:
155-
s = dev.create_stream()
156-
try:
157-
# Create input array on GPU
158-
arr_gpu = cp.ones(1024, dtype=cp.int32)
159-
print(f"before: {arr_gpu[:10]=}")
160-
161-
# Run the workload
162-
my_func(arr_gpu, s)
163-
164-
# Check the result
165-
print(f"after: {arr_gpu[:10]=}")
166-
assert cp.allclose(arr_gpu, 1 + cp.arange(1024, dtype=cp.int32))
167-
finally:
168-
s.close()
169-
170-
# This takes the CPU path
171-
if FFI:
172-
try:
173115
# Create input array on CPU
174116
arr_cpu = np.zeros(1024, dtype=np.int32)
175117
print(f"before: {arr_cpu[:10]=}")
176118

177119
# Run the workload
178-
my_func(arr_cpu, None)
120+
my_func(arr_cpu)
179121

180122
# Check the result
181123
print(f"after: {arr_cpu[:10]=}")
182124
assert np.allclose(arr_cpu, np.arange(1024, dtype=np.int32))
183125
finally:
126+
sys.path = saved_sys_path
184127
# to allow FFI module to unload, we delete references to
185128
# to cpu_func
186129
del cpu_func, my_func
187130
# clean up temp directory
188131
shutil.rmtree(temp_dir)
132+
133+
134+
if __name__ == "__main__":
135+
run()

‎cuda_core/examples/strided_memory_view_gpu.py‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
# ################################################################################
6+
#
7+
# This demo aims to illustrate two takeaways:
8+
#
9+
# 1. The similarity between CPU and GPU JIT-compilation with C++ sources
10+
# 2. How to use StridedMemoryView to interface with foreign C/C++ functions
11+
#
12+
# To facilitate this demo, we use cffi (https://cffi.readthedocs.io/) for the CPU
13+
# path, which can be easily installed from pip or conda following their instructions.
14+
# We also use NumPy/CuPy as the CPU/GPU array container.
15+
#
16+
# ################################################################################
17+
18+
import string
19+
import sys
20+
21+
try:
22+
import cupy as cp
23+
except ImportError:
24+
print("cupy is not installed, the GPU example will be skipped", file=sys.stderr)
25+
cp = None
26+
import numpy as np
27+
28+
from cuda.core.experimental import Device, LaunchConfig, Program, ProgramOptions, launch
29+
from cuda.core.experimental.utils import StridedMemoryView, args_viewable_as_strided_memory
30+
31+
# ################################################################################
32+
#
33+
# Usually this entire code block is in a separate file, built as a Python extension
34+
# module that can be imported by users at run time. For illustrative purposes we
35+
# use JIT compilation to make this demo self-contained.
36+
#
37+
# Here we assume an in-place operation, equivalent to the following NumPy code:
38+
#
39+
# >>> arr = ...
40+
# >>> assert arr.dtype == np.int32
41+
# >>> assert arr.ndim == 1
42+
# >>> arr += np.arange(arr.size, dtype=arr.dtype)
43+
#
44+
# is implemented for both CPU and GPU at low-level, with the following C function
45+
# signature:
46+
func_name = "inplace_plus_arange_N"
47+
func_sig = f"void {func_name}(int* data, size_t N)"
48+
49+
# Now we are prepared to run the code from the user's perspective!
50+
#
51+
# ################################################################################
52+
53+
54+
# Below, as a user we want to perform the said in-place operation on either CPU
55+
# or GPU, by calling the corresponding function implemented "elsewhere" (done above).
56+
57+
58+
# We assume the 0-th argument supports either DLPack or CUDA Array Interface (both
59+
# of which are supported by StridedMemoryView).
60+
@args_viewable_as_strided_memory((0,))
61+
def my_func(arr, work_stream, gpu_ker):
62+
# Create a memory view over arr (assumed to be a 1D array of int32). The stream
63+
# ordering is taken care of, so that arr can be safely accessed on our work
64+
# stream (ordered after a data stream on which arr is potentially prepared).
65+
view = arr.view(work_stream.handle if work_stream else -1)
66+
assert isinstance(view, StridedMemoryView)
67+
assert len(view.shape) == 1
68+
assert view.dtype == np.int32
69+
assert view.is_device_accessible
70+
71+
size = view.shape[0]
72+
# DLPack also supports host arrays. We want to know if the array data is
73+
# accessible from the GPU, and dispatch to the right routine accordingly.
74+
block = 256
75+
grid = (size + block - 1) // block
76+
config = LaunchConfig(grid=grid, block=block)
77+
launch(work_stream, config, gpu_ker, view.ptr, np.uint64(size))
78+
# Here we're being conservative and synchronize over our work stream,
79+
# assuming we do not know the data stream; if we know then we could
80+
# just order the data stream after the work stream here, e.g.
81+
#
82+
# data_stream.wait(work_stream)
83+
#
84+
# without an expensive synchronization (with respect to the host).
85+
work_stream.sync()
86+
87+
88+
def run():
89+
global my_func
90+
if not cp:
91+
return None
92+
# Here is a concrete (very naive!) implementation on GPU:
93+
gpu_code = string.Template(r"""
94+
extern "C"
95+
__global__ $func_sig {
96+
const size_t tid = threadIdx.x + blockIdx.x * blockDim.x;
97+
const size_t stride_size = gridDim.x * blockDim.x;
98+
for (size_t i = tid; i < N; i += stride_size) {
99+
data[i] += i;
100+
}
101+
}
102+
""").substitute(func_sig=func_sig)
103+
104+
# To know the GPU's compute capability, we need to identify which GPU to use.
105+
dev = Device(0)
106+
dev.set_current()
107+
arch = "".join(f"{i}" for i in dev.compute_capability)
108+
gpu_prog = Program(gpu_code, code_type="c++", options=ProgramOptions(arch=f"sm_{arch}", std="c++11"))
109+
mod = gpu_prog.compile(target_type="cubin")
110+
gpu_ker = mod.get_kernel(func_name)
111+
112+
s = dev.create_stream()
113+
try:
114+
# Create input array on GPU
115+
arr_gpu = cp.ones(1024, dtype=cp.int32)
116+
print(f"before: {arr_gpu[:10]=}")
117+
118+
# Run the workload
119+
my_func(arr_gpu, s, gpu_ker)
120+
121+
# Check the result
122+
print(f"after: {arr_gpu[:10]=}")
123+
assert cp.allclose(arr_gpu, 1 + cp.arange(1024, dtype=cp.int32))
124+
finally:
125+
s.close()
126+
127+
128+
if __name__ == "__main__":
129+
run()

0 commit comments

Comments
 (0)