Proofs: Winograd Fourier transform algorithm

WFTA computes the DFT of length q_1 ⋯ q_d with ∏ m_i − ∏ u_i multiplications and Σ (a_i + c_i) ∏_{j<i} m_j ∏_{j>i} q_j additions full

This proof shows: correctness, complexity.

Setting. $R$ is $\mathbb{C}$ or $\mathbb{Z}/p\mathbb{Z}$, $n = q_1 \cdots q_d$ with pairwise coprime $q_i \in \{2, 3, 4, 5, 7, 8\}$, and $\omega \in R$ has order exactly $n$. Put $\omega_q = \omega^{n/q}$, of order exactly $q$, and $F_q = (\omega_q^{jk})_{0 \le k, j < q}$. In $\mathbb{Z}/p\mathbb{Z}$ an element of order $q$ exists only if $q \mid p - 1$; so $2$ is invertible when some $q_i \in \{3, 5, 7, 8\}$ (then $p \ge 7$), and $3$ is invertible when $7 \mid n$ (then $p \ge 29$). These are the only divisions in the constants. For a root $w$ of order $q$ write $C(e) = (w^e + w^{-e})/2$ and $S(e) = (w^e - w^{-e})/2$.

Statement. WFTA returns $X_k = \sum_{j<n} x_j \omega^{jk}$ for all $0 \le k < n$. It does $\prod_i m_{q_i} - \prod_i u_{q_i}$ multiplications and $\sum_{i=1}^d (a_{q_i} + c_{q_i}) \prod_{j<i} m_{q_j} \prod_{j>i} q_j$ additions and subtractions, with $(m, a, c, u)$ equal to $(2, 2, 0, 2)$, $(3, 3, 3, 1)$, $(4, 6, 2, 3)$, $(6, 8, 9, 1)$, $(9, 17, 19, 1)$ and $(8, 16, 10, 4)$ for $q = 2, 3, 4, 5, 7, 8$.

Lemma 1 (index maps). The map $(j_1, \ldots, j_d) \mapsto j = \sum_i j_i n/q_i \bmod n$ is a bijection from $\prod_i [0, q_i)$ to $[0, n)$, and $k \mapsto (k \bmod q_1, \ldots, k \bmod q_d)$ is a bijection from $[0, n)$ to $\prod_i [0, q_i)$. With $a[j_1, \ldots, j_d] = x_j$ and $k_i = k \bmod q_i$,

Proof. For $l \ne i$, $q_i$ divides $n/q_l$, so $j \equiv j_i\, n/q_i \pmod{q_i}$; since $n/q_i$ is coprime to $q_i$, $j$ determines each $j_i$. The map is injective between sets of size $n$, hence bijective. The second map is bijective by the Chinese remainder theorem. Finally $\omega^{jk} = \prod_i \omega^{(n/q_i) j_i k} = \prod_i \omega_{q_i}^{j_i k}$, because $\omega^n = 1$, and $\omega_{q_i}^{j_i k} = \omega_{q_i}^{j_i k_i}$ because $\omega_{q_i}^{q_i} = 1$. Summing over $j$, that is over all tuples, gives the formula. $\square$

Lemma 2 (pairing). Let $q$ be an odd prime, $w$ of order $q$, $t^\pm_j = x_j \pm x_{q-j}$ for $1 \le j \le (q - 1)/2$. Then for $1 \le k < q$, $X_k = x_0 + \sum_j \bigl(C(jk)\, t^+_j + S(jk)\, t^-_j\bigr)$, and $X_0 = x_0 + \sum_j t^+_j$. Moreover $C$ is even and $S$ odd in $e$, and both depend on $e$ modulo $q$.

Proof. $x_j w^{jk} + x_{q-j} w^{-jk} = C(jk)(x_j + x_{q-j}) + S(jk)(x_j - x_{q-j})$, since $C + S = w^{jk}$ and $C - S = w^{-jk}$. Summing over the pairs $\{j, q - j\}$, which partition $[1, q)$, gives the claim. $\square$

