Numba: solve triangular C-contiguous B without copy by jessegrabowski · Pull Request #2359 · pymc-devs/pytensor

jessegrabowski · 2026-08-17T05:08:57Z

Every rhs a graph produces is c-contiguous, so SolveTriangular always fell through to a transposing copy of B before calling trtrs. It now calls trsm with the side picked from the rhs layout: side="R" solves the transposed system straight out of a c-contiguous B, side="L" takes an f-contiguous one. That takes a (128, 100_000) solve from 39ms to 14ms.

trtrs drops out of the path entirely. Reference LAPACK's trtrs is a zero-pivot scan followed by a side="L" trsm, and the side="R" path already needed that scan in our own code, so calling trsm directly costs nothing; benchmarked equal within one sigma on an f-contiguous rhs across four shapes.

A c-contiguous rhs is now consumed in place under overwrite_b, where it previously survived because trtrs copied it.

numba_xtrtrs in _LAPACK is unused after this and left in place; happy to drop it here if you'd rather.

The first two commits are unrelated: the numba linalg tests only pass under floatX=float32 with them.

Closes #2358

…rtrs Every graph-produced right-hand side is c-contiguous, so SolveTriangular always fell through to a transposing copy of B before calling trtrs; the flag-flipping trick that avoids the copy for A has no analog for B, since trtrs cannot consume a row-major rhs. trsm with side="R" solves the equivalent transposed system straight out of B's memory, cutting a (128, 100_000) solve from 39ms to 14ms. trans=2 keeps the trtrs path because op(A)^T is conj(A), which trsm has no mode for, and trsm reports no INFO, so a zero pivot has to be caught before the call to preserve the fill-with-NaN behavior. Closes pymc-devs#2358

ricardoV94 · 2026-08-17T06:15:20Z

We should at some point explore outputting in F order (or rather batched F order) directly. Question, is this the same speed for C arrays as the "regular solver" is for F arrays.

Tbis question isn't a blocker, because there will always be arrays we don't control (and right now we always emit C anyway)

jessegrabowski · 2026-08-18T01:32:17Z

I think you're asking if we could just drop trtrs entirely? The answer is yes. The layout of A basically doesn't matter because we already do transpose tricks on that. There's a difference based on layout of B. But they seem to boil down to the same kernel when both are comparable. Timings:

-------------------------------------- benchmark 'N=1024 NRHS=1024': 8 tests ---------------------------------------
Name (time in ms)                        Min              Mean            StdDev            Median            Rounds
--------------------------------------------------------------------------------------------------------------------
test_bench[trsm_Ac_Bc-1024-1024]      2.2122 (1.17)     2.3266 (1.18)     0.0969 (3.06)     2.3054 (1.19)         20
test_bench[trsm_Ac_Bf-1024-1024]      1.9029 (1.01)     1.9988 (1.02)     0.1390 (4.38)     1.9484 (1.01)         20
test_bench[trsm_Af_Bc-1024-1024]      2.1925 (1.16)     2.3513 (1.19)     0.0950 (2.99)     2.3602 (1.22)         20
test_bench[trsm_Af_Bf-1024-1024]      1.9265 (1.02)     1.9816 (1.01)     0.0530 (1.67)     1.9685 (1.02)         20
test_bench[trtrs_Ac_Bc-1024-1024]     5.2846 (2.80)     5.7961 (2.95)     0.1767 (5.57)     5.8142 (3.00)         20
test_bench[trtrs_Ac_Bf-1024-1024]     1.8869 (1.0)      1.9850 (1.01)     0.1381 (4.35)     1.9355 (1.0)          20
test_bench[trtrs_Af_Bc-1024-1024]     5.3833 (2.85)     5.8179 (2.96)     0.2298 (7.24)     5.8161 (3.00)         20
test_bench[trtrs_Af_Bf-1024-1024]     1.9195 (1.02)     1.9681 (1.0)      0.0317 (1.0)      1.9668 (1.02)         20
--------------------------------------------------------------------------------------------------------------------

