[Perf] Tiles 2b: add triangular solve

[hughperkins]

Issue: #

Brief Summary

This PR adds triangular solve to Tile16x16, building on the Cholesky factorization from PR 2a. Together, cholesky_ + solve_triangular_ give users the two primitives needed for
blocked Cholesky-based linear solves entirely in registers.
Two methods are added to _tile16.py:
• _trsm(self, L) -- the internal workhorse. Solves X @ L^T = B in-place (self starts as B, ends as X) by forward-substituting column-by-column. Each column c computes a dot
product of previously-solved columns against row c of L (fetched via subgroup.shuffle), subtracts from the current column, and divides by the diagonal. Uses the
_get_col/_set_col helpers from PR 2a.
• solve_triangular_(self, B, lower=True) -- the public API wrapper. Validates lower=True at compile time, then delegates to B._trsm(self). The argument order follows the
convention "self is the factored matrix, B is modified in-place."

The test loads a pre-computed Cholesky factor L and a random matrix B, runs L.solve_triangular_(B), and verifies X @ L^T ≈ B against scipy. Parametrized over f32/f64.
Good points:
• Very small and focused: 68 lines total (28 impl + 40 test), easy to review.
• Reuses _get_col/set_col from PR 2a, so the algorithm reads cleanly as a loop rather than 16-way unrolled register access.
• The algorithm structure mirrors cholesky (column-by-column forward substitution with shuffles), making it easy to verify correctness by analogy.
• Test validates against scipy with tight f64 tolerance.

Bad points:
• Only lower=True is supported; lower=False raises at compile time rather than implementing upper-triangular solve.
• Test only exercises the full 16x16 case (offset 0, full tile). No partial-tile or offset-based parametrization like the Cholesky test in PR 2a has.
• No test for the lower=False error path.

copilot:summary

Walkthrough

copilot:walkthrough

[hughperkins]

Questions to opus:

Is there any code or code change in this PR that might be embarrassing in some way if published in a public github repo?

  Yes, two issues:
  1. The test reference computation is wrong. The _trsm solves X @ L^T = B, but the scipy reference computes (L^{-T} @ B)^T = B^T @ L^{-1} instead of B @ L^{-T}. This test
     actually fails (we saw this earlier). The fix is B.T not trans="T":


  # Wrong (current):
  X_ref = scipy.linalg.solve_triangular(L_ref, B, lower=True, trans="T").T
  # Correct:
  X_ref = scipy.linalg.solve_triangular(L_ref, B.T, lower=True).T

  2. Missing `# type: ignore[reportOperatorIssue]` in _trsm -- the same DSL operator expressions (+=, *, /) that needed suppression in cholesky_ appear here without it. Pyright
     CI will fail on this branch too.

  Let me also check if the protocol stubs were updated.

  ⬢ Grepped 4 greps
    Grepped "_trsm|solve_triangular_" in quadrants/python/quadrants/lang/simt/_tile16.py
    Grepped "def _trsm|def solve_triangular_" in quadrants/python/quadrants/lang/simt/_tile16.py
    Grepped "class _Tile16x16Proto" in quadrants/python/quadrants/lang/simt/_tile16.py
    Grepped "_Tile16x16Proto" in quadrants/python/quadrants/lang/simt/_tile16.py

  3. `solve_triangular_` is missing `@qd.func` -- cholesky_ and _ger_sub both have it, but solve_triangular_ does not. It'll fail when called from a kernel.

  Want me to fix all three issues?

why dont the tests catch this?

  Good question -- let me reconsider each:
  Issue 1 (wrong reference): The test does fail -- we saw it earlier when running on hp/tiles-2b. So this IS caught. It's a test bug (wrong expected value), not a code bug.
  Issue 2 (missing pyright suppression): This would be caught by pyright CI, same as the hp/tiles-2a failure you just linked.
  Issue 3 (missing `@qd.func`): Let me check whether this actually matters.

  ⬢ Grepped "..._eye_|_load|_store|_get_col|_set_col)" in quadrants/python/quadrants/lang/simt/_tile16.py
    Found 398 matches

  I was wrong about issue #3. The tests on tiles-4 call L.solve_triangular_(B) inside a kernel and all 785 tests pass. So solve_triangular_ works without @qd.func -- Quadrants'
  AST transformer inlines dataclass method bodies regardless of the decorator. The if not lower: raise TypeError(...) acts as a compile-time guard, not runtime, which is why it
  works.
  So to correct my answer: there are really two issues, not three:
  1. Wrong test reference -- the test does fail (we saw it earlier), so it IS caught. It's just a test bug that was shipped. Fix: B.T instead of trans="T".
  2. Missing pyright suppression in `_trsm` -- pyright CI will catch this, same as it caught the cholesky_ issue on 2a.

  Want me to fix both on hp/tiles-2b and propagate?

yes

  Done. Both fixes applied to hp/tiles-2b and propagated to all downstream branches:
  • Test fix: solve_triangular(L, B.T, lower=True).T instead of solve_triangular(L, B, lower=True, trans="T").T
  • Pyright fix: # type: ignore[reportOperatorIssue] on the two operator lines in _trsm

[hughperkins]

add test for the lower-False error path

Done. Added test_tile16_solve_triangular_upper_raises which verifies that calling solve_triangular_(..., lower=False) raises TypeError with the expected message.

[claude]

LGTM — small, focused PR with correct triangular solve implementation.

Extended reasoning...

Overview

Adds two methods to _tile16.py: the internal _trsm workhorse (in-place triangular solve via subgroup shuffles) and the public solve_triangular_ wrapper. Test validates against scipy over f32/f64. Changes touch 2 files, ~68 lines total.