Lemma 3 (cyclic convolution of length 3). For $h, v \in R^3$, let $y_b = \sum_{a<3} h_{(b - a) \bmod 3} v_a$, $H = h_0 + h_1 + h_2$, $g_e = h_e - H/3$, $s = v_0 + v_1 + v_2$, $d_0 = v_0 - v_1$, $d_1 = v_1 - v_2$, $e = v_0 - v_2$, $P_1 = g_0 d_0$, $P_2 = g_1 d_1$, $P_3 = (g_0 + g_1) e$. Then $y = \frac H3 s\,(1, 1, 1) + (P_1 - P_2,\ P_3 - P_1,\ P_2 - P_3)$.

Proof. Convolution is linear in the kernel, and the constant kernel $\frac H3 (1, 1, 1)$ gives $\frac H3 s$ in every position. For $g$, $g_2 = -g_0 - g_1$. Then $g_0 v_0 + g_2 v_1 + g_1 v_2 = g_0 (v_0 - v_1) - g_1 (v_1 - v_2) = P_1 - P_2$; $g_1 v_0 + g_0 v_1 + g_2 v_2 = g_1 (v_0 - v_2) + g_0 (v_1 - v_2) = (g_0 + g_1) e - g_0 d_0 = P_3 - P_1$, using $e = d_0 + d_1$; and the third entry is minus the sum of the first two, because $\sum_b \sum_a g_{b-a} v_a = (\sum g) s = 0$, which is $P_2 - P_3$. $\square$

Lemma 4 (modules). For each $q$ in the list and each $w$ of order $q$, the module of the page (the code of _pre, _post and _mu) computes $F_q x$ with root $w$.

Proof. $q = 2$: $w = -1$, and the module returns $(x_0 + x_1, x_0 - x_1)$.

$q = 4$: $w^2 = -1$ (an element of order $2$ in a field), $w^3 = -w$. Then $X_0 = (x_0 + x_2) + (x_1 + x_3)$, $X_2 = (x_0 + x_2) - (x_1 + x_3)$, $X_1 = (x_0 - x_2) + w (x_1 - x_3)$ and $X_3 = (x_0 - x_2) - w(x_1 - x_3)$, which are the four outputs.

$q = 3$: by Lemma 2 with $t^\pm = x_1 \pm x_2$, $X_1 = x_0 + C(1) t^+ + S(1) t^-$ and $X_2 = x_0 + C(1) t^+ - S(1) t^-$, using $C(2) = C(-1) = C(1)$ and $S(2) = -S(1)$. The module has $m_0 = x_0 + t^+ = X_0$, $m_1 = (C(1) - 1) t^+$, $m_2 = S(1) t^-$, and $m_0 + m_1 = x_0 + C(1) t^+$, so its outputs $m_0$, $m_0 + m_1 \pm m_2$ are right.

$q = 5$: with $t_0, t_1 = x_1 + x_4, x_2 + x_3$, $u_0, u_1 = x_1 - x_4, x_2 - x_3$, and $C_e = C(e)$, $S_e = S(e)$, Lemma 2 and $C_4 = C_1$, $S_4 = -S_1$ give $X_1 = x_0 + C_1 t_0 + C_2 t_1 + S_1 u_0 + S_2 u_1$ and $X_2 = x_0 + C_2 t_0 + C_1 t_1 + S_2 u_0 - S_1 u_1$, and $X_4$, $X_3$ are the same with the signs of the $S$ terms reversed. The module computes $s = t_0 + t_1$, $m_0 = x_0 + s = X_0$, $m_1 = (\frac{C_1 + C_2}2 - 1) s$, $m_2 = \frac{C_1 - C_2}2 (t_0 - t_1)$, $m_3 = S_2 (u_0 + u_1)$, $m_4 = (S_1 - S_2) u_0$, $m_5 = -(S_1 + S_2) u_1$. Then $m_0 + m_1 \pm m_2$ is $x_0 + C_1 t_0 + C_2 t_1$ or $x_0 + C_2 t_0 + C_1 t_1$; $m_3 + m_4 = S_1 u_0 + S_2 u_1$ and $m_3 + m_5 = S_2 u_0 - S_1 u_1$. The outputs $m_0$, $r_0 + i_0$, $r_1 + i_1$, $r_1 - i_1$, $r_0 - i_0$ are $X_0, \ldots, X_4$.

