Tensor-Product Presentation of NTT and Constant-Geometry NTT

2026-08-03 · 19 min read · ntt-deep-dive #1 of 1

#math #hardware

Most iterative NTT hardware paper eventually hits the same wall: the butterflies are trivial, and the addressing is where all the design effort goes. Write an NTT as a nest of for loops and you get index expressions like j + (i << (k-l)) scattered across the datapath, changing shape at every stage. The loops are correct but they are not analyzable — you cannot look at them and see which memory access pattern falls out.

Tensor products fix this. They let you write the whole transform as a product of structured matrices, where each factor names one concrete hardware resource: a bank of butterflies, a permutation network, a ROM read. Once the algorithm is an algebraic object, restructuring the dataflow becomes matrix manipulation rather than index bookkeeping.

This note builds that machinery from scratch and then uses it to derive the constant-geometry NTT — the variant where the wiring between stages is literally the same permutation every stage. I follow the derivation in Zhang et al. (FCCM 2023), which presents this in an FHE accelerator context, and I flag at the end where the structure is actually much older than that paper.

Everything below I checked numerically (symbolic DFT over $\mathbb{C}$ up to $N=64$, plus integer NTT mod $q$); the figures are generated from the sparsity patterns of the actual matrices, not drawn by hand.

1. The two patterns that matter

For square matrices $A_{M\times M}$ and $B_{L\times L}$, the tensor (Kronecker) product $A_M \otimes B_L$ replaces every entry $a_{ij}$ of $A$ with the block $a_{ij}B$, giving a matrix of order $ML$.

The entire formalism rests on the fact that two special cases mean two different machines. Let $\mathbf{x}$ have length $ML$.

$I_L \otimes A_M$ is block-diagonal with $L$ copies of $A_M$. To see what that demands of memory, you have to say explicitly how $\mathbf{x}$ is cut into the blocks the copies act on. Partition it into $L$ contiguous segments of length $M$:

so segment $p$ starts at offset $pM$ and walks with stride $1$. Under that blocking the product is literally the segments processed one by one:

The index $p$ never appears inside the sum. Segment $p$ is closed under the operator — PE $p$ reads only $\mathbf{x}^{(p)}$, writes only $\mathbf{x}^{(p)}$, and never talks to its neighbours. That is $L$ processors doing the same job on adjacent data, and if you give each PE its own memory bank holding one segment, no two PEs ever address the same bank.

The second pattern is easier to build in the opposite direction: start from the operation you want, and let the matrix fall out of it.

Cut $\mathbf{x}$ the other way — not $L$ segments of length $M$, but $M$ subvectors of length $L$:

Now treat each subvector as though it were a single number, and combine the whole set of them using the coefficients of $A_M$. Output subvector $i$ is the weighted sum of all $M$ input subvectors, the weights being row $i$ of $A$:

There is no matrix product here yet — $(*)$ is just $M$ scalar-times-vector products added together, done $M$ times over. Note the roles: $M$ is the size of the transform, and $L$ is how many vectors ride through it together, the width of the machine.

To turn $(*)$ into a single operator on the stacked $\mathbf{x}$, use the one observation that does all the work: scaling a length-$L$ vector by the scalar $A[i,j]$ is the same as multiplying it by the $L \times L$ matrix $A[i,j]\,I_L$. So one output block is a row of such matrices applied to $\mathbf{x}$,

and stacking all $M$ block rows gives the whole thing:

Now look at what that matrix is. Every entry $A[i,j]$ of $A_M$ has been replaced by the block $A[i,j]\,I_L$ — which is precisely the tensor product defined at the top of this section, taken with $B = I_L$. The operator that forms weighted sums of whole subvectors is the tensor product:

Read by index, $(*)$ says

and the position $p$ is untouched — it sits identically on both sides of the sum. Element $p$ of every subvector is combined only with element $p$ of the others, never with the elements beside it. Collect those matching positions and you get $L$ vectors of length $M$,

