Skip to content

Commit b14869a

Browse files
xsjkxsjk
committed
perf: reduce CUDA host synchronization overhead
1 parent 7a9e3e7 commit b14869a

6 files changed

Lines changed: 254 additions & 30 deletions

File tree

‎compyle/jit.py‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -370,11 +370,12 @@ def __call__(self, *args, **kw):
370370
c_func(*c_args, **kw)
371371
self.queue.finish()
372372
elif self.backend == 'cuda':
373-
import pycuda.driver as drv
374-
event = drv.Event()
375373
c_func(*c_args, **kw)
376-
event.record()
377-
event.synchronize()
374+
if get_config().profile:
375+
import pycuda.driver as drv
376+
event = drv.Event()
377+
event.record()
378+
event.synchronize()
378379

379380

380381
class ReductionJIT(parallel.ReductionBase):
@@ -449,11 +450,12 @@ def __call__(self, *args, **kw):
449450
self.queue.finish()
450451
return result.get()
451452
elif self.backend == 'cuda':
452-
import pycuda.driver as drv
453-
event = drv.Event()
454453
result = c_func(*c_args, **kw)
455-
event.record()
456-
event.synchronize()
454+
if get_config().profile:
455+
import pycuda.driver as drv
456+
event = drv.Event()
457+
event.record()
458+
event.synchronize()
457459
return result.get()
458460

459461

@@ -569,8 +571,9 @@ def __call__(self, **kwargs):
569571
c_func(*[c_args_dict[k] for k in output_arg_keys])
570572
self.queue.finish()
571573
elif self.backend == 'cuda':
572-
import pycuda.driver as drv
573-
event = drv.Event()
574574
c_func(*[c_args_dict[k] for k in output_arg_keys])
575-
event.record()
576-
event.synchronize()
575+
if get_config().profile:
576+
import pycuda.driver as drv
577+
event = drv.Event()
578+
event.record()
579+
event.synchronize()

‎compyle/low_level.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -260,15 +260,16 @@ def __call__(self, *args, **kw):
260260
self.knl(*c_args)
261261
self.queue.finish()
262262
elif self.backend == 'cuda':
263-
import pycuda.driver as drv
264263
shared_mem_size = int(self._get_local_size(args, ls[0]))
265264
num_blocks = int((n + ls[0] - 1) / ls[0])
266265
num_tpb = int(ls[0])
267-
event = drv.Event()
268266
self.knl(*c_args, block=(num_tpb, 1, 1), grid=(num_blocks, 1),
269267
shared=shared_mem_size)
270-
event.record()
271-
event.synchronize()
268+
if get_config().profile:
269+
import pycuda.driver as drv
270+
event = drv.Event()
271+
event.record()
272+
event.synchronize()
272273

273274

274275
class _prange(Extern):

‎compyle/parallel.py‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -547,11 +547,12 @@ def __call__(self, *args, **kw):
547547
self.c_func(*c_args, **kw)
548548
self.queue.finish()
549549
elif self.backend == 'cuda':
550-
import pycuda.driver as drv
551-
event = drv.Event()
552550
self.c_func(*c_args, **kw)
553-
event.record()
554-
event.synchronize()
551+
if get_config().profile:
552+
import pycuda.driver as drv
553+
event = drv.Event()
554+
event.record()
555+
event.synchronize()
555556

556557

557558
class Elementwise(object):
@@ -809,11 +810,12 @@ def __call__(self, *args):
809810
self.queue.finish()
810811
return result.get()
811812
elif self.backend == 'cuda':
812-
import pycuda.driver as drv
813-
event = drv.Event()
814813
result = self.c_func(*c_args)
815-
event.record()
816-
event.synchronize()
814+
if get_config().profile:
815+
import pycuda.driver as drv
816+
event = drv.Event()
817+
event.record()
818+
event.synchronize()
817819
return result.get()
818820

819821

@@ -1229,11 +1231,12 @@ def __call__(self, **kwargs):
12291231
self.c_func(*[c_args_dict[k] for k in output_arg_keys])
12301232
self.queue.finish()
12311233
elif self.backend == 'cuda':
1232-
import pycuda.driver as drv
1233-
event = drv.Event()
12341234
self.c_func(*[c_args_dict[k] for k in output_arg_keys])
1235-
event.record()
1236-
event.synchronize()
1235+
if get_config().profile:
1236+
import pycuda.driver as drv
1237+
event = drv.Event()
1238+
event.record()
1239+
event.synchronize()
12371240

12381241

12391242
class Scan(object):

‎compyle/tests/test_jit.py‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
from math import sin
22
import unittest
33
import numpy as np
4+
from unittest.mock import patch
45