$q = 7$: let $p_j, u_j = x_j \pm x_{7-j}$ for $j = 1, 2, 3$. By Lemma 2, with $C_e$ and $S_e$ taken modulo $7$, the real parts $R_k = x_0 + C_k p_1 + C_{2k} p_2 + C_{3k} p_3$ are

and the imaginary parts $I_k = S_k u_1 + S_{2k} u_2 + S_{3k} u_3$ are $I_1 = S_1 u_1 + S_2 u_2 + S_3 u_3$, $I_3 = S_3 u_1 - S_1 u_2 + S_2 u_3$ and $I_2 = S_2 u_1 - S_3 u_2 - S_1 u_3$, with $X_k = R_k + I_k$ and $X_{7-k} = R_k - I_k$. The convolution $y$ of Lemma 3 with kernel $(C_1, C_3, C_2)$ and input $(p_1, p_2, p_3)$ is $(R_1, R_3, R_2) - x_0$, by expanding $y_b = \sum_a h_{b-a} v_a$. The convolution with kernel $(S_1, -S_3, S_2)$ and input $(u_1, u_2, -u_3)$ is $(I_1, -I_3, I_2)$, by the same expansion. The module's inputs are those of Lemma 3 for both: $s, d_0, d_1, e$ from $(p_1, p_2, p_3)$, and $u_1 + u_2 - u_3$, $u_1 - u_2$, $u_2 + u_3$, $u_1 + u_3$ from $(u_1, u_2, -u_3)$. Its constants are $1$ for $m_0 = x_0 + s = X_0$, $H/3 - 1$ for $s$, so that $m_0 + m_1 = x_0 + \frac H3 s$, then $g_0, g_1, g_0 + g_1$, then $T/3$ and the three values of the second kernel. So $r_b = m_0 + m_1 + w_b = x_0 + y_b$ gives $R_1, R_3, R_2$, and $u_b$ gives $I_1, -I_3, I_2$. The outputs $m_0$, $r_0 + u_0$, $r_2 + u_2$, $r_1 - u_1$, $r_1 + u_1$, $r_2 - u_2$, $r_0 - u_0$ are $X_0, X_1, X_2, X_3, X_4, X_5, X_6$.

$q = 8$: let $J = w^2$, of order $4$, so $J^2 = -1$ and $1 + J \ne 0$; put $c = w/(1 + J)$. Write $s_j = x_j + x_{j+4}$ and $d_j = x_j - x_{j+4}$ for $j < 4$. Since $w^4 = -1$, $X_{2k'} = \sum_{j<4} s_j J^{jk'}$, the DFT of length $4$ of $s$ with root $J$, which the first four products and the $q = 4$ formulas compute. And $X_{2k'+1} = \sum_{j<4} z_j J^{jk'}$ with $z_j = w^j d_j$. Now $w = c(1 + J)$ and $w^3 = wJ = c(J - 1)$, so

using $J^2 = -1$. With $a, b = d_0 \pm J d_2$, $P = m_6 + m_7$ and $Q = m_7 - m_6$, where $m_6 = c(d_1 - d_3)$ and $m_7 = cJ(d_1 + d_3)$, the DFT of length $4$ of $z$ is $(a + P, b + Q, a - P, b - Q)$, which are $X_1, X_3, X_5, X_7$. $\square$

