Implement BandedDot Op #1416
Implement BandedDot Op #1416
Conversation
jessegrabowski added enhancement 
jessegrabowski force-pushed the banded-dot branch from 912a88d to 2282161 Compare 
| I added trust_input and I also load the BLAS functions once on import and save them. So that should reduce some of the most obvious sources of python overhead. New benchmarks (note that they're in ns now, not us): |
pytensor/tensor/slinalg.py Outdated
| A = np.asarray(A) | ||
| m, n = A.shape | ||
| ab = np.zeros((kl + ku + 1, n), dtype=A.dtype, order="C") | ||
| | ||
| for i, k in enumerate(range(ku, -kl - 1, -1)): | ||
| padding = (k, 0) if k >= 0 else (0, -k) | ||
| diag = np.pad(np.diag(A, k=k), padding) | ||
| ab[i, :] = diag | ||
| | ||
| return ab |
There was a problem hiding this comment.
I imagine this explains most of the python overhead for small cases?
There was a problem hiding this comment.
one way or another we have to do that though as part of the cost of the Op. Unless we demand users have inputs ready in that form.
There was a problem hiding this comment.
Yeah it's fine, I was just thinking out loud.
There was a problem hiding this comment.
This rearrangement could be done symbolically in a wrapper Op that calls the blas Op (which expects things to be ready in the correct form)
It might also be better to do smart column indexing on ab instead of using pad
There was a problem hiding this comment.
Yeah it's similar to the Solve, in that you can also do it once and reuse many times possibly, but I think that's too much micro-optimization for now. We also don't want to autodiff through it
pytensor/tensor/slinalg.py Outdated
| _dgbmv = scipy_linalg.get_blas_funcs("gbmv", dtype="float64") | ||
| _sgbmv = scipy_linalg.get_blas_funcs("gbmv", dtype="float32") |
There was a problem hiding this comment.
This will cause import time overhead to PyTensor.
I'm okay paying the extra 3us at runtime instead since virtually nobody will ever use this (or use it in a case where they need those extra us)
There was a problem hiding this comment.
I thought about this as well. It won't stay in the final verison.
There was a problem hiding this comment.
You can exploit prepare_node and add the function to node.tag, which the perform method can then retrieve from. That's two attribute accesses instead of a string check / scipy caching...
There was a problem hiding this comment.
Or you can sidestep perform and use make_thunk instead
| I think the Op is fine, specially if we are not trying to introduce it automatically via rewrites. If we are we may consider the backend (once we have it in numba I suspect it will win for smaller matrices) and/or static shapes if we think the worse-case penalty is still too big |
| Benchmark after tuning up the |
| That looks much better! |
| I agree numba will probably be better across the board. I'd really like this Op to win on the 100x100 case, that's already a pretty big matrix. 1000x1000 and 10,000x10,000 doesn't really show up in nature too often |
| 100x100 is 1us, you are at the edge of python overhead there. Calling an identity PyTensor function and no trust_input is 300-500ns. Calling np.zeros is like 100-200ns. That means you would basically need to have no python overhead whatsoever Edit: those are on my machine, don't know about yours |
| This is the best I think we can get out of this in python? def make_thunk(self, node, storage_map, compute_map, no_recycling, impl): kl = self.lower_diags ku = self.upper_diags if node.outputs[0].dtype == "float64": gbmv = scipy_linalg.get_blas_funcs("gbmv", dtype="float64") else: gbmv = scipy_linalg.get_blas_funcs("gbmv", dtype="float32") ab_size = kl + ku + 1 a_storage = storage_map[node.inputs[0]] b_storage = storage_map[node.inputs[1]] out_storage = storage_map[node.outputs[0]] out_computed = compute_map[node.outputs[0]] if compute_map is not None else [False] def thunk( a_storage=a_storage, b_storage=b_storage, out_storage=out_storage, out_computed=out_computed, kl=kl, ku=ku, ab_size=ab_size, gbmv=gbmv, ): A = a_storage[0] b = b_storage[0] m, n = A.shape ab = np.zeros((ab_size, n), dtype=A.dtype, order="C") for i, k in enumerate(range(ku, -kl - 1, -1)): if k > 0: ab[i, k:] = diag(A, k=k) else: ab[i, :n + k] = diag(A, k=k) out_storage[0] = gbmv(m, n, kl, ku, 1, ab, b) out_computed[0] = True return thunk |
pytensor/tensor/slinalg.py Outdated
| A = as_tensor_variable(A) | ||
| B = as_tensor_variable(b) | ||
| | ||
| out_dtype = pytensor.scalar.upcast(A.dtype, B.dtype) |
There was a problem hiding this comment.
I suspect this is wrong for integer types
| I'm not saying we should do that, but it gives you a lower bound on what to expect from your micro-optimizations |
| Here's what the thunk version benchmarks as for me: I'm curious if it's possible to destroy A and make it into A_banded in-place. If it's possible, it doesn't seem trivial. BLAS doesn't have an Frankly my time would be better served thinking about how to do this in C at this point. |
| Also we should probably be benchmarking this against sparse_dot -- this all might be a waste of time? |
Well SparseDot doesn't work with batch inputs, but I'm curious. Also I don't think the code is too complex or performing too bad. I don't agree with your sentiment, should be thinking of a C impl. A numba one is more interesting... |
| Por que não os dois? Seriously though my feeling is that if we're putting this stuff into a PyMC model the code has to be ultra-performant. It's going to be called umptillion times, the inner-loop of a PDE solver times the MCMC loop. I'll work on the numba dispatch next at any rate |
| By that argument you can't really add any specialized Op that doesn't have a C implementation (unless it's replacing an Op that also doesn't have C implementation). Ignoring the general user, you can have code to decide whether to use this Op or not based on the size (or a rewrite). Also how are you sampling / getting A, can you avoid the boxing/unboxing of the diagonals? |
| well the point is the specialization isn't adding anything over good ol' |
jessegrabowski force-pushed the banded-dot branch from 497721e to 976422f Compare jessegrabowski force-pushed the banded-dot branch from 60d90c4 to c9cc8ce Compare | I just pushed a major refactor to this PR, which:
Regarding point (4), here are the benchmarks using the numba GEMV overload vs what we current get with It's about the same or maybe slightly better, but at the cost that we can't cache the compiled function anymore due to the function pointer. Also note that the test is very sensitive to the detection of the alpha parameter. I had to write: In order for the GEMV rewrite to correctly find |
jessegrabowski force-pushed the banded-dot branch from c9cc8ce to 976fd5b Compare 2 participants
Description
This PR adds a
BandedDotOpthat usesgbmvto do matrix-vector multiplication for the case that A is a banded matrix.In my testing, I found that this case sped up computation significantly. Benchmarking against Pytensor's dot, however, the current implementation is significantly slower:
I guess there's some major overhead from doing the diagonal extractions and looking up the blas function in python? This could and should probably be a C Op, but I'm not sure I have time to realistically dig into all that anytime soon. Help wanted, at any rate.
Related Issue
linalg.BandedDot#1415Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1416.org.readthedocs.build/en/1416/