56
from pytest import importorskip
67

78
from ..config import get_config, use_config
89
from ..array import wrap
9-
from ..jit import get_binop_return_type, AnnotationHelper
10+
from ..jit import (
11+
AnnotationHelper, ElementwiseJIT, ReductionJIT, ScanJIT,
12+
get_binop_return_type
13+
)
1014
from ..types import annotate
1115
from ..parallel import Elementwise, Reduction, Scan
1216

@@ -32,6 +36,91 @@ def undeclared_f(a, b):
3236
return g(h_ab)
3337

3438

39+
class TestCUDAJITSynchronization(unittest.TestCase):
40+
def _patch_cuda_event(self):
41+
sync_calls = []
42+
43+
class FakeEvent:
44+
def record(self):
45+
pass
46+
47+
def synchronize(self):
48+
sync_calls.append("sync")
49+
50+
return sync_calls, patch("pycuda.driver.Event", FakeEvent)
51+
52+
def test_cuda_elementwise_jit_does_not_synchronize_without_profile(self):
53+
importorskip("pycuda")
54+
55+
@annotate
56+
def axpb(i, x):
57+
x[i] = x[i] + 1.0
58+
59+
kernel = ElementwiseJIT(axpb, backend="cuda")
60+
kernel._generate_kernel = lambda *args: lambda *c_args, **kw: None
61+
sync_calls, event_patch = self._patch_cuda_event()
62+
63+
with use_config(profile=False), event_patch:
64+
kernel(np.zeros(8))
65+
66+
assert sync_calls == []
67+
68+
def test_cuda_scan_jit_does_not_synchronize_without_profile(self):
69+
importorskip("pycuda")
70+
71+
@annotate(input="doublep", return_="double")
72+
def input_expr(i, input):
73+
return input[i]
74+
75+
@annotate(output="doublep", item="double")
76+
def output_expr(i, item, output):
77+
output[i] = item
78+
79+
scan = ScanJIT(input=input_expr, output=output_expr, backend="cuda")
80+
output_expr.arg_keys = {scan._get_backend_key(): ["input", "output"]}
81+
scan._generate_kernel = lambda **kwargs: lambda *c_args: None
82+
sync_calls, event_patch = self._patch_cuda_event()
83+
84+
with use_config(profile=False), event_patch:
85+
scan(input=np.zeros(8), output=np.zeros(8))
86+
87+
assert sync_calls == []
88+
89+
def test_cuda_reduction_jit_does_not_event_synchronize_without_profile(self):
90+
importorskip("pycuda")
91+
92+
class FakeResult:
93+
def get(self):
94+
return 1.0
95+
96+
reduction = ReductionJIT("a+b", backend="cuda")
97+
reduction._generate_kernel = (
98+
lambda *args: lambda *c_args, **kw: FakeResult()
99+
)
100+
sync_calls, event_patch = self._patch_cuda_event()
101+
102+
with use_config(profile=False), event_patch:
103+
assert reduction(np.zeros(8)) == 1.0
104+
105+
assert sync_calls == []
106+
107+
def test_cuda_elementwise_jit_synchronizes_with_profile(self):
108+
importorskip("pycuda")
109+
110+
@annotate
111+
def axpb(i, x):
112+
x[i] = x[i] + 1.0
113+
114+
kernel = ElementwiseJIT(axpb, backend="cuda")
115+
kernel._generate_kernel = lambda *args: lambda *c_args, **kw: None
116+
sync_calls, event_patch = self._patch_cuda_event()
117+
118+
with use_config(profile=True), event_patch:
119+
kernel(np.zeros(8))
120+
121+
assert sync_calls == ["sync"]
122+
123+
35124
class TestAnnotationHelper(unittest.TestCase):
36125
def test_const_as_call_arg(self):
37126
# Given

‎compyle/tests/test_low_level.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import unittest
22
import numpy as np
3+
from unittest.mock import patch
34

45
from pytest import importorskip
56

@@ -13,6 +14,52 @@
1314

1415