Lemma 5 (axes and Kronecker products). Store an array of shape $(s_1, \ldots, s_d)$ in row-major order. Applying a linear map $L: R^{s_i} \to R^{r}$ to every line along axis $i$ is multiplication by $I_{s_1 \cdots s_{i-1}} \otimes L \otimes I_{s_{i+1} \cdots s_d}$, and multiplying entry $(t_1, \ldots, t_d)$ by $\prod_i \mu_i[t_i]$ is multiplication by $\mathrm{diag}(\mu_1) \otimes \cdots \otimes \mathrm{diag}(\mu_d)$.

Proof. In row-major order the entry $(t_1, \ldots, t_d)$ has position $\sum_i t_i \prod_{l > i} s_l$, which is the indexing of the Kronecker product: the entry of $P \otimes Q$ in row $(r, r')$ and column $(c, c')$ is $P_{rc} Q_{r'c'}$. Applying $L$ along axis $i$ changes $t_i$ by $L$ and keeps the other coordinates, which is the stated product, and the diagonal case is the same statement for the entries. $\square$

Correctness. By Lemma 5, line L3 multiplies $\mathrm{vec}(a)$ by $\prod_{i=d}^{1} (I_{m_1 \cdots m_{i-1}} \otimes A_i \otimes I_{q_{i+1} \cdots q_d})$, which by the mixed-product rule $(P \otimes Q)(P' \otimes Q') = PP' \otimes QQ'$ equals $A_1 \otimes \cdots \otimes A_d$. Line L4 multiplies by $\bigotimes_i \mathrm{diag}(\mu_i)$; skipping the products whose factors are all $1$ does not change any value. Line L5 multiplies by $C_1 \otimes \cdots \otimes C_d$ in the same way, in the reverse order. By the mixed-product rule again, the whole is $\bigotimes_i C_i \mathrm{diag}(\mu_i) A_i = \bigotimes_i F_{q_i}$ by Lemma 4 with $w = \omega_{q_i}$. The entry of $(\bigotimes_i F_{q_i}) \mathrm{vec}(a)$ at $(k_1, \ldots, k_d)$ is $\sum a[j_1, \ldots, j_d] \prod_i \omega_{q_i}^{j_i k_i}$, which is $X_k$ by Lemma 1, and line L6 reads it at $k_i = k \bmod q_i$. The inverse transform is the same computation with $\omega^{-1}$, which has order $n$, followed by division by $n$.

Operation count. A module application costs $a_q$ additions in $A_q$ and $c_q$ in $C_q$, read off from the straight-line code; the table of the page lists them. Applying $A_{q_i}$ along axis $i$ costs $a_{q_i}$ times the number of lines, the product of the other axis lengths. When line L3 treats axis $i$, the axes $j < i$ already have length $m_{q_j}$ and the axes $j > i$ still have length $q_j$. When line L5 treats axis $i$ (in decreasing order), the axes $j > i$ have been shrunk back to $q_j$ and the axes $j < i$ still have length $m_{q_j}$. So both passes over axis $i$ cost $(a_{q_i} + c_{q_i}) \prod_{j<i} m_{q_j} \prod_{j>i} q_j$, which sums to the stated count. Line L4 does one multiplication for each of the $\prod_i m_{q_i}$ tuples except the $\prod_i u_{q_i}$ whose constants are all $1$. Lines L1, L2 and L6 do no arithmetic on the data; the constants and their products depend only on $n$. The index maps and the passes over the array cost $O(d \prod_i m_{q_i})$ steps, and the array has at most $\prod_i m_{q_i}$ entries.

Real and imaginary constants. Over $\mathbb{C}$ with $w = e^{-2\pi i/q}$, $w^{-e} = \overline{w^e}$, so every $C(e)$ is real and every $S(e)$ purely imaginary; so are $H$ (a sum of $C$ values), $T$ (a sum of $S$ values), the $g$ and $h$ values, and the constants $1$. For $q = 8$, $J = -i$, $1 + J = 1 - i$ and $w = (1 - i)/\sqrt 2$, so $c = 1/\sqrt 2$ is real and $cJ$ imaginary. Hence every product $\prod_i \mu_i[t_i]$ is real or purely imaginary, and multiplying a complex number by it costs $2$ real multiplications. $\blacksquare$