Dot and blas slowed by negative strides

[TeemuSailynoja]

Context for the issue:

Currently, cumulative sum is implemented as a wrapper for the corresponding Numpy function.
When testing with vectors, instead of using

pt.cumsum(x)

using

pt.dot(pt.tril(pt.ones((d,d))), x)

where d is the length of vector x, seems to lead to considerably faster sampling performance.

See this gist for a quick demo.

Proposal:

Keep the API un changed, but change the internals to compute the dot product with the lower diagonal matrix for the dimension specified in the axis argument of cumsum.

[aseyboldt]

Thanks for investigating this.
If the matrix multiplication is really faster, than I think something strange must be going on with cumsum.
cumsum should be an $O(n)$ operation, while the matrix multiplication is $O(n^2)$ both in memory and number of floating point operations, so if both are implemented properly, the matrix multiplication should pretty much never beat the cumsum implementation.

[TeemuSailynoja]

It does sound odd indeed. In the gist I linked, it's not even close for larger number of summands. I could still try with different backends. At least this looks like something that needs a deeper look into (or I'm missing something fairly obvious). I don't know, if Pytensor is able to optimize the matrix multiplication version somehow differently to the Numpy function.

[ricardoV94]

I suggest benchmarking a direct PyTensor function of cumsum (and the gradient), not going through PyMC.

[ricardoV94]

With n=1000, dot is much slower on my machine:

import pytensor
import pytensor.tensor as pt
import numpy as np

x = pt.vector("x", shape=(1000,))
out = pt.cumsum(x)
d = x.shape[0]
out2 = pt.dot(pt.tril(pt.ones((d,d))), x)

fn = pytensor.function([x], [out, pt.grad(out.sum(), x)], trust_input=True)
fn2 = pytensor.function([x], [out2, pt.grad(out2.sum(), x)], trust_input=True)
x_test = np.random.normal(size=x.type.shape)
np.testing.assert_allclose(fn(x_test)[1], fn2(x_test)[1])
%timeit fn(x_test)  # 7.44 μs ± 178 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
%timeit fn2(x_test)  # 356 μs ± 21.3 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

With n=100, it's kind of a wash

[TeemuSailynoja]

Thank you. This issue should definitely be moved to a discussion, where I investigate what is happening on my machine. Perhaps that could be useful for posteriority.

[ricardoV94]

I think I found the culprit, the CGemV operation seems to do very badly with negative strides on the x input. Example:

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.tensor.blas_c import CGemv

alloc = pt.empty((100,))
x = pt.matrix("x", shape=(100, 1000))
y = pt.vector("y", shape=(1000,))
out = CGemv(inplace=True)(
    alloc,
    1.0,
    x,
    y,
    0.0
)
fn = pytensor.function([x, y], out, accept_inplace=True, trust_input=True)

x_test = np.zeros(x.type.shape[::-1]).T
print(x_test.strides)  # (8, 800)
y_test = np.zeros(y.type.shape)
%timeit fn(x_test, y_test)  # 32.1 μs ± 3.33 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

x_flipped = x_test[::-1]
print(x_flipped.strides)  # (-8, 800)
%timeit fn(x_flipped, y_test)  # 92.2 μs ± 4.34 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

The gradient of CumOp introduces the negative strides when it flips the incoming gradient. My example above didn't have this second dot step

[ricardoV94]

Numpy dot also seems to have this problem, so maybe it's an issue that goes all the way down to blas itself:

import numpy as np

y_test = np.zeros(1000)

x_test = np.zeros((1000, 100)).T
print(x_test.strides)  # (8, 800)

x_test_flipped = x_test[::-1]
print(x_test_flipped.strides)  # (-8, 800)

%timeit np.dot(x_test, y_test)  # 27.1 μs ± 7.56 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit np.dot(x_test_flipped, y_test)  # 86.3 μs ± 828 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[ricardoV94]

#1389 should help quite a bit. It brings down the difference in the PyMC model, although the dot variant still does slightly better. Could you test and confirm on your end @TeemuSailynoja?

[TeemuSailynoja]

Straight out of the box, with just updating pytensor to the latest commit, there is some change, but the difference is still very noticeable for the benefit of the dot product.
My test setup wasn't very rigorous, so I should probably make it more reproducible for a proper comparison.
https://gist.github.com/TeemuSailynoja/9ab551ab1ec9f95bedcaa1052a216b3d
In the merged PR the change looks quite substantial though!

[ricardoV94]

You sure it's picking up the right environment? For me it brought the heavy case within 10s difference