---------------------------------------- benchmark 'N=128 NRHS=100000': 8 tests ----------------------------------------
Name (time in ms)                          Min               Mean            StdDev             Median            Rounds
------------------------------------------------------------------------------------------------------------------------
test_bench[trsm_Ac_Bc-128-100000]       7.0265 (1.11)      7.2238 (1.12)     0.1430 (3.20)      7.1976 (1.12)         20
test_bench[trsm_Ac_Bf-128-100000]       6.4635 (1.02)      6.5521 (1.02)     0.0662 (1.48)      6.5426 (1.02)         20
test_bench[trsm_Af_Bc-128-100000]       7.2090 (1.14)      7.3812 (1.15)     0.0896 (2.01)      7.3797 (1.15)         20
test_bench[trsm_Af_Bf-128-100000]       6.4580 (1.02)      6.5396 (1.02)     0.0446 (1.0)       6.5455 (1.02)         20
test_bench[trtrs_Ac_Bc-128-100000]     33.6925 (5.33)     35.7314 (5.56)     0.9195 (20.60)    35.8557 (5.59)         20
test_bench[trtrs_Ac_Bf-128-100000]      6.3445 (1.00)      6.4237 (1.0)      0.0458 (1.03)      6.4153 (1.0)          20
test_bench[trtrs_Af_Bc-128-100000]     34.5266 (5.47)     35.5161 (5.53)     0.5664 (12.69)    35.4475 (5.53)         20
test_bench[trtrs_Af_Bf-128-100000]      6.3171 (1.0)       6.5350 (1.02)     0.2149 (4.82)      6.4756 (1.01)         20
------------------------------------------------------------------------------------------------------------------------

-------------------------------------- benchmark 'N=32 NRHS=100000': 8 tests ---------------------------------------
Name (time in ms)                        Min              Mean            StdDev            Median            Rounds
--------------------------------------------------------------------------------------------------------------------
test_bench[trsm_Ac_Bc-32-100000]      1.5998 (1.0)      1.7807 (1.0)      0.1297 (2.86)     1.7288 (1.0)          20
test_bench[trsm_Ac_Bf-32-100000]      1.8771 (1.17)     2.1439 (1.20)     0.0758 (1.67)     2.1556 (1.25)         20
test_bench[trsm_Af_Bc-32-100000]      1.6026 (1.00)     1.7821 (1.00)     0.1131 (2.49)     1.7715 (1.02)         20
test_bench[trsm_Af_Bf-32-100000]      1.8120 (1.13)     2.1851 (1.23)     0.1010 (2.23)     2.1964 (1.27)         20
test_bench[trtrs_Ac_Bc-32-100000]     7.8987 (4.94)     8.2649 (4.64)     0.2026 (4.47)     8.2921 (4.80)         20
test_bench[trtrs_Ac_Bf-32-100000]     1.9898 (1.24)     2.0676 (1.16)     0.0453 (1.0)      2.0636 (1.19)         20
test_bench[trtrs_Af_Bc-32-100000]     7.2659 (4.54)     8.1056 (4.55)     0.2981 (6.58)     8.1318 (4.70)         20
test_bench[trtrs_Af_Bf-32-100000]     1.8752 (1.17)     2.1146 (1.19)     0.0864 (1.91)     2.1072 (1.22)         20
--------------------------------------------------------------------------------------------------------------------

------------------------------------------ benchmark 'N=512 NRHS=128': 8 tests ------------------------------------------
Name (time in us)                        Min                Mean             StdDev              Median            Rounds
-------------------------------------------------------------------------------------------------------------------------
test_bench[trsm_Ac_Bc-512-128]      127.4580 (1.0)      130.8249 (1.0)       4.7251 (2.58)     129.8329 (1.0)          20
test_bench[trsm_Ac_Bf-512-128]      152.6251 (1.20)     156.0478 (1.19)      5.5057 (3.01)     153.4999 (1.18)         20
test_bench[trsm_Af_Bc-512-128]      130.7498 (1.03)     138.0167 (1.05)      8.5175 (4.65)     135.6045 (1.04)         20
test_bench[trsm_Af_Bf-512-128]      159.6250 (1.25)     165.1063 (1.26)      7.4495 (4.07)     162.5626 (1.25)         20
test_bench[trtrs_Ac_Bc-512-128]     333.6659 (2.62)     349.4582 (2.67)     13.2918 (7.26)     346.8744 (2.67)         20
test_bench[trtrs_Ac_Bf-512-128]     157.3330 (1.23)     159.3167 (1.22)      1.8299 (1.0)      158.6045 (1.22)         20
test_bench[trtrs_Af_Bc-512-128]     328.4160 (2.58)     342.3771 (2.62)     11.2343 (6.14)     339.2086 (2.61)         20
test_bench[trtrs_Af_Bf-512-128]     157.1248 (1.23)     159.6709 (1.22)      2.3559 (1.29)     158.0212 (1.22)         20
-------------------------------------------------------------------------------------------------------------------------

ricardoV94 · 2026-08-18T05:45:48Z

So let's go with that and drop the other. It sounds like trsm is what trtrs uses internally anyway.

The zero check is also something we may want to over optimize later down the road when we know that it can't have one by construction (if that's a thing).

It also exposes a scalar alpha parameter. Not sure where that shows up in our graphs but another potential benefit.

ricardoV94 · 2026-08-18T05:50:42Z