one for each $p = 0,\dots,L-1$. Run that relation backwards and it says the same thing from the other side: subvector $\mathbf{u}^{(j)}$ collects the $j$-th element of each of the $L$ vectors — the first subvector is built from the first element of every $\mathbf{v}_p$, the second from every second element, and so on.

In those terms the whole operation is just $A_M\mathbf{v}_p = \mathbf{v}'_p$, one vector at a time. Concretely at $p=0$: gather the first element of each of the $M$ subvectors, multiply that length-$M$ vector by each row of $A_M$ in turn, and the $M$ numbers that come out are the first elements of the $M$ output subvectors. Then repeat for $p=1,\dots,L-1$. So $L$ separate $M$-point operations — one per vector, each reading with stride $L$.

That is the whole difference. Both forms are $L$ independent copies of $A_M$; the only question is which $M$ elements one copy sees — offset $pM$ with stride $1$ above, offset $p$ with stride $L$ here.

The figure makes the point better than the algebra does. Both panels contain exactly $L$ copies of $A_M$ and do exactly the same arithmetic. What changes is the access pattern: contiguous runs on the left, interleaved vectors on the right. Same operator $A$, same data, completely different memory behaviour — and on hardware that difference is a conflict-free bank read versus a strided gather. The whole game is choosing which form you want at each stage.

One more identity we need for scalability. If you have fewer than $L$ processors, split the identity:

Read right to left: the inner $I_{L_2}\otimes A_M$ is a chunk of work, and you hand $L_1$ of those chunks to $L_1$ physical processors. This single line is why the resulting architecture scales to any PE count — more on that in §7.

2. Stride permutation

The data movement operator is the stride permutation $P(N,M)$, with $N = ML$. Conventions differ between authors, so let me pin mine down explicitly:

In words: read $\mathbf{x}$ with stride $M$, $L$ elements at a time. Output block $i$ is $(x_i, x_{M+i}, x_{2M+i},\dots)$. For $N=8, M=2$ this is the perfect shuffle — evens then odds:

Three properties carry the derivation.

Factorization. A stride splits into two strides:

Inverse. Taking $M_1 = M$, $M_2 = N/M$ and noting $P(N,N)=I$:

so strides $M$ and $N/M$ undo each other. This is what makes the twiddle matrix stay diagonal in §5.

Commutation theorem. This is the bridge between the two processing styles:

Parallel form and vector form compute the same thing up to a stride permutation. Any algorithm expressed in tensor products can therefore be retargeted — pushed toward parallel or toward vector — by moving permutations around. That is the lever we pull in §4.

3. NTT as a matrix factorization

The NTT of length $N$ over $\mathbb{Z}_q$ maps coefficients $(x_0,\dots,x_{N-1})$ by

with $\omega_0$ a primitive $N$-th root of unity mod $q$ (so $\omega_0^N \equiv 1$). Collecting this into a matrix, with $k=\log_2 N$:

It is the DFT with $e^{-2\pi i/N}$ swapped for $\omega_0$ — same algebra, no rounding error. So every Cooley–Tukey-style factorization transfers verbatim; only the coefficient ring changes.

For the twiddles, define the diagonal matrix

where $\omega$ is a primitive $S$-th root of unity. The general Cooley–Tukey construction stacks $S/R$ powers of $D_R(S)$ down the diagonal, but radix 2 never needs that generality. The only twiddle matrix (5) below ever instantiates is $T_{2^{k-l}}(2^{k-l+1})$, and $2^{k-l+1} = 2\cdot 2^{k-l}$, so $S = 2R$ at every stage and there are always exactly two blocks:

That is just the twiddle structure of a decimation-in-time stage written out: the first half of the coefficients pass through untouched, the second half pick up $\omega^0,\dots,\omega^{R-1}$.

The decimation-in-time factorization into radix-2 butterflies $F(2) = \left[\begin{smallmatrix}1&1\\1&-1\end{smallmatrix}\right]$ then reads

