Skip to content

Commit 8aad4eb

Browse files
authored
[Perf] Streams 1-4 (#410)
1 parent 1ac4cb8 commit 8aad4eb

40 files changed

Lines changed: 1740 additions & 100 deletions

docs/source/user_guide/index.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ tile16
5959
6060
fastcache
6161
graph
62+
streams
6263
perf_dispatch
6364
init_options
6465
```

docs/source/user_guide/streams.md

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
# Streams
2+
3+
Streams allow concurrent execution of GPU operations. By default, all Quadrants kernels launch on the default stream, which serializes everything. With streams, you can run multiple top-level for loops in parallel.
4+
5+
## Supported platforms
6+
7+
| Backend | Supported |
8+
|---------|-----------|
9+
| CUDA | Yes |
10+
| AMDGPU | Yes |
11+
| CPU | No-op |
12+
| Metal | No-op |
13+
| Vulkan | No-op |
14+
15+
On backends without native stream support, stream operations are no-ops and for loops run serially. Code using streams is portable across all backends — it will run without modifications, but serially.
16+
17+
## Stream parallelism
18+
19+
Inside a `@qd.kernel`, each `with qd.stream_parallel():` block runs on its own GPU stream.
20+
21+
```python
22+
import quadrants as qd
23+
24+
qd.init(arch=qd.cuda)
25+
26+
N = 1024
27+
a = qd.field(qd.f32, shape=(N,))
28+
b = qd.field(qd.f32, shape=(N,))
29+
c = qd.field(qd.f32, shape=(N,))
30+
31+
@qd.kernel
32+
def compute_ab():
33+
with qd.stream_parallel():
34+
for i in range(N):
35+
a[i] = compute_a(i)
36+
with qd.stream_parallel():
37+
for j in range(N):
38+
b[j] = compute_b(j)
39+
40+
@qd.kernel
41+
def combine():
42+
for i in range(N):
43+
c[i] = a[i] + b[i]
44+
45+
compute_ab() # the two stream_parallel blocks run concurrently
46+
combine() # runs after compute_ab() returns — a[] and b[] are ready
47+
```
48+
49+
Consecutive `with qd.stream_parallel():` blocks run concurrently. Multiple for loops within a single block share a stream and run serially on it. All streams are synchronized before the kernel returns.
50+
51+
### Restrictions
52+
53+
- All top-level statements in a kernel must be either all `stream_parallel` blocks or all regular statements. Mixing the two at the top level is a compile-time error.
54+
- Nesting `stream_parallel` blocks is not supported.
55+
56+
## Explicit streams
57+
58+
For cases that require manual control — such as launching separate kernels on different streams or interoperating with PyTorch — you can create and manage streams directly.
59+
60+
### Creating and using streams
61+
62+
Any `@qd.kernel` function accepts a special `qd_stream` keyword argument — you do not need to declare it in the kernel signature. The `@qd.kernel` decorator handles it automatically.
63+
64+
```python
65+
@qd.kernel
66+
def my_kernel():
67+
for i in range(N):
68+
a[i] = i
69+
70+
s1 = qd.create_stream()
71+
s2 = qd.create_stream()
72+
73+
my_kernel(qd_stream=s1)
74+
my_kernel(qd_stream=s2)
75+
76+
s1.synchronize()
77+
s2.synchronize()
78+
79+
s1.destroy()
80+
s2.destroy()
81+
```
82+
83+
Kernels on different streams may execute concurrently. Call `synchronize()` to block until all work on a stream completes.
84+
85+
### Events
86+
87+
Events let you express dependencies between streams without full synchronization.
88+
89+
```python
90+
s1 = qd.create_stream()
91+
s2 = qd.create_stream()
92+
93+
@qd.kernel
94+
def produce():
95+
for i in range(N):
96+
a[i] = 10.0
97+
98+
@qd.kernel
99+
def consume():
100+
for i in range(N):
101+
b[i] = a[i]
102+
103+
produce(qd_stream=s1)
104+
105+
e = qd.create_event()
106+
e.record(s1) # record when s1 finishes produce()
107+
e.wait(qd_stream=s2) # s2 waits for that event before proceeding
108+
109+
consume(qd_stream=s2) # safe to read a[] — produce() is guaranteed complete
110+
s2.synchronize()
111+
112+
e.destroy()
113+
s1.destroy()
114+
s2.destroy()
115+
```
116+
117+
`e.record(stream)` captures the point in `stream`'s execution. `e.wait(qd_stream=stream)` makes `stream` wait until the recorded point is reached. If `qd_stream` is omitted, the default stream waits.
118+
119+
### Context managers
120+
121+
Streams and events support `with` blocks for automatic cleanup:
122+
123+
```python
124+
with qd.create_stream() as s:
125+
some_func1(qd_stream=s)
126+
# s.destroy() called automatically — waits for in-flight work
127+
```
128+
129+
## Synchronization notes
130+
131+
- **`qd.sync()` only waits on the default stream.** It does not drain explicit streams. Call `stream.synchronize()` on each stream you need to wait for.
132+
- **No automatic synchronization with explicit streams.** When using explicit streams, you are responsible for inserting events or `synchronize()` calls when one stream's output is another stream's input. `stream_parallel` handles this automatically.
133+
134+
## Limitations
135+
136+
- **Not compatible with graphs.** Do not pass `qd_stream` to a kernel decorated with `graph=True` (if you do, a `RuntimeError` will be raised).
137+
- **Not compatible with autodiff.** Do not pass `qd_stream` to a kernel that uses reverse-mode or forward-mode differentiation, or inside a `qd.ad.Tape` context (if you do, a `RuntimeError` will be raised).

python/quadrants/lang/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from quadrants.lang.runtime_ops import *
1717
from quadrants.lang.snode import *
1818
from quadrants.lang.source_builder import *
19+
from quadrants.lang.stream import *
1920
from quadrants.lang.struct import *
2021
from quadrants.types.enums import DeviceCapability, Format, Layout # noqa: F401
2122

@@ -47,6 +48,7 @@
4748
"shell",
4849
"snode",
4950
"source_builder",
51+
"stream",
5052
"struct",
5153
"util",
5254
]

python/quadrants/lang/ast/ast_transformer.py

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,11 @@ def build_AnnAssign(ctx: ASTTransformerFuncContext, node: ast.AnnAssign):
119119

120120
@staticmethod
121121
def build_assign_annotated(
122-
ctx: ASTTransformerFuncContext, target: ast.Name, value, is_static_assign: bool, annotation: Type
122+
ctx: ASTTransformerFuncContext,
123+
target: ast.Name,
124+
value,
125+
is_static_assign: bool,
126+
annotation: Type,
123127
):
124128
"""Build an annotated assignment like this: target: annotation = value.
125129
@@ -165,7 +169,10 @@ def build_Assign(ctx: ASTTransformerFuncContext, node: ast.Assign) -> None:
165169

166170
@staticmethod
167171
def build_assign_unpack(
168-
ctx: ASTTransformerFuncContext, node_target: list | ast.Tuple, values, is_static_assign: bool
172+
ctx: ASTTransformerFuncContext,
173+
node_target: list | ast.Tuple,
174+
values,
175+
is_static_assign: bool,
169176
):
170177
"""Build the unpack assignments like this: (target1, target2) = (value1, value2).
171178
The function should be called only if the node target is a tuple.
@@ -591,7 +598,8 @@ def build_Return(ctx: ASTTransformerFuncContext, node: ast.Return) -> None:
591598
else:
592599
raise QuadrantsSyntaxError("The return type is not supported now!")
593600
ctx.ast_builder.create_kernel_exprgroup_return(
594-
expr.make_expr_group(return_exprs), _qd_core.DebugInfo(ctx.get_pos_info(node))
601+
expr.make_expr_group(return_exprs),
602+
_qd_core.DebugInfo(ctx.get_pos_info(node)),
595603
)
596604
else:
597605
ctx.return_data = node.value.ptr
@@ -1520,6 +1528,24 @@ def build_Continue(ctx: ASTTransformerFuncContext, node: ast.Continue) -> None:
15201528
ctx.ast_builder.insert_continue_stmt(_qd_core.DebugInfo(ctx.get_pos_info(node)))
15211529
return None
15221530

1531+
@staticmethod
1532+
def build_With(ctx: ASTTransformerFuncContext, node: ast.With) -> None:
1533+
if len(node.items) != 1:
1534+
raise QuadrantsSyntaxError("'with' in Quadrants kernels only supports a single context manager")
1535+
item = node.items[0]
1536+
if item.optional_vars is not None:
1537+
raise QuadrantsSyntaxError("'with ... as ...' is not supported in Quadrants kernels")
1538+
if not isinstance(item.context_expr, ast.Call):
1539+
raise QuadrantsSyntaxError("'with' in Quadrants kernels requires a call expression")
1540+
if not FunctionDefTransformer._is_stream_parallel_with(node, ctx.global_vars):
1541+
raise QuadrantsSyntaxError("'with' in Quadrants kernels only supports qd.stream_parallel()")
1542+
if not ctx.is_kernel:
1543+
raise QuadrantsSyntaxError("qd.stream_parallel() can only be used inside @qd.kernel, not @qd.func")
1544+
ctx.ast_builder.begin_stream_parallel()
1545+
build_stmts(ctx, node.body)
1546+
ctx.ast_builder.end_stream_parallel()
1547+
return None
1548+
15231549
@staticmethod
15241550
def build_Pass(ctx: ASTTransformerFuncContext, node: ast.Pass) -> None:
15251551
return None

python/quadrants/lang/ast/ast_transformers/function_def_transformer.py

Lines changed: 77 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,13 @@
2626
from quadrants.lang.ast.ast_transformer_utils import (
2727
ASTTransformerFuncContext,
2828
)
29+
from quadrants.lang.ast.symbol_resolver import ASTResolver
2930
from quadrants.lang.buffer_view import BufferView
3031
from quadrants.lang.exception import (
3132
QuadrantsSyntaxError,
3233
)
3334
from quadrants.lang.matrix import MatrixType
35+
from quadrants.lang.stream import stream_parallel
3436
from quadrants.lang.struct import StructType
3537
from quadrants.lang.util import to_quadrants_type
3638
from quadrants.types import annotations, buffer_view_type, ndarray_type, primitive_types
@@ -317,7 +319,11 @@ def _transform_func_arg(
317319
# polymorphic).
318320
if field.type is not _TensorClass and hasattr(field.type, "check_matched"):
319321
field.type.check_matched(data_child.get_type(), field.name)
320-
_cache = getattr(getattr(ctx, "global_context", None), "ndarray_to_any_array", None)
322+
_cache = getattr(
323+
getattr(ctx, "global_context", None),
324+
"ndarray_to_any_array",
325+
None,
326+
)
321327
promoted = _cache.get(id(data_child)) if _cache else None
322328
ctx.create_variable(flat_name, promoted if promoted is not None else data_child)
323329
elif dataclasses.is_dataclass(data_child):
@@ -336,7 +342,13 @@ def _transform_func_arg(
336342
# Ndarray arguments are passed by reference.
337343
if isinstance(argument_type, (ndarray_type.NdarrayType)):
338344
if not isinstance(
339-
data, (_ndarray.ScalarNdarray, matrix.VectorNdarray, matrix.MatrixNdarray, any_array.AnyArray)
345+
data,
346+
(
347+
_ndarray.ScalarNdarray,
348+
matrix.VectorNdarray,
349+
matrix.MatrixNdarray,
350+
any_array.AnyArray,
351+
),
340352
):
341353
raise QuadrantsSyntaxError(f"Argument {argument_name} of type {argument_type} is not recognized.")
342354
argument_type.check_matched(data.get_type(), argument_name)
@@ -443,7 +455,70 @@ def build_FunctionDef(
443455
else:
444456
FunctionDefTransformer._transform_as_func(ctx, node, args)
445457

458+
if ctx.is_kernel:
459+
FunctionDefTransformer._validate_stream_parallel_exclusivity(node.body, ctx.global_vars)
460+
446461
with ctx.variable_scope_guard():
447462
build_stmts(ctx, node.body)
448463

449464
return None
465+
466+
@staticmethod
467+
def _is_stream_parallel_with(stmt: ast.stmt, global_vars: dict[str, Any]) -> bool:
468+
if not isinstance(stmt, ast.With):
469+
return False
470+
if len(stmt.items) != 1:
471+
return False
472+
item = stmt.items[0]
473+
if not isinstance(item.context_expr, ast.Call):
474+
return False
475+
func_node = item.context_expr.func
476+
if ASTResolver.resolve_to(func_node, stream_parallel, global_vars):
477+
return True
478+
resolved = ASTResolver.resolve_value(func_node, global_vars)
479+
if resolved is not None:
480+
return getattr(resolved, "__name__", None) == "stream_parallel" and getattr(
481+
resolved, "__module__", ""
482+
).startswith("quadrants")
483+
if isinstance(func_node, ast.Attribute) and func_node.attr == "stream_parallel":
484+
return True
485+
if isinstance(func_node, ast.Name) and func_node.id == "stream_parallel":
486+
return True
487+
return False
488+
489+
@staticmethod
490+
def _is_docstring(stmt: ast.stmt, index: int) -> bool:
491+
return index == 0 and isinstance(stmt, ast.Expr) and isinstance(stmt.value, (ast.Constant, ast.Str))
492+
493+
@staticmethod
494+
def _is_coverage_probe(stmt: ast.stmt) -> bool:
495+
if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1:
496+
return False
497+
target = stmt.targets[0]
498+
return (
499+
isinstance(target, ast.Subscript)
500+
and isinstance(target.value, ast.Name)
501+
and target.value.id.startswith("_qd_cov")
502+
)
503+
504+
@staticmethod
505+
def _validate_stream_parallel_exclusivity(body: list[ast.stmt], global_vars: dict[str, Any]) -> None:
506+
if not any(FunctionDefTransformer._is_stream_parallel_with(s, global_vars) for s in body):
507+
return
508+
for i, stmt in enumerate(body):
509+
if FunctionDefTransformer._is_docstring(stmt, i):
510+
continue
511+
if FunctionDefTransformer._is_coverage_probe(stmt):
512+
continue
513+
if not FunctionDefTransformer._is_stream_parallel_with(stmt, global_vars):
514+
stmt_desc = f"{type(stmt).__name__}"
515+
if isinstance(stmt, ast.With) and stmt.items:
516+
ctx_expr = stmt.items[0].context_expr
517+
if isinstance(ctx_expr, ast.Call) and isinstance(ctx_expr.func, ast.Attribute):
518+
stmt_desc += f"(with {ast.dump(ctx_expr.func)})"
519+
raise QuadrantsSyntaxError(
520+
"When using qd.stream_parallel(), all top-level statements "
521+
"in the kernel must be 'with qd.stream_parallel():' blocks. "
522+
f"Move non-parallel code to a separate kernel. "
523+
f"[stmt {i}: {stmt_desc}, body_len={len(body)}]"
524+
)

python/quadrants/lang/ast/symbol_resolver.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,3 +55,35 @@ def resolve_to(node, wanted, scope):
5555
return False
5656
# The name ``scope`` here could be a bit confusing
5757
return scope is wanted
58+
59+
@staticmethod
60+
def resolve_value(node, scope):
61+
"""Resolve an AST Name/Attribute node to a Python object.
62+
63+
Same traversal as resolve_to but returns the resolved object (or None) instead of comparing against a wanted
64+
value.
65+
"""
66+
if isinstance(node, ast.Name):
67+
return scope.get(node.id) if isinstance(scope, dict) else None
68+
69+
if not isinstance(node, ast.Attribute):
70+
return None
71+
72+
v = node.value
73+
chain = [node.attr]
74+
while isinstance(v, ast.Attribute):
75+
chain.append(v.attr)
76+
v = v.value
77+
if not isinstance(v, ast.Name):
78+
return None
79+
chain.append(v.id)
80+
81+
for attr in reversed(chain):
82+
try:
83+
if isinstance(scope, dict):
84+
scope = scope[attr]
85+
else:
86+
scope = getattr(scope, attr)
87+
except (KeyError, AttributeError):
88+
return None
89+
return scope

0 commit comments

Comments
 (0)