From 8274d3433a141f54694d561ca2a1ff9b58efa63d Mon Sep 17 00:00:00 2001 From: T-Karu-smaecs Date: Wed, 16 Sep 2026 21:53:50 +0200 Subject: [PATCH] Add GPU workgroup barrier via __sync_threads and Barrier.arrive_and_wait. Lower both call forms to the same GLSL barrier and memoryBarrierShared through a CallPlugin, wire Call expr-stmts on the GPU path, and export a sync stub. --- .../compiler/translation/plugins/__init__.py | 2 + .../translation/plugins/sync_threads.py | 68 +++++++++++++++++++ .../gpu/compiler/translation/syntax/Flow.py | 9 ++- src/cthreads/python/cthreads/sync/__init__.py | 21 ++++++ tests/unit/test_gpu_syntax.py | 28 ++++++++ 5 files changed, 126 insertions(+), 2 deletions(-) create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/plugins/sync_threads.py diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py index f8b218e..093ea55 100644 --- a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py @@ -102,6 +102,8 @@ def lower_attr( # Side-effect: register concrete plugins. from .indexes import IndexAttrPlugin # noqa: E402 from .math_calls import MathCallPlugin # noqa: E402 +from .sync_threads import SyncThreadsPlugin # noqa: E402 register_attr(IndexAttrPlugin()) register_call(MathCallPlugin()) +register_call(SyncThreadsPlugin()) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/sync_threads.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/sync_threads.py new file mode 100644 index 0000000..8d33404 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/sync_threads.py @@ -0,0 +1,68 @@ +""" +Workgroup barrier lowering for @Gpu. + +Maps CUDA-style `__sync_threads()` and beginner-friendly +`Barrier.arrive_and_wait()` (no instance) to the same GLSL workgroup barrier. +Does not implement host lock semantics. +""" + +from __future__ import annotations + +import ast + +from ..context import GpuTranslationContext +from .base import CallPlugin, TranslateExpr + +# Single GLSL fragment used as an expression-statement body (trailing ; added +# by GpuFlow.expr_stmt). Workgroup scope only — not grid-wide. +_GPU_BARRIER_GLSL = "barrier(); memoryBarrierShared()" + + +def _require_no_args(node: ast.Call, ctx: GpuTranslationContext, label: str) -> None: + if node.keywords or node.args: + raise TypeError( + f"GPU function {ctx.func_name}: {label} takes no arguments" + ) + + +class SyncThreadsPlugin(CallPlugin): + """ + `__sync_threads()` / `Barrier.arrive_and_wait()` -> GLSL workgroup barrier. + + Rejects `Barrier(...)` construction inside @Gpu (use arrive_and_wait). + """ + + def try_lower( + self, + node: ast.Call, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, + ) -> str | None: + del translate_expr # barrier forms take no nested exprs + fn = node.func + + # CUDA-style free call. + if isinstance(fn, ast.Name) and fn.id == "__sync_threads": + _require_no_args(node, ctx, "__sync_threads()") + return _GPU_BARRIER_GLSL + + # Beginner form: Barrier.arrive_and_wait() — not Barrier(...). + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "arrive_and_wait" + and isinstance(fn.value, ast.Name) + and fn.value.id == "Barrier" + ): + _require_no_args(node, ctx, "Barrier.arrive_and_wait()") + return _GPU_BARRIER_GLSL + + # Clear error if someone writes Barrier(n) or Barrier() in a @Gpu body. + if isinstance(fn, ast.Name) and fn.id == "Barrier": + raise TypeError( + f"GPU function {ctx.func_name}: " + "Barrier(...) construction is not valid inside @Gpu; " + "use Barrier.arrive_and_wait() or __sync_threads() " + "(workgroup barrier)" + ) + + return None diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py index 5e2d56e..3f4f733 100644 --- a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py @@ -164,17 +164,22 @@ def expr_stmt(node: ast.Expr, ctx: GpuTranslationContext) -> list[str]: """ Lower expression statements; string doc-exprs are ignored. - Calls are not supported until GPU builtins / math mirrors exist. + Call expressions go through GpuSyntax.expr (CallPlugins), e.g. + `__sync_threads()` -> GLSL barrier. #### Args: - node: ast.Expr = expression statement - ctx: GpuTranslationContext = current GPU translation state #### Returns - - list[str] = empty, or a comment for unsupported forms + - list[str] = GLSL statement lines, empty for docstrings, or a comment """ + from .Syntax import GpuSyntax + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str): return [] + if isinstance(node.value, ast.Call): + return [f" {GpuSyntax.expr(node.value, ctx)};"] return [ f" // unsupported statement: Expr ({type(node.value).__name__})" ] diff --git a/src/cthreads/python/cthreads/sync/__init__.py b/src/cthreads/python/cthreads/sync/__init__.py index e8d82d4..b7869fa 100644 --- a/src/cthreads/python/cthreads/sync/__init__.py +++ b/src/cthreads/python/cthreads/sync/__init__.py @@ -4,6 +4,8 @@ - Annotation: `TBuffer[...]` (from types) - Host alloc: `create_tbuffer` / `TBufferHandle` / … - Native locks/events: re-exported from `cthreads._ext.sync` when present +- GPU workgroup barrier stub: `__sync_threads`; inside `@Gpu` also + `Barrier.arrive_and_wait()` (same GLSL lowering, no Barrier(...) call) """ from __future__ import annotations @@ -19,6 +21,24 @@ tbuffer_read_copy_ptr, ) + +def __sync_threads() -> None: + """ + Workgroup barrier stub (CUDA-style). + + Only valid inside `@Gpu` bodies; compiled to GLSL barrier() / + memoryBarrierShared(). Same device sync as Barrier.arrive_and_wait() + on the GPU path. + + #### Raises + - RuntimeError = called from ordinary Python (not compiled @Gpu) + """ + raise RuntimeError( + "cthreads.sync.__sync_threads() is only valid inside @Gpu bodies " + "(workgroup barrier; compiled to GLSL barrier())" + ) + + try: from cthreads import _ext as _ext except ImportError: @@ -52,4 +72,5 @@ "RWLock", "Barrier", "TBufferI64", + "__sync_threads", ] diff --git a/tests/unit/test_gpu_syntax.py b/tests/unit/test_gpu_syntax.py index fb617ef..3b7546a 100644 --- a/tests/unit/test_gpu_syntax.py +++ b/tests/unit/test_gpu_syntax.py @@ -600,6 +600,34 @@ def k(n: int) -> None: assert lines[0].startswith(" // unsupported") +def test_sync_threads_expr_stmt_lowers_to_barrier(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("__sync_threads()", ctx) + assert any("barrier()" in line for line in lines) + assert any("memoryBarrierShared()" in line for line in lines) + + +def test_barrier_arrive_and_wait_expr_stmt_same_barrier(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("Barrier.arrive_and_wait()", ctx) + assert any("barrier()" in line for line in lines) + + +def test_barrier_constructor_rejected_in_gpu(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="Barrier\\(\\.\\.\\.\\) construction"): + _stmt("Barrier(4)", ctx) + + def test_unsupported_stmt_comment(): def k(n: int) -> None: pass