with $Q(2^k)$ the bit-reversal permutation applied to the input. Write $Z_0(l) = I_{2^{l-1}} \otimes T_{2^{k-l}}(2^{k-l+1})$ for the stage-$l$ twiddle matrix.

The product runs $l=1\ldots k$ left to right, so applied to a vector the stage $l=k$ executes first. Keep this in mind — it is why the figures below run $l=3,2,1$.

A worked example: $N = 8$

Take $k = 3$, so three stages. Everything below stays in terms of $\omega_0$ alone — a primitive $8$-th root of unity mod $q$ — so none of it depends on which $q$ you pick.

At stage $l$ we have $R = 2^{k-l}$, $S = 2R$, and the order-$S$ root is $\omega = \omega_0^{2^{l-1}}$. Writing out (6) and (7) stage by stage:

Two things fall straight out. The stage that executes first has $R = 1$, so $T_1(2) = I_2$ — no twiddles at all, which is the familiar fact that the opening decimation-in-time stage is pure add/subtract. And the twiddle set doubles every stage until the last one carries the full $\omega_0^0,\dots,\omega_0^3$.

The butterfly operator changes shape just as fast. $I_{2^{l-1}} \otimes F(2) \otimes I_{2^{k-l}}$ pairs coefficients at distance $2^{k-l}$:

Bit reversal $Q(8)$ first puts the input in the order

and because stage $l=3$ carries no twiddles, it is nothing but four 2-point transforms on adjacent pairs:

Stages $l=2$ and $l=1$ then merge those four 2-point transforms into two 4-point ones and finally a single 8-point one, picking up $\omega_0^2$ and then $\omega_0,\omega_0^2,\omega_0^3$ along the way. That last column of the second table — stride $1, 2, 4$ — is the whole difficulty, and it is what the next section has to get rid of.

Now look at the operator in the middle:

This is neither of our two clean patterns. It is a mixed three-factor product: $2^{l-1}$ independent copies of the vector operation $F(2)\otimes I_{2^{k-l}}$. Worse, both exponents move with $l$. At $l=1$ it is one wide vector operation; at $l=k$ it is $N/2$ independent scalar butterflies. A fixed piece of hardware cannot be both.

4. Flattening to pure parallelism

Getting rid of the mixed form is exactly what stride permutations are for. The three-factor identity is

with $N = N_1N_2N_3$. The mixed operator becomes a pure parallel operator sandwiched between two permutations. Substituting into (5) with $N_1 = 2^{l-1}$, $N_2 = 2$, $N_3 = 2^{k-l}$:

This is real progress: the compute operator $I_{2^{k-1}} \otimes F(2)$ is now the same at every stage — $N/2$ independent radix-2 butterflies, the ideal parallel form.

But the addressing got worse. Each stage now carries two permutations, and both depend on $l$:

Three stages, three different wiring patterns, two readdressing passes each. In hardware that means either a general-purpose interconnect (expensive, and it degrades badly as $N$ and PE count grow) or reconfiguring the address generation per stage. This is the memory-access bottleneck that dominates scalable NTT designs.

5. Deriving constant geometry

Here is the trick. Unfold (9) along $l$ and look at the boundaries between adjacent stages:

The leading permutation of stage $l$ is $P(2^k,2^l)$. Split it with rule (2):

Only $P(2^k,2)$ — a plain perfect shuffle, no $l$ in it — stays with stage $l$. The leftover $P(2^k,2^{l-1})$ is pushed left, out of stage $l$ and onto the tail of stage $l-1$. Do this at every boundary and the shifted factors telescope:

Every stage now has an identical prefix $P(2^k,2)(I_{2^{k-1}}\otimes F(2))$, and the $l$-dependence has been squeezed into the region around the twiddle matrix. Absorb it:

Why $Z_{2^k,l}$ is still diagonal. Because $2^{k-l}\cdot 2^{l} = 2^k$, the inverse property from §2 gives $P(2^k,2^{l}) = P(2^k,2^{k-l})^{-1}$. So (13) is a similarity transform of a diagonal matrix by a permutation:

Conjugating a diagonal matrix by a permutation just reorders its diagonal — it can never introduce off-diagonal mass. $Z_{2^k,l}$ is diagonal, therefore it moves no data. It is a pure element-wise multiply.

That is the whole point. All stage-dependent data movement has been converted into stage-dependent twiddle-factor selection. Coefficients now flow through the same fixed shuffle every stage; only the ROM addresses change.

Compare against the previous figure: same butterflies, same number of stages, but one wiring pattern instead of three. This is a constant-geometry algorithm.

The trade is favourable because it is asymmetric: coefficient movement costs you on both the read and the write side every stage, whereas twiddle selection is a read-only lookup in the multiplier stage that never has to relocate anything in memory.

6. Collapsing the twiddle storage

Constant geometry is only a win if the twiddle addressing does not become the new mess. It doesn’t — it collapses to one line of index arithmetic.

Let $\mathbf{z}_0$ be the diagonal of $Z_0(l)$. From (6) and (7) it is $2^{l-1}$ repetitions of a $T_{2^{k-l}}(2^{k-l+1})$ block:

with $\omega$ primitive of order $2^{k-l+1}$. By the similarity argument the diagonal of $Z_{2^k,l}$ is $P(2^k,2^{k-l})\mathbf{z}_0$ — segment $\mathbf{z}_0$ into $2^{l}$ groups and deal one element from each group round-robin. The result interleaves into

Two things fall out.

One root for all stages. $\omega$ has order $2^{k-l+1}$ while $\omega_0$ has order $2^k$, so

Every stage’s twiddles are powers of the same $\omega_0$. No per-stage root tables.

A closed-form address. Every even-indexed coefficient pairs with $\omega^0 = 1$. For odd index $t$ at stage $l$, the required exponent is

Since $Z_r < 2^{k-1}$ always, a single table of $N/2$ powers $\omega_0^0,\dots,\omega_0^{N/2-1}$ — stored once, in order, never rewritten — serves the entire transform. The exponents in the second figure are computed from (16); I verified it reproduces the exact diagonal of (13) for all stages up to $N=64$.

So the “hard” part of the memory access pattern reduces to a shift and a mask.

7. What this buys in hardware

Scalable PE count. Split the butterfly bank using identity (1). With $\alpha$ PEs and $\beta = 2^{k-1}/\alpha$:

Each PE handles $\beta$ butterflies. $\alpha$ is now a free parameter — the same algorithm covers a 1-PE area-minimal design and a 64-PE throughput design with no restructuring. Each stage takes $\beta$ cycles, so latency is $\approx \frac{N}{2\alpha}\log_2 N$ butterfly slots.

Fixed interconnect. $P(2^k,2)$ never changes, so read/write addressing is generated once at elaboration time from $N$ and $\alpha$. The control unit shrinks to a counter plus (16) — which is what actually makes routing tractable as $\alpha$ grows, since a stage-varying permutation network is what usually caps the achievable frequency.

Ping-pong memory. The transform is not in-place under a fixed shuffle, so you need two banks, alternating by the parity of $l$. That is the honest cost of constant geometry: $2N$ coefficient storage instead of $N$. You are buying wiring simplicity with memory.

The paper’s Algorithm 1, which is (12) written as loops:

k = log2(N);  beta = 2^(k-1) / alpha
x = bit_reverse(x)
for l = k down to 1:
    for i = 0 to alpha-1:          # parallel across PEs
        for j = 0 to beta-1:       # sequential within a PE
            d  = i*beta + j
            w  = w0 ^ ( floor((2*d+1) / 2^l) * 2^(l-1) )    # eq (16)
            v[d]           = x[2*d] + x[2*d+1] * w   mod q
            v[d + 2^(k-1)] = x[2*d] - x[2*d+1] * w   mod q
    x = v
