Cache-friendly bit-order reversal about plonky2 HOT 2 CLOSED

def _transpose(x):
    n, n_, _ = x.shape
    assert n == n_
    assert n == 1 << n.bit_length() - 1

    if n == 1:
        return
    assert n & 1 == 0

    _transpose(x[:n>>1, :n>>1])
    _transpose_and_swap(x[:n>>1, n>>1:], x[n>>1:, :n>>1])
    _transpose(x[n>>1:, n>>1:])


def _transpose_and_swap(x, y):
    assert x.shape == y.shape
    n, n_, _ = x.shape
    assert n == n_
    assert n == 1 << n.bit_length() - 1

    if n == 1:
        x[...], y[...] = y.copy(), x.copy()
        return

    _transpose_and_swap(x[:n>>1, :n>>1], y[:n>>1, :n>>1])
    _transpose_and_swap(x[:n>>1, n>>1:], y[n>>1:, :n>>1])
    _transpose_and_swap(x[n>>1:, :n>>1], y[:n>>1, n>>1:])
    _transpose_and_swap(x[n>>1:, n>>1:], y[n>>1:, n>>1:])


def bit_reverse(x):
    n, = x.shape
    lb_n = n.bit_length() - 1
    assert n == 1 << lb_n

    if lb_n <= 1:
        return

    half_lb_n = lb_n >> 1
    leftover = lb_n & 1

    x = x.reshape(1 << half_lb_n, 1 << half_lb_n + leftover)
    for row in x:
        bit_reverse(row)

    _transpose(x.reshape(1 << half_lb_n, 1 << half_lb_n, 1 << leftover))

    for row in x:
        bit_reverse(row)