Security risks

None. No auth, I/O, or user-facing surfaces touched.

Level of scrutiny

Low. This is internal GPU tile arithmetic — sandbox/library code, not production-critical infrastructure. The algorithm mirrors the existing cholesky_ pattern (same shuffle mechanics), which lowers the verification bar.

Other factors

Known issues (wrong test reference, missing pyright suppressions) were fixed prior to this review based on the PR timeline. Two nits flagged by the bug hunter (misleading Lkj variable name and unnecessary GPU decorator on a pure-Python test) are cosmetic/minor — neither affects correctness of the implementation or the primary test coverage.

[claude]

🟡 The variable Lkj on line 433 of _trsm is a copy-paste artifact from cholesky_ where the outer loop variable was k; in _trsm the outer loop variable is c, so the shuffle actually fetches L[c, j] and the variable should be named Lcj. No runtime effect — purely a naming issue.

Extended reasoning...

What the bug is and how it manifests

In _trsm, the outer loop iterates over column index c (for c in range(_TILE)). Inside the inner loop, the subgroup shuffle qd.simt.subgroup.shuffle(L._get_col(j), qd.u32(c)) broadcasts to all threads the value held by thread c in register j — in other words, it retrieves L[c, j]. The result is stored in a local variable named Lkj. The name strongly implies the outer loop variable is k, not c.

The specific code path that triggers it

Lines 433–434 of python/quadrants/lang/simt/_tile16.py:

for c in range(_TILE):
    ...
    for j in range(_TILE):
        if c > j:
            Lkj = qd.simt.subgroup.shuffle(L._get_col(j), qd.u32(c))  # line 433
            dot += self._get_col(j) * Lkj                              # line 434

Why existing code doesn't prevent it

This is a pure naming issue. Python does not enforce any relationship between a variable name and the loop variable it conceptually tracks, and the Quadrants DSL compiler is only interested in the value of the shuffle result, not its name. No static analysis or test will flag a misleading identifier name.

What the impact would be

The code is functionally correct. The misleading name causes confusion only for readers who compare _trsm against cholesky_, where Lkj correctly names L[k, j] because the outer loop variable there is k. A developer reading _trsm expecting k to be the outer loop variable (by analogy) will be briefly confused.

How to fix it

Rename Lkj to Lcj on the two lines where it appears (assignment and use):

Lcj = qd.simt.subgroup.shuffle(L._get_col(j), qd.u32(c))
dot += self._get_col(j) * Lcj  # type: ignore[reportOperatorIssue]

Step-by-step proof

1. In cholesky_ (lines ~411–418), the outer loop is for k in range(_TILE). The shuffle qd.simt.subgroup.shuffle(my_col, qd.u32(k)) fetches L[k, j] from thread k. Naming this Lkj is accurate.
2. In _trsm (lines ~430–434), the outer loop is for c in range(_TILE). The shuffle qd.simt.subgroup.shuffle(L._get_col(j), qd.u32(c)) fetches L[c, j] from thread c. Naming this Lkj is inaccurate — there is no variable k anywhere in _trsm.
3. The name was carried over verbatim when _trsm was written by analogy with cholesky_, but the loop variable was renamed from k to c without updating the derived variable name.

[claude]

🟡 test_tile16_solve_triangular_upper_raises is decorated with @test_utils.test(arch=qd.gpu), causing it to be silently skipped on non-GPU machines. Since the test only exercises a pure-Python TypeError raised before any GPU code runs, the decorator is unnecessary and should be removed.

Extended reasoning...

The test test_tile16_solve_triangular_upper_raises (tests/python/test_tile16.py:437-441) is decorated with @test_utils.test(arch=qd.gpu). Looking at test_utils.py lines 226-234, when qd.gpu is not present in expected_archs() (i.e., no GPU is available), the arch list becomes empty and pytest.mark.skip is applied to the test function. This means the test is silently skipped on any CI agent or developer machine without GPU hardware.

However, the test body is entirely pure Python. It calls make_tile16x16(qd.f32) to get a dataclass type, calls Tile() to instantiate it, and then calls solve_triangular(Tile(), lower=False). The solve_triangular_ method (in _tile16.py) immediately checks "if not lower: raise TypeError(...)" before delegating to B._trsm(self) or doing anything GPU-related. No @qd.func decorated path is ever entered, no kernel is compiled or launched.

Existing pure-Python tests in the same file correctly omit the GPU decorator. For example, test_tile16_size_constant (which just checks Tile.SIZE == 16) and test_tile16_make_caching (which checks the caching behavior of _make_tile16x16) both run as plain pytest functions with no decorator, and conftest.py lines 79-83 confirm that undecorated functions receive (None, None) for req_arch/req_options and run unconditionally as pure Python.

The impact is silent coverage loss on CPU-only CI agents. The lower=False error-path guard in solve_triangular_ goes entirely untested in those environments, at zero cost to fix.

Step-by-step proof on a non-GPU machine: (1) pytest collects test_tile16_solve_triangular_upper_raises; (2) the @test_utils.test(arch=qd.gpu) decorator runs, finds qd.gpu not in expected_archs(), so arch=[], and applies pytest.mark.skip; (3) pytest marks the test as skipped; (4) the body -- make_tile16x16, Tile(), solve_triangular(lower=False) -- never executes; (5) the TypeError guard in solve_triangular_ is never validated.

The fix is to remove the @test_utils.test(arch=qd.gpu) decorator from test_tile16_solve_triangular_upper_raises, exactly as was done for the other pure-Python tests in the same file. No other changes are needed.

[erizmr]

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

[hughperkins]

Thanks! 🙌