1516
class TestKernel(unittest.TestCase):
17+
def _patch_cuda_event(self):
18+
sync_calls = []
19+
20+
class FakeEvent:
21+
def record(self):
22+
pass
23+
24+
def synchronize(self):
25+
sync_calls.append("sync")
26+
27+
return sync_calls, patch("pycuda.driver.Event", FakeEvent)
28+
29+
def _make_cuda_kernel(self):
30+
class FakeArray:
31+
data = np.zeros(8)
32+
33+
kernel = object.__new__(Kernel)
34+
kernel.backend = "cuda"
35+
kernel.knl = lambda *c_args, **kw: None
36+
kernel._get_workgroup_size = lambda n: ((8,), (128,))
37+
kernel._get_args = lambda args, workgroup_size: []
38+
kernel._get_local_size = lambda args, workgroup_size: 0
39+
return kernel, FakeArray()
40+
41+
def test_cuda_kernel_does_not_synchronize_without_profile(self):
42+
importorskip("pycuda")
43+
44+
kernel, fake_array = self._make_cuda_kernel()
45+
sync_calls, event_patch = self._patch_cuda_event()
46+
47+
with use_config(profile=False), event_patch:
48+
kernel(fake_array)
49+
50+
assert sync_calls == []
51+
52+
def test_cuda_kernel_synchronizes_with_profile(self):
53+
importorskip("pycuda")
54+
55+
kernel, fake_array = self._make_cuda_kernel()
56+
sync_calls, event_patch = self._patch_cuda_event()
57+
58+
with use_config(profile=True), event_patch:
59+
kernel(fake_array)
60+
61+
assert sync_calls == ["sync"]
62+
1663
def test_simple_kernel_opencl(self):
1764
importorskip('pyopencl')
1865

‎compyle/tests/test_parallel.py‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,100 @@
11
from math import sin
22
import unittest
33
import numpy as np
4+
from unittest.mock import patch
45

56
from pytest import importorskip
67

78
from ..config import get_config, use_config
89
from ..array import wrap, zeros
910
from ..types import annotate, declare
10-
from ..parallel import Elementwise, Reduction, Scan
11+
from ..parallel import (
12+
Elementwise, ElementwiseBase, Reduction, ReductionBase, Scan, ScanBase
13+
)
1114
from ..low_level import atomic_inc, atomic_dec
1215
from .test_jit import g
1316

1417
MY_CONST = 42
1518

1619

20+
class TestCUDAParallelSynchronization(unittest.TestCase):
21+
def _patch_cuda_event(self):
22+
sync_calls = []
23+
24+
class FakeEvent:
25+
def record(self):
26+
pass
27+
28+
def synchronize(self):
29+
sync_calls.append("sync")
30+
31+
return sync_calls, patch("pycuda.driver.Event", FakeEvent)
32+
33+
def test_cuda_elementwise_base_does_not_synchronize_without_profile(self):
34+
importorskip("pycuda")
35+
36+
kernel = object.__new__(ElementwiseBase)
37+
kernel.backend = "cuda"
38+
kernel.c_func = lambda *c_args, **kw: None
39+
sync_calls, event_patch = self._patch_cuda_event()
40+
41+
with use_config(profile=False), event_patch:
42+
kernel(np.zeros(8))
43+
44+
assert sync_calls == []
45+
46+
def test_cuda_reduction_base_does_not_event_synchronize_without_profile(self):
47+
importorskip("pycuda")
48+
49+
class FakeResult:
50+
def get(self):
51+
return 1.0
52+
53+
reduction = object.__new__(ReductionBase)
54+
reduction.backend = "cuda"
55+
reduction.c_func = lambda *c_args: FakeResult()
56+
sync_calls, event_patch = self._patch_cuda_event()
57+
58+
with use_config(profile=False), event_patch:
59+
assert reduction(np.zeros(8)) == 1.0
60+
61+
assert sync_calls == []
62+
63+
def test_cuda_scan_base_does_not_synchronize_without_profile(self):
64+
importorskip("pycuda")
65+
66+
class OutputFunc:
67+
pass
68+
69+
scan = object.__new__(ScanBase)
70+
scan.backend = "cuda"
71+
scan._config = get_config()
72+
scan.c_func = lambda *c_args: None
73+
scan.output_func = OutputFunc()
74+
scan.output_func.arg_keys = {
75+
scan._get_backend_key(): ["input", "output"]
76+
}
77+
sync_calls, event_patch = self._patch_cuda_event()
78+
79+
with use_config(profile=False), event_patch:
80+
scan(input=np.zeros(8), output=np.zeros(8))
81+
82+
assert sync_calls == []
83+
84+
def test_cuda_elementwise_base_synchronizes_with_profile(self):
85+
importorskip("pycuda")
86+
87+
kernel = object.__new__(ElementwiseBase)
88+
kernel.backend = "cuda"
89+
kernel.c_func = lambda *c_args, **kw: None
90+
sync_calls, event_patch = self._patch_cuda_event()
91+
92+
with use_config(profile=True), event_patch:
93+
kernel(np.zeros(8))
94+
95+
assert sync_calls == ["sync"]
96+
97+
1798
@annotate(x='int', return_='int')
1899
def external(x):
19100
return x

0 commit comments

Comments
 (0)