From 54fa88064ea76b3e9abe5b2da7dc54b67e98ac0c Mon Sep 17 00:00:00 2001 From: Aryaman Gupta Date: Fri, 24 Jul 2026 14:48:40 +0000 Subject: [PATCH 1/2] [fix] Autotuner: propagate the tuned function's return value _run_with_hints called self.fn(*args, **kwargs) on both branches without returning it, so Autotuner.__call__ always yielded None even though _run_config already returns _run_with_hints(...). Any autotuned function that produces its result as a return value (rather than writing into a caller-owned output buffer) silently got None. Add the missing return on both branches; backward compatible for callers that ignore the return value. --- python/flydsl/autotune.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/flydsl/autotune.py b/python/flydsl/autotune.py index 3b4cad77b..ce0848cd8 100644 --- a/python/flydsl/autotune.py +++ b/python/flydsl/autotune.py @@ -350,9 +350,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) def _run_config(self, config, args, kwargs): """Run the chosen config as a real (non-benchmark) call. Re-applies From d759f6c8010426902420d1b2d880df9bc7290e7d Mon Sep 17 00:00:00 2001 From: Aryaman Gupta Date: Fri, 24 Jul 2026 15:03:02 +0000 Subject: [PATCH 2/2] [test] Autotuner: cover return-value propagation Add a regression test that the tuned function's return value propagates through __call__ on both the search and cache-hit paths, plus a guard that in-place kernels (which return None) still return None. The first test fails without the _run_with_hints fix. --- tests/unit/test_autotune.py | 42 +++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/unit/test_autotune.py b/tests/unit/test_autotune.py index 8fd90d514..9229fefe5 100644 --- a/tests/unit/test_autotune.py +++ b/tests/unit/test_autotune.py @@ -454,5 +454,47 @@ def bench(call, warmup, rep): assert args[1]._data[0] == 64.0 +# ── 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], + ) + 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) + + +def test_call_returns_none_when_fn_returns_none(monkeypatch): + """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 + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"]))