A functional implementation of Fast Walsh Hadamard Transform using JAX

[shailesh1729]

Hi,

I have come up with a JAX-based functional implementation of Walsh Hadamard Transform which works nicely with @jit. However, it involves three levels of lax.while_loop. I wanted to show the implementation and ask for comments if there is a way to express the logic in a better and more efficient manner.

Code

First is a reference version which I wrote. It uses for, while loops and is slower. Took me a day to get the logic right. It may differ from other implementations by a scaling factor.

def fwht_ref(X):
    """Computes the Fast Walsh Hadamard Transform over columns of :math:`X`
    """
    n = X.shape[0]
    # number of stages
    s = (n-1).bit_length()

    def init():
        Y = jnp.empty(X.shape, dtype=X.dtype)
        A  = X[0::2]
        B = X[1::2]
        Y = Y.at[0::2].set(A + B)
        Y = Y.at[1::2].set(A - B)
        return (Y, 1, 2, 4)

    def body(state):
        # gap between x entries
        # number of x entries
        X, count, gap, step = state
        Y = jnp.empty(X.shape, dtype=X.dtype)
        J = 0
        k = 0
        while k < n -1:
            for j in range(J, J+gap-1, 2):
                # compute the four parts
                a = X[j]
                b = X[j+gap]
                c = X[j+1]
                d = X[j+1+gap]
                Y = Y.at[k].set(a+b)
                Y = Y.at[k+1].set(a-b)
                Y = Y.at[k+2].set(c-d)
                Y = Y.at[k+3].set(c+d)
                k += 4
            J += step
        return (Y, count+1, 2*gap, 2*step)

    def cond(state):
        count = state[1]
        return count < s

    state = init()
    while cond(state):
        state = body(state)
    #state = lax.while_loop(cond, body, init())
    return state[0]

After this, I converted it to a functional version:

@jit
def fwht(X):
    """Computes the Fast Walsh Hadamard Transform over columns of :math:`X`
    """
    n = X.shape[0]
    # number of stages
    s = (n-1).bit_length()

    def init1():
        Y = jnp.empty(X.shape, dtype=X.dtype)
        A  = X[0::2]
        B = X[1::2]
        Y = Y.at[0::2].set(A + B)
        Y = Y.at[1::2].set(A - B)
        return (Y, 1, 2, 4)

    def body1(state):
        # gap between x entries
        # number of x entries
        X, count, gap, step = state
        Y = jnp.empty(X.shape, dtype=X.dtype)
        J = 0
        k = 0
        def body2(state):
            Y, J, k = state
            def body3(state):
                Y, j, k = state
                # compute the four parts
                a = X[j]
                b = X[j+gap]
                c = X[j+1]
                d = X[j+1+gap]
                Y = Y.at[k].set(a+b)
                Y = Y.at[k+1].set(a-b)
                Y = Y.at[k+2].set(c-d)
                Y = Y.at[k+3].set(c+d)
                return (Y, j+2, k+4)
            def cond3(state):
                j = state[1]
                return j <  J+gap-1
            # the loop
            init3 = (Y, J, k)
            Y, j, k = lax.while_loop(cond3, body3, init3)
            return (Y, J + step, k)

        def cond2(state):
            k = state[2]
            return k < n - 1

        init2 = Y, J, 0
        Y, J, k = lax.while_loop(cond2, body2, init2)

        return (Y, count+1, 2*gap, 2*step)

    def cond1(state):
        count = state[1]
        return count < s

    state = lax.while_loop(cond1, body1, init1())
    return state[0]

The three levels of lax.while_loop reflect the original logic however, is it possible to write it in a more expressive manner?

Benchmarks

Some code that I ran in Jupyter notebook for benchmarking:

n = 32
I = jnp.eye(n, dtype=jnp.int32)
nonjit_time = %timeit -o crdsp.fwht_ref(I).block_until_ready()
# 278 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
jit_time = %timeit -o crdsp.fwht(I).block_until_ready()
# 24.7 µs ± 298 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
nonjit_time.average / jit_time.average
# 11182.88999174172

The gain of 11000x seems surreal. Normally, I get gains in the range of 50-200x during JIT compilation. Is this expected?

Here is another benchmark comparing a direct matrix-vector product to a transform operation

key = jax.random.PRNGKey(0)
n = 1024
y = jax.random.normal(key, (n,))
F = crdsp.fwht(jnp.eye(n))
mat_time = %timeit -o  (F @ y).block_until_ready()
# 423 µs ± 6.11 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
op_time = %timeit -o  (crdsp.fwht(y)).block_until_ready()
# 63.5 µs ± 498 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
mat_time.average / op_time.average
# 6.668208241947298

The gain for the matrix-vector product is relatively modest (just 6x). As I increase the number of columns in y, the gain reduces further. This seems to indicate that there is some opportunity to improve the JAX implementation.

This implementation is available as part of CR-Sparse.