[Perf] Tiles 4b: Migrate tiles16 tests to enable fastcache

[hughperkins]

Issue: #

Brief Summary

copilot:summary

Summary

Migrate all 31 tile16 test kernels to fastcache=True by passing the tile size as a qd.Template parameter and explicitly parameterizing all closure-captured variables.
Base branch: hp/tiles-4a

What's in this PR

A single file change to tests/python/test_tile16.py (+273 / -222 lines):

• All 31 @qd.kernel decorators changed to @qd.kernel(fastcache=True).
• Adds tile_size: qd.Template as the last parameter to every kernel, passed as _TILE at each call site. Inside kernels, _TILE references are replaced with tile_size
  (for block_dim, range(), and _load/_store args).
• Captured integer variables (e.g. src_offset, src_row_end, batch, NCOLS, NBATCH) are added as qd.i32 kernel parameters to satisfy purity requirements.
• Removes test_tile16_size_constant_in_kernel (it tests Tile.SIZE access inside a kernel, which requires the public Tile16x16 proxy not yet available on this branch).

Good points

• No library code changes: the tile_size: qd.Template approach avoids needing the Tile16x16Proxy, simt/__init__.py changes, or purity checker modifications. The
  entire PR is test-only.
• Exercises the purity checker: all kernels now run under fastcache=True, catching any accidental closure captures that would break caching.
• Clean layering: builds directly on tiles-4a (tensor_type parametrization), keeping each PR focused on one concern.
• All 276 tests pass with -W error::UserWarning (80 skipped: f64 on Vulkan).

Bad points

• Verbose call sites: kernels like test_tile16_load_store now have 11 arguments at the call site, which is noisy.
• Removes a test: test_tile16_size_constant_in_kernel is dropped rather than adapted, reducing coverage for the Tile.SIZE access path.
• tile_size is redundant: it's always _TILE (16) — every kernel receives the same constant. This is a testing-only workaround; production code would use the
  Tile16x16 proxy directly.

Walkthrough

copilot:walkthrough

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4c88b3abea

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Restore in-kernel Tile.SIZE coverage

This change removes the test_tile16_size_constant_in_kernel case and leaves only a host-side Tile.SIZE assertion, so we no longer verify that Tile.SIZE can be referenced from kernel code. That behavior is where purity/lowering regressions would surface, so a break could now slip through CI even though Python-side Tile.SIZE still equals 16. Please keep an in-kernel assertion (fastcache-compatible if needed) rather than dropping the check.

Useful? React with 👍 / 👎.

[claude]

🔴 This PR introduces silent cache key collisions in test_tile16_zeros and test_tile16_eye: both kernels close over a boolean (use_zeros_alias and inplace respectively) and branch on it via qd.static(), but the fastcache disk-cache key is derived solely from the kernel source text, explicit argument types/values, config, filepath, and start_lineno — none of which capture the closed-over variable. Two parameterized variants with the same dtype and tensor_type but opposite closure values hash to the same cache key; whichever runs first caches its compiled kernel, and the other silently loads the wrong kernel on subsequent pytest runs. Fix by converting use_zeros_alias and inplace into explicit qd.Template parameters (so they enter the args hash), or by removing fastcache=True from these two tests.

Extended reasoning...

What the bug is and how it manifests

In test_tile16_zeros, kernel k1 closes over the pytest parameter use_zeros_alias and uses qd.static(use_zeros_alias) to select between two structurally different compiled kernel bodies: the True branch calls Tile.zeros() (guaranteed zero-init), while the False branch calls Tile() (uninitialized tile). Similarly, test_tile16_eye closes over inplace and uses qd.static(inplace) to select between an in-place eye_() path (load + eye_() + store) and a factory path (Tile.eye() + store). Because qd.static() is a compile-time branch, each distinct value of the closed-over boolean produces a different compiled GPU kernel. The bug manifests across pytest sessions: on the second run with a warm disk cache, the second parameterized variant (by whichever closure value) loads the wrong compiled kernel silently.

The specific code path that triggers it

The fastcache disk-cache key is created by src_hasher.create_cache_key(). Per the verified implementation, this key is composed of: (1) the quadrants version string, (2) a kernel_hash = sha256 of the source lines of k1 at the same file and start_lineno — which is textually identical for all pytest parameterizations of the same test function, (3) an args_hash derived from the explicit kernel parameters (dst_arr: Ann, tile_size: qd.Template=16 for test_tile16_zeros; src_arr, dst_arr, tile_size for test_tile16_eye) — also identical across both closure values when dtype and tensor_type match, (4) config_hash, (5) filepath, (6) start_lineno. The closed-over use_zeros_alias and inplace variables are not explicit kernel parameters and therefore contribute nothing to the cache key.

Why existing code does not prevent it

The purity checker in ast_transformer.py skips violations for variables referenced inside qd.static() scope (line 92: violations suppressed when is_in_static_scope is True). This means no warning is raised at kernel-definition time. The in-memory kernel cache (materialize()) keys on the Python function object itself, so within a single test variant there is no collision. Only the disk cache (fastcache / PythonSideCache, accessed via _try_load_fastcache()) uses the source-hash-based key — and it is the disk cache that collides.

What the impact would be

On any CI run or developer machine where the disk cache is warm (i.e., a prior pytest run of the same test suite has populated it), the second parameterized variant of test_tile16_zeros[..., use_zeros_alias=True] may silently execute the use_zeros_alias=False compiled kernel (Tile() uninitialized), and vice-versa. In test_tile16_zeros both code paths coincidentally zero-initialize the tile, so assertions pass — the test appears green while secretly exercising the wrong code path, hiding potential regressions in Tile.zeros(). In test_tile16_eye the two compiled kernels happen to produce the same identity-matrix output, so assertions also pass silently. The danger is not current test failures but future regressions in the compiled code paths that will be invisible behind the wrong cached kernel.

Step-by-step proof of collision

Consider test_tile16_zeros parametrized over qd_dtype=qd.f32, tensor_type=qd.ndarray, use_zeros_alias=False (variant A) and use_zeros_alias=True (variant B), running in two separate pytest sessions:

Session 1, variant A runs first: k1 is defined at file line L (identical for both variants). src_hasher.create_cache_key computes key K from (version, hash(lines_of_k1), hash(NDArray[f32,2], Template=16), config, filepath, L). k1 is compiled with use_zeros_alias=False (Tile() body), result written to disk at key K.

Session 1, variant B runs second: k1 is defined at the same file line L with identical source text. create_cache_key again produces key K (use_zeros_alias is not in any component). _try_load_fastcache(K) finds the cache entry, loads variant A's compiled kernel (Tile() body), and skips recompilation. The test passes because Tile() happens to zero the memory, but the Tile.zeros() code path was never actually compiled or exercised.

Session 2 and beyond: same collision repeats indefinitely. Changing Tile.zeros() would not be caught by this test. The same pattern applies symmetrically to test_tile16_eye with closed-over inplace.

How to fix it

The correct fix is to make the compile-time branch selector an explicit qd.Template parameter so it enters the args_hash component of the cache key. For test_tile16_zeros: add use_zeros_alias: qd.Template to k1's signature and call k1(dst, _TILE, use_zeros_alias). For test_tile16_eye: add inplace: qd.Template to k1's signature and call k1(src, dst, _TILE, inplace). Alternatively, remove fastcache=True from these two tests specifically, accepting the per-run compilation cost in exchange for correctness.

[erizmr]

Approving on the basis that I have reviewed the design and the public facing API, tests, and they look reasonable to me.

[hughperkins]

Thank you 🙌