[Feat] Add fx.known_block_size() as constant accessor (#968)

* [Feat] Add fx.known_block_size() as constant accessor

* accept None as the trace_option value

Author: sjfeng1999

python/flydsl/compiler/jit_function.py

@@ -22,7 +22,6 @@
 from .._mlir.passmanager import PassManager
 from ..expr.meta import tracing_context
 from ..expr.typing import Constexpr, Stream
-from ..expr.utils.arith import fastmath as fastmath_ctx
 from ..utils import env, log
 from ..utils.file import atomic_write
 from .ast_rewriter import ASTRewriter
@@ -1511,12 +1510,11 @@ def __call__(self, *args, **kwargs):
                                 log().info(f"dsl_args={dsl_args}")
                                 named_args = dict(zip(param_names, dsl_args))
                                 named_args.update(constexpr_values)
-                                fastmath_flag = effective_fastmath_hint(CompilationContext.get_compile_hints())
-                                fastmath_scope = (
-                                    fastmath_ctx(fastmath_flag) if fastmath_flag is not None else nullcontext()
-                                )
                                 # Bound the call-site boundary at the jit body.
-                                with tracing_context(self.func), fastmath_scope:
+                                with tracing_context(
+                                    self.func,
+                                    fastmath=effective_fastmath_hint(CompilationContext.get_compile_hints()),
+                                ):
                                     if bound_self is not None:
                                         self.func(bound_self, **named_args)
                                     else:

python/flydsl/compiler/kernel_function.py

@@ -3,7 +3,7 @@
 
 import inspect
 import threading
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -12,7 +12,6 @@
 from ..expr.meta import capture_user_location, file_location, tracing_context
 from ..expr.numeric import Index, Integer
 from ..expr.typing import Constexpr, as_ir_value
-from ..expr.utils.arith import fastmath as fastmath_ctx
 from .ast_rewriter import ASTRewriter
 from .diagnostics import install_excepthook, warn_annotation_value_mismatch, warn_invalid_annotations
 from .jit_argument import is_type_param_annotation, resolve_signature
@@ -590,10 +589,13 @@ def _emit_kernel(
 
                     dsl_args.update(constexpr_values)
 
-                    fastmath_flag = effective_fastmath_hint(CompilationContext.get_compile_hints())
-                    fastmath_scope = fastmath_ctx(fastmath_flag) if fastmath_flag is not None else nullcontext()
-                    # Bound the call-site boundary at the kernel body.
-                    with tracing_context(self._func), fastmath_scope:
+                    # Bound the call-site boundary at the kernel body and carry
+                    # the ambient tracing options into it.
+                    with tracing_context(
+                        self._func,
+                        fastmath=effective_fastmath_hint(CompilationContext.get_compile_hints()),
+                        known_block_size=tuple(known_block_size) if known_block_size is not None else None,
+                    ):
                         if bound_self is not None:
                             self._func(bound_self, **dsl_args)
                         else:

python/flydsl/expr/gpu.py

@@ -21,7 +21,7 @@
 from .._mlir.dialects._fly_enum_gen import AddressSpace
 from ..compiler.protocol import dsl_align_of, dsl_size_of
 from .math import dsl_math_wrap_result
-from .meta import dsl_loc_tracing
+from .meta import dsl_loc_tracing, tracing_option
 from .numeric import Int32, Numeric, Uint8
 from .primitive import get_dyn_shared, make_ptr
 from .struct import (
@@ -45,6 +45,7 @@
     "shuffle_up",
     "shuffle_down",
     "shuffle_idx",
+    "known_block_size",
     "SharedAllocator",
 ]
 
@@ -103,6 +104,17 @@ def shuffle_idx(value, lane, width):
     return shuffle(value, lane, width, mode="idx")
 
 
+def known_block_size():
+    """Return the compile-time block dimensions as ``(x, y, z)`` Python ints.
+
+    Raises ``RuntimeError`` when no block size is in scope.
+    """
+    size = tracing_option("known_block_size")
+    if size is None:
+        raise RuntimeError("no compile-time block size is in scope.")
+    return size
+
+
 thread_idx = Tuple3D(gpu.thread_id)
 block_idx = Tuple3D(gpu.block_id)
 block_dim = Tuple3D(gpu.block_dim)

python/flydsl/expr/meta.py

@@ -22,6 +22,7 @@
     "dsl_loc_tracing",
     "dsl_wrap_result",
     "tracing_context",
+    "tracing_option",
 ]
 
 # Package root for the ``flydsl`` Python package: ``.../python/flydsl``.
@@ -54,8 +55,20 @@ def _stack():
 
 
 @contextlib.contextmanager
-def tracing_context(func):
-    _stack().append(getattr(func, "__code__", None))
+def tracing_context(func=None, **options):
+    """Push a tracing frame for the duration of the ``with`` block.
+
+    *func*, when given, marks the call-site boundary used by
+    ``capture_user_location``; a frame entered without one inherits the
+    enclosing boundary.  Keyword *options* become ambient settings readable
+    with ``tracing_option(key)``.  Every key passed is set, ``None`` included
+    -- an explicit ``None`` shadows an enclosing setting.
+    """
+    stack = _stack()
+    boundary = getattr(func, "__code__", None)
+    if boundary is None and stack:
+        boundary = stack[-1][0]
+    stack.append((boundary, options))
     try:
         yield
     finally:
@@ -64,6 +77,16 @@ def tracing_context(func):
             stack.pop()
 
 
+def tracing_option(key, default=None):
+    """Return the ambient value of *key*, innermost tracing frame first."""
+    stack = getattr(_tls, "stack", None)
+    if stack:
+        for _, options in reversed(stack):
+            if key in options:
+                return options[key]
+    return default
+
+
 def file_location(filename: str, line: int, col: int = 0, context=None) -> ir.Location:
     ctx = context or ir.Context.current
     if filename and not filename.startswith("<"):
@@ -79,7 +102,7 @@ def capture_user_location() -> ir.Location:
     tracing boundary.
     """
     stack = getattr(_tls, "stack", None)
-    boundary = stack[-1] if stack else None
+    boundary = stack[-1][0] if stack else None
     max_depth = env.debug.max_loc_depth
     ctx = ir.Context.current
     locs = []

python/flydsl/expr/utils/arith.py

@@ -3,18 +3,16 @@
 
 import builtins
 import contextlib
-import threading
 from functools import partialmethod
 
 from ..._mlir import ir
 from ..._mlir.dialects import arith, math
 from ..._mlir.extras import types as T
-from ..meta import dsl_loc_tracing
+from ..meta import dsl_loc_tracing, tracing_context, tracing_option
 
 # --------------------------------------------------------------------------- #
-# Ambient fastmath context (thread-local)
+# Ambient fastmath context (a ``tracing_context`` option)
 # --------------------------------------------------------------------------- #
-_fm_tls = threading.local()
 
 
 def _normalize_fastmath(flags):
@@ -37,7 +35,7 @@ def _normalize_fastmath(flags):
 
 def current_fastmath():
     """Return the ambient fastmath flags set by ``fastmath(...)``, or ``None``."""
-    return getattr(_fm_tls, "value", None)
+    return _normalize_fastmath(tracing_option("fastmath"))
 
 
 def resolve_fastmath(explicit):
@@ -48,12 +46,8 @@ def resolve_fastmath(explicit):
 @contextlib.contextmanager
 def fastmath(flags):
     """Apply *flags* to floating-point ops built inside the ``with`` block."""
-    prev = getattr(_fm_tls, "value", None)
-    _fm_tls.value = _normalize_fastmath(flags)
-    try:
+    with tracing_context(fastmath=_normalize_fastmath(flags)):
         yield
-    finally:
-        _fm_tls.value = prev
 
 
 def element_type(ty) -> ir.Type:

tests/unit/test_kernel_known_block_size.py

@@ -285,3 +285,74 @@ def _launch_big(x: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
 
         source_ir = _get_source_ir(_launch_big, self.x)
         assert "known_block_size = array<i32: 512, 1, 1>" in source_ir
+
+
+class TestKnownBlockSizeTraceAccessor:
+    """Verify the trace-time ``fx.known_block_size()`` view of the same value."""
+
+    @pytest.fixture(autouse=True)
+    def _setup(self):
+        self.x = torch.zeros(64, device="cuda", dtype=torch.float32)
+
+    def test_raises_outside_a_kernel(self):
+        with pytest.raises(RuntimeError, match="no compile-time block size"):
+            fx.known_block_size()
+
+    def test_sees_the_declared_size(self):
+        seen = []
+
+        @flyc.kernel(known_block_size=[128, 4, 2])
+        def _kn(x: fx.Tensor):
+            seen.append(fx.known_block_size())
+
+        @flyc.jit
+        def _launch(x: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
+            _kn(x).launch(grid=(1, 1, 1), block=(128, 4, 2), stream=stream)
+
+        _get_source_ir(_launch, self.x)
+        assert seen == [(128, 4, 2)]
+
+    def test_sees_the_size_inferred_from_static_launch_dims(self):
+        seen = []
+
+        @flyc.kernel
+        def _kn(x: fx.Tensor):
+            seen.append(fx.known_block_size())
+
+        @flyc.jit
+        def _launch(x: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
+            _kn(x).launch(grid=(1, 1, 1), block=(256, 1, 1), stream=stream)
+
+        _get_source_ir(_launch, self.x)
+        assert seen == [(256, 1, 1)]
+
+    def test_raises_for_a_dynamic_launch(self):
+        @flyc.kernel
+        def _kn(x: fx.Tensor, nthreads: fx.Int32):
+            fx.known_block_size()
+
+        @flyc.jit
+        def _launch(x: fx.Tensor, nthreads: fx.Int32):
+            _kn(x, nthreads).launch(grid=(1, 1, 1), block=(nthreads, 1, 1))
+
+        with pytest.raises(RuntimeError, match="no compile-time block size"):
+            _launch(self.x, 64)
+
+    def test_is_visible_from_a_nested_jit_helper(self):
+        """An inner tracing frame carries no size of its own and inherits the kernel's."""
+        seen = []
+
+        @flyc.jit
+        def _helper():
+            seen.append(fx.known_block_size())
+
+        @flyc.kernel(known_block_size=[64, 1, 1])
+        def _kn(x: fx.Tensor):
+            _helper()
+
+        @flyc.jit
+        def _launch(x: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
+            _kn(x).launch(grid=(1, 1, 1), block=(64, 1, 1), stream=stream)
+
+        _get_source_ir(_launch, self.x)
+        assert seen == [(64, 1, 1)]

tests/unit/test_tracing_context.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Tests for the ambient key/value options carried by ``tracing_context``."""
+
+import pytest
+
+from flydsl.expr.meta import tracing_context, tracing_option
+
+pytestmark = pytest.mark.l0_backend_agnostic
+
+
+def _traced():
+    pass
+
+
+def test_option_is_visible_inside_the_scope_only():
+    assert tracing_option("fastmath") is None
+    with tracing_context(_traced, fastmath="fast"):
+        assert tracing_option("fastmath") == "fast"
+    assert tracing_option("fastmath") is None
+
+
+def test_inner_scope_shadows_and_restores():
+    with tracing_context(fastmath="fast", unroll=4):
+        with tracing_context(fastmath="contract"):
+            assert tracing_option("fastmath") == "contract"
+            # Keys the inner frame does not set stay visible.
+            assert tracing_option("unroll") == 4
+        assert tracing_option("fastmath") == "fast"
+
+
+def test_explicit_none_shadows_an_outer_value():
+    with tracing_context(fastmath="fast"):
+        with tracing_context(_traced, fastmath=None):
+            assert tracing_option("fastmath") is None
+        assert tracing_option("fastmath") == "fast"
+
+
+def test_unset_key_inherits_the_outer_value():
+    with tracing_context(fastmath="fast"):
+        with tracing_context(_traced, unroll=4):
+            assert tracing_option("fastmath") == "fast"
+
+
+def test_default_is_returned_for_unset_keys():
+    with tracing_context(_traced):
+        assert tracing_option("missing", "fallback") == "fallback"