return x

The write pattern d and d + 2^(k-1) is $P(2^k,2)$ — the shuffle never appears as an explicit permutation step, it is just where the two butterfly outputs get written. I implemented this against a naive $O(N^2)$ NTT mod $q$ for $N \in \{8,16,32,64\}$ and $\alpha \in \{1,2,N/2\}$; all match.

8. Two caveats

This is Pease’s algorithm. Zhang et al. state theirs is “the first work to apply tensor products to NTT designs.” The tensor-product formalism for FFTs is standard and long-established — Tolimieri, An & Lu is the canonical treatment, and it is the paper’s own reference [25]; the SPIRAL project is built on it. More specifically, equation (12) is the Pease FFT (1968), transplanted from $\mathbb{C}$ to $\mathbb{Z}_q$: constant geometry with a fixed stride permutation and bit-reversed input is precisely what Pease derived. That does not diminish the accelerator — applying it to FHE-scale parameters with a scalable PE count and getting good silicon out is real work — but the derivation in §4–§6 is classical FFT theory, and it is worth knowing that so you can read the older literature for the variants.

This is the cyclic NTT, not the negacyclic one. Algorithm 1 advertises $\mathbb{Z}[x]/(x^N+1)$ but computes the plain length-$N$ cyclic transform — I confirmed this by testing against the cyclic definition, which is what it matches. RLWE-based FHE needs negacyclic convolution, which requires folding in powers of $\psi$ with $\psi^2 = \omega_0$: pre-scale $x_n \mapsto \psi^n x_n$, post-scale by $\psi^{-m}$, or use the standard merged-$\psi$ butterflies that absorb the scaling into the twiddle tables. That last option is what real designs do, and it changes the twiddle table contents but not the geometry — so §5 survives intact. Worth stating explicitly, because the omission is easy to miss.

Where I want to take this next: the same manipulation with a different split point gives the four-step / six-step decompositions, which is how you get NTTs whose working set fits in on-chip memory at FHE parameter sizes. Same algebra, different factorization — which is the real argument for the tensor-product view.

References

• Y. Zhang, S. R. Sathi, Z. Kou, S. Sinha, W. Zhang, “Tensor-Product-Based Accelerator for Area-efficient and Scalable Number Theoretic Transform,” FCCM, 2023. DOI: 10.1109/FCCM57271.2023.00027 — the source of the derivation in §3–§7.
• R. Tolimieri, M. An, C. Lu, Algorithms for Discrete Fourier Transform and Convolution, Springer, 1997 — the standard reference for the tensor-product/stride-permutation formalism. §1–§2 here is a compressed version of its opening chapters.
• M. C. Pease, “An Adaptation of the Fast Fourier Transform for Parallel Processing,” JACM 15(2), 1968 — the original constant-geometry FFT.
• C. Van Loan, Computational Frameworks for the Fast Fourier Transform, SIAM, 1992 — the clearest catalogue of FFT variants in Kronecker notation.
• M. Püschel et al., “SPIRAL: Code Generation for DSP Transforms,” Proc. IEEE 93(2), 2005 — tensor-product factorizations as a machine-searchable space.
• A. C. Mert, E. Karabulut, E. Öztürk, E. Savaş, A. Aysu, “An Extensive Study of Flexible Design Methods for the Number Theoretic Transform,” IEEE Trans. Computers, 2020 — the scalable iterative-NTT baseline Zhang et al. compare against.
• A. C. Mert, E. Öztürk, E. Savaş, “Design and Implementation of Encryption/Decryption Architectures for BFV Homomorphic Encryption Scheme,” IEEE TVLSI 28(2), 2020 — the 64-PE iterative and four-step comparison points.
• P. Longa, M. Naehrig, “Speeding up the Number Theoretic Transform for Faster Ideal Lattice-Based Cryptography,” CANS, 2016 — the merged-$\psi$ negacyclic construction referenced in §8.

← More from Research