Add BLAS dispatches and Dot-based lowering for the JIT backend - #2406
jessegrabowski wants to merge 21 commits into
Conversation
jessegrabowski
changed the title
97bf342 to
5ac9cb9
Compare
| return batched_dot | ||
|
|
||
|
|
||
| @jax_funcify.register(Gemm) |
There was a problem hiding this comment.
this is an argument to not bother with these ops in jax, like we don't bother with Fusion/Inplace?
There was a problem hiding this comment.
This feels a lot more straight-forward though. We canonicalize to one form them represent it however the backend is able to. In this case it's just a naive form.
There was a problem hiding this comment.
This is not a canonicalization, it's a specialization. It doesn't make sense to specialize (and slow compilation) if we're throwing it away next.
With the very same argument you'd say do fusion and inplace rewrites in jax
There was a problem hiding this comment.
Went with your way
| return expm, cache_version | ||
|
|
||
|
|
||
| def _gemm(A, B, C, transa=False, transb=False, alpha=1.0, beta=0.0): |
There was a problem hiding this comment.
there was a big slowdown in gemv(?) C-code with negative strides, where a copy could be avoided. we may want to do the same trick for numba
|
|
||
|
|
||
| @overload(_ger) | ||
| def _ger_impl(alpha, x, y, A): |
There was a problem hiding this comment.
isn't ger the more useless one, compared to gemv?
There was a problem hiding this comment.
I was getting good results from ger and gemv across the board, but gemm was pretty useless unless the matrices were specific shapes.
There was a problem hiding this comment.
I'll post some real benchmarks.
| b = beta.item() | ||
| if b == 1.0: | ||
| out += Z | ||
| elif b != 0.0: | ||
| out += b * Z |
There was a problem hiding this comment.
isn't this defeat the point of the scalar/mul fusion of gemm?
| def dot(x, y, out=None): | ||
| if out is None: | ||
| out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) | ||
| return _gemm(x, y, out, False, False, 1.0, 0.0) |
There was a problem hiding this comment.
I'd be surprised if numba doesn't emit blas for np.dot already
There was a problem hiding this comment.
it does, but it doesn't have support for inplace, alpha, or beta.
There was a problem hiding this comment.
np.dot(out=x)?
alpha and beta I suspect are just cuteness. We probably make better use with fusion/reduction downstream than trying very hard to merge with blas
| if numba_dot_dtype in _GEMM_DTYPES: | ||
| # `gemm` reads each operand's memory order as a transpose flag, so an | ||
| # operand that reaches here transposed costs nothing, where `np.dot` would | ||
| # have to be handed a contiguous copy of it. |
| return destroy_dependencies | ||
|
|
||
|
|
||
| def get_static_scalar(node: Apply | None, input_index: int) -> float | None: |
There was a problem hiding this comment.
I don't think this helper earns its keep, the logic to make it general is more complex than inlining it
There was a problem hiding this comment.
Kept, five callers
| ) | ||
|
|
||
|
|
||
| @node_rewriter([AllocEmpty]) |
There was a problem hiding this comment.
should apply to alloc of zeros as well
| destroyers = dict.fromkeys( | ||
| client | ||
| for client, input_index in clients[1:] | ||
| if not isinstance(client.op, Output) |
There was a problem hiding this comment.
the refactor and bootstrap code for the c impl dispatcher looks neat.
OTOH I'm concerned IFF we are making the blas pipeline run by default in jax and numba? For jax it reads just like rewrite overhead, since we end up emitting the naive code. For numba I'd need a more exhaustive reproducible benchmark than "it speeds up pytensor-ml by 10%", as this is a fundamental change that touches most graphs we work with. Or proving that np.dot always lowers to blas by numba anyway and we are just skipping some indirection (which ones?).
| def ger(A, alpha, x, y): | ||
| # `A` is only broadcast against the outer product, so the buffer the update | ||
| # writes into takes the product's shape rather than `A`'s. Copying also leaves | ||
| # `A` intact, which is the whole difference between this op and its inplace form. | ||
| out = np.empty((x.shape[0], y.shape[0]), dtype=dtype) | ||
| out[:] = A | ||
| return _ger(alpha.item(), x, y, out) |
There was a problem hiding this comment.
Should this be return _ger(alpha.item(), x, y, A.copy())? or copy_asfortran_order?
|
Evaluation time on the numba backend, mean +/- stddev of a pytest-benchmark run, A bare Gemm
Every batched row wins, 1.4x to 3.0x, with a median of 1.8x on forward and about 2.1x on backward and forward+backward. 46 of the 54 batched cells clear one stddev. The unbatched Ger
The unbatched Nothing loses. Batched forward is unchanged Gemv
Batched backward wins by the most of anything here, 1.5x to 17x with a median of 4.1x and all nine cells significant, and batched forward+backward follows at 1.7x to 12x, median 2.9x. The largest ratios are at square-512, where main spends 3.6 ms on a 2 MB rank-1 update. Batched forward is a milder 1.0x to 1.9x. Nothing loses. The unbatched SGD step on a two-layer MLP
Neutral, 0.97x to 1.04x, all inside one stddev. The fused inplace |
757f97e to
bb78ad3
Compare
|
I'm running an adversarial check on those claims. I am highly skeptical of the "when batched we now avoid copy" and all that arguing because we were making use of |
|
Bot analysis and guided reply, It argues against BLAS pipeline on the pytensor graph vs just giving the right hints? Wanna confirm these results on your side @jessegrabowski and push back again? I reproduced the large batched-matvec backward speedups. My main takeaways are that ordinary GEMM is roughly a wash, the GEMV workload benefits mainly from avoiding input copies, and GER should be dropped from the Numba backend (I was on the fence on C, I think also there) in favor of fused Elemwise outer products. These checks use full GEMM: ordinary dense matrix-matrix accumulation showed no substantial improvement over the previous code: approximately 4.52 ms before versus 4.53 ms with the PR for the tested 512² case. The large regression below is specifically an outer-product case. GEMV: the batched backward gains are real, but a layout guard around the existing Numba infers those transposed matrix slices as arbitrary-layout arrays and copies them before calling BLAS, even though each slice is F-contiguous at runtime. I verified this in the actual compiled function's types and LLVM. Adding The guard inside Blockwise's core Dot lowering looks like this ( if a.flags.f_contiguous:
np.dot(np.asfortranarray(a), b, out)
else:
np.dot(np.ascontiguousarray(a), b, out)The helpers give Numba a C/F-typed operand; checking the flags alone leaves its inferred type unchanged and still triggers the copy. For the full backward function with eight 512² matrices, one interleaved run gives:
Guarded dot matches PR Gemv at 128² and beats it at 512² in these checks. Blockwise(Gemv) also retains the gains when starting from zeros. Unbatched Gemv itself was roughly neutral. BatchedDot is therefore unnecessary for retaining these gains and can be removed in a follow-up PR, provided its Blockwise replacement handles input layout. GER: I would remove this specialization and lower outer-product updates to fused broadcast multiplication, including alpha/beta. For a 512² update with the accumulator unchanged, the previous path takes 100.2 µs, PR Ger 44.6 µs, and fused Elemwise 43.8 µs. The non-inplace Ger implementation is itself a fused loop. Actual inplace BLAS GER also loses to inplace Elemwise here: 31.6 versus 28.7 µs. This also addresses the major regression: Since Numba's |
|
I reproduced your timings. So the PR is much less of a win than I was hoping. What the guard can't do is a slice that's neither C nor F.
Inplace accumulation My preference is to merge it as-is because that's less work for me and it seems like a push at worst with maybe some fringe benefits. The machinery is definitely cleaner. I'm open to making changes if you insist. |
Introduce the singledispatch c_funcify registry returning detached CImpl implementations, resolve CLinker through it, and route OpWiseCLinker, the VM, and DebugMode to the dispatched C thunk with a Python fallback.
The deleted make_c_gemv_destructive also duplicated a shared AllocEmpty so each Gemv could destroy its own buffer, so test_multiple_inplace fails until the generic replacement lands two commits later.
Without this the rewrite that introduces Ger regresses against the elemwise it replaces, since numba falls back to object mode for it.
Copying the accumulator in and letting BLAS scale it on top touches the output twice, which cost more than the elemwise dot-and-add these ops replace.
1b05695 to
471bb1a
Compare
|
blas can handle strided inputs? Is that general across Blas implementations? I was under the impression from the C work that it must be contiguous and the only clever thing we could do is handle negative [::-1] strides. On my end we can go ahead but I'd not include any blas rewrites in numba vs just having the better dot/batched dot dispatch. |
|
One other advantage of our own dot dispatch is that we avoid the numba spurious contiguity warning? |
2 participants
All of this measures netural against pymc-model-catalog. I saw 10-15% speedup in the backward pass of linear layers in pytensor-ml. Rewriting graphs of the form
a + B @ Cinto GEMM isn't perfect, it still requires the user to put parenthesis. We can try to tune it up in follow-up work, but it's not super clear to me it's an obvious win.