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

[aryaman-gupta]

Summary

Autotuner._run_with_hints calls self.fn(*args, **kwargs) on both branches (with and without compiler hints) but does not return it. As a result Autotuner.__call__ always yields None, even though _run_config does return self._run_with_hints(...) and __call__ returns _run_config(...) — the value is threaded all the way up to this one method, which drops it.

Any autotuned function that produces its result as a return value (rather than writing into a caller-owned output buffer) therefore silently gets None. Callers currently work around this by allocating an output tensor themselves, passing it in, and returning that buffer instead of trusting the autotuner's return.

Fix

Add the missing return on both branches of _run_with_hints.

Backward compatible:

• Callers that ignore the return value (in-place kernels writing into out) are unaffected — the fn still returns None, so the tuner still returns None.
• Callers whose fn returns its result now receive it.

Test

Adds two cases to tests/unit/test_autotune.py (GPU-free):

• test_call_returns_tuned_fn_value — the tuned fn's return value propagates through __call__ on both the search path and the cache-hit path. Fails without the fix (assert None == ('result', 64)).
• test_call_returns_none_when_fn_returns_none — an in-place fn that returns nothing still yields None, confirming the fix only forwards the value and doesn't fabricate one.

Run:

pytest tests/unit/test_autotune.py -k return -v

[coderfeli]

LGTM

[jhinpan]

Thanks for catching this — the bug is real and the fix is in the right place. I see you have already rebased onto current main; I re-verified everything against b69d3549. Notes below, then a few suggestions inline. None of them are objections to the two-line change itself.

The bug reproduces and the fix works. Against the pre-fix code a tuned fn returning ("result", 64) yields None from __call__; with the patch it yields the value. Reverting just the two returns makes test_call_returns_tuned_fn_value fail and every other test pass, so the regression is genuinely pinned.

The fix is complete. _run_with_hints holds the only two self.fn(...) call sites in the module, and _run_config already did return self._run_with_hints(...), so those two branches were the entire gap. _bench_one still discards the value, which is correct — benchmark reps only need timing, and not retaining per-rep outputs is what you want.

It also covers a path that appeared after you opened this. #786 added a fourth __call__ exit, the offline-artifact hit at main:576. Since all four exits funnel through _run_config, your two lines cover it for free — but nothing pins it; see the last comment.

No behaviour change for the only in-repo consumer. rmsnorm_direct ends in a bare launch(...) and returns None, so rmsnorm_autotuned returns None before and after. The rebased tests/unit/test_autotune.py passes 42/42 for me, with #786's artifact block and your new block coexisting cleanly.

Happy to approve once the two return-value tests are tightened — the suggestions below are all small.

[jhinpan]

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.

[jhinpan]

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)

[jhinpan]

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.

[jhinpan]

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.

[jhinpan]

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.