Skip to content

[fix] Autotuner: propagate the tuned function's return value #899

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aryaman-gupta wants to merge 2 commits into
base: main
Choose a base branch
from
+44 −2
Open

[fix] Autotuner: propagate the tuned function's return value #899

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1ed9250
[fix] Autotuner: propagate the tuned function's return value
aryaman-gupta Jul 24, 2026
b69d354
[test] Autotuner: cover return-value propagation
aryaman-gupta Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions python/flydsl/autotune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -405,9 +405,9 @@ def _run_with_hints(self, compiler_opts, args, kwargs):
from .compiler.kernel_function import CompilationContext

with CompilationContext.compile_hints(compiler_opts):
self.fn(*args, **kwargs)
return self.fn(*args, **kwargs)
else:
self.fn(*args, **kwargs)
return self.fn(*args, **kwargs)

@jhinpan jhinpan Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you document the contract this establishes? There is a sharp edge worth writing down: this only yields a value when @autotune wraps a plain Python callable. When it wraps a @flyc.jit function directly — the form the autotune() docstring shows — the host launch path structurally cannot return anything:

  • the generated entry point has zero results: func.FuncOp(self.func.__name__, (ir_types, [])) (jit_function.py:1507), terminated by func.ReturnOp([]) (:1534)
  • host invocation ends at ctypes.CFUNCTYPE(None, ctypes.c_void_p) (jit_executor.py:363), so restype is None

(JitFunction.__call__ does have a return self.func(...) at jit_function.py:1369, but that is the ir.Context.current is not None trace-time inlining branch, not a host call.)

That does not make the fix pointless — Autotuner explicitly accepts plain callables (autotune.py:231, source = fn.func if hasattr(fn, "func") else fn), so "autotune a Python wrapper that allocates and returns its output" is a real and useful pattern, and this unblocks it. It just is not the @autotune + @flyc.jit stack the docstring describes. A sentence or two on autotune() saying that the wrapper forwards the tuned callable's return value, and that a directly-wrapped jit kernel always returns None, would save the next reader the trace I just did.


def _run_config(self, config, args, kwargs):
"""Run the chosen config as a real (non-benchmark) call. Re-applies
Expand Down
42 changes: 42 additions & 0 deletions tests/unit/test_autotune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -711,5 +711,47 @@ def test_artifact_key_must_name_a_kernel_parameter():
_make_artifact_tuner(key=["mni"])


# ── return-value propagation ─────────────────────────────────────────────
def test_call_returns_tuned_fn_value(monkeypatch):
"""The tuned function's return value must propagate through __call__ on both
the search path and the cache-hit path."""
monkeypatch.setenv("FLYDSL_AUTOTUNE", "1") # force the search path first

def fn(a, out, BLOCK):
return ("result", BLOCK)

tuner = _make_tuner(
fn=fn,
configs=[Config(BLOCK=64), Config(BLOCK=128)],
do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1],

@jhinpan jhinpan Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This stub returns a constant 1.0 for every config, so both configs tie at exactly 1.000 ms and min(results, key=lambda x: x[1]) returns the first list element. BLOCK=64 "wins" because it is listed first, not because it is fastest — so the assertion four lines down is really pinning CPython's min tie-break plus the literal ordering of configs.

Reversing the list is enough to flip the result:

order [64,128], constant 1.0 -> ('result', 64)
order [128,64], constant 1.0 -> ('result', 128)

Two consequences: the test cannot distinguish "returns the best config's value" from "returns the first config's value", and any future change to selection (sorting results, batching the benchmark, a different tie-break) would break it for reasons unrelated to return propagation.

Giving the configs distinct times makes the winner unambiguous and order-independent. I checked that this variant returns ('result', 128) under both list orderings:

times = iter([2.0, 1.0])  # the second config is genuinely faster
tuner = _make_tuner(
    fn=fn,
    configs=[Config(BLOCK=64), Config(BLOCK=128)],
    do_bench_fn=lambda call, warmup, rep: (call(), next(times))[1],
)
...
assert searched == ("result", 128)

)
args = (FakeTensor((8,)), FakeTensor((1,)))

searched = tuner(*args) # search path -> runs the winning config
assert searched == ("result", 64)

monkeypatch.setenv("FLYDSL_AUTOTUNE", "0")
hit = tuner(*args) # cache-hit path
assert hit == ("result", 64)

@jhinpan jhinpan Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This half does not actually exercise the cache-hit path. _make_tuner is called without default=, so on a cache miss __call__ falls straight through to a full re-search and returns the same value — the assertion is satisfied either way. Wiping the cache right before this line keeps it green:

search -> ('result', 64)
[autotune] tuning 2 configs...          <-- re-searched, did not hit the cache
after cache.clear() -> ('result', 64)   <-- assertion still passes

The file already has the idiom for this: test_cache_hit_precedes_default_and_search installs a fail_bench that blows up if the search runs. Swapping tuner._do_bench for something that calls pytest.fail before the second call would make this assertion mean what it says.



def test_call_returns_none_when_fn_returns_none(monkeypatch):

@jhinpan jhinpan Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor, just so it is not miscounted as regression coverage: this one passes against the unpatched code too, because the buggy __call__ also returns None. Reverting the two returns gives:

test_call_returns_tuned_fn_value             FAILED   (assert None == ('result', 64))
test_call_returns_none_when_fn_returns_none  PASSED

Still worth keeping as a guard against an over-correction — someone later "improving" this to hand back out or the winning Config when fn returns nothing — but test_call_returns_tuned_fn_value is the only one of the two that actually detects the bug.

"""In-place kernels (write into `out`, return nothing) still return None --
the fix only forwards whatever the tuned fn returns, it doesn't fabricate."""
monkeypatch.setenv("FLYDSL_AUTOTUNE", "1")

def fn(a, out, BLOCK):
out._data[0] = float(BLOCK) # writes in place, returns None

tuner = _make_tuner(
fn=fn,
configs=[Config(BLOCK=64)],
do_bench_fn=lambda call, warmup, rep: (call(), 1.0)[1],
)
args = (FakeTensor((8,)), FakeTensor((1,)))
assert tuner(*args) is None
assert args[1]._data[0] == 64.0

@jhinpan jhinpan Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now that this is rebased there are four __call__ exits, and these tests cover two of them:

exit main line covered
cache hit 570 yes, weakly (see above)
offline artifact hit 576 no
default heuristic 579 no
post-search winner 604 yes

The default path is the one I would most want covered: it is what the only real consumer uses (rmsnorm_autotune.py passes default=_default_config), and it is its own return statement, so it is exactly the line a future refactor could silently re-break. The artifact path is the one that did not exist when you opened this.

Both are cheap now — #786 already ships _make_artifact_tuner and the artifact fixture just above your block. A single parametrized test walking all four exits would pin the contract more tightly than the two tests as written.



if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))
Loading