+        self, lower: bool, trans: int, A_order: Literal["C", "F"]
+    ):
+        # test_solve_triangular covers trans=0 through the graph, which is all the op ever passes.
+        # The overload is called directly to reach trans=1 and 2, with a complex dtype so that a


why add the test now? presumably if we find a need later then we can add such tests? for now you could just hardcode trans and mention in comment that there's no test for other cases?

Reference LAPACK's trtrs is a zero-pivot scan followed by a side="L" trsm, so it buys nothing over calling trsm directly once we do that scan ourselves, which the side="R" path already had to. Benchmarked equal within one sigma on an f-contiguous rhs across four shapes.

jessegrabowski · 2026-08-19T06:52:32Z

I removed trtrs. The pr message has been updated accordingly.

ricardoV94 · 2026-08-20T12:38:13Z

+        np.testing.assert_allclose(result, expected, rtol=rtol)
+
+    @pytest.mark.parametrize("b_shape", [(5, 3), (5,)], ids=["b_matrix", "b_vec"])
+    def test_solve_triangular_singular_returns_nan(self, b_shape: tuple[int, ...]):


should be covered by a test using pytensor graphs, if there's none already

Sure. This one does matter because trsm doesn't return info, so we're responsible for detecting failure now.

ricardoV94

Looks good just nit with tests

It was allocated per call, which is 19ns and a fifth of the runtime when the systems are small enough for the batch to be dominated by per-call overhead.

ricardoV94 · 2026-08-23T08:00:15Z

Ok if I squash merge?

jessegrabowski added 3 commits August 16, 2026 23:44

Scale numba solver test tolerances to the dtype under test

95b49e7

Build test_empty inputs with the dtype of the graph they feed

16c8edc

jessegrabowski requested a review from ricardoV94 August 17, 2026 05:11

jessegrabowski added numba performance linalg Linear algebra labels Aug 17, 2026

ricardoV94 reviewed Aug 18, 2026

View reviewed changes

jessegrabowski added 2 commits August 19, 2026 01:33

Drop solve_triangular docstring reference to a nonexistent b_ndim arg

c15c023

jessegrabowski changed the title ~~Numba: solve a c-contiguous rhs with trsm instead of copying for trtrs~~ Numba: solve triangular systems with trsm, picking side by rhs layout Aug 19, 2026

Numba: drop the now-unused trtrs binding

2afbe2c

ricardoV94 reviewed Aug 20, 2026

View reviewed changes

Comment thread tests/link/numba/linalg/test_solvers.py Outdated

ricardoV94 reviewed Aug 20, 2026

View reviewed changes

Comment thread tests/link/numba/linalg/test_solvers.py Outdated

ricardoV94 reviewed Aug 20, 2026

View reviewed changes

jessegrabowski added 4 commits August 20, 2026 22:35

Numba: exercise solve_triangular edge cases through graphs

5c8e56d

Numba: build the trsm alpha scalar once per specialization

0ef44dc

It was allocated per call, which is 19ns and a fifth of the runtime when the systems are small enough for the batch to be dominated by per-call overhead.

Numba: scan for a zero pivot before copying A to Fortran order

0344b12

Numba: drop the unit_diagonal zero-diagonal solve_triangular test

fb10af6

jessegrabowski requested a review from ricardoV94 August 22, 2026 20:21

ricardoV94 approved these changes Aug 23, 2026

View reviewed changes

ricardoV94 changed the title ~~Numba: solve triangular systems with trsm, picking side by rhs layout~~ Numba: solve triangular C-contiguous B without copy Aug 23, 2026

Navigation Menu

Uh oh!

Numba: solve triangular C-contiguous B without copy - #2359

Numba: solve triangular C-contiguous B without copy#2359
jessegrabowski wants to merge 10 commits into
pymc-devs:mainfrom
jessegrabowski:numba_trsm_dispatch

jessegrabowski commented Aug 17, 2026 •

edited

Loading

Uh oh!

ricardoV94 commented Aug 17, 2026

Uh oh!

jessegrabowski commented Aug 18, 2026

Uh oh!

ricardoV94 commented Aug 18, 2026 •

edited

Loading

Uh oh!

ricardoV94 Aug 18, 2026

Uh oh!

jessegrabowski commented Aug 19, 2026

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Aug 20, 2026

Uh oh!

jessegrabowski Aug 20, 2026

Uh oh!

ricardoV94 left a comment

Uh oh!

ricardoV94 commented Aug 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

jessegrabowski commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Aug 17, 2026

Uh oh!

jessegrabowski commented Aug 18, 2026

Uh oh!

ricardoV94 commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Aug 19, 2026

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Aug 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jessegrabowski commented Aug 17, 2026 •

edited

Loading

ricardoV94 commented Aug 18, 2026 •

edited

Loading