A Parallel Scan Algorithm in the Tensor Core Unit Model

We present a parallel scan (prefix sum) algorithm in the Tensor Core Unit (TCU) model of computation. The TCU model assumes that multiplication between two square matrices of constant size s is a basic operation. In the (s2, ℓ)-TCU model, we show that for inputs of size n, the algorithm has depth at most 2lfloor logs (n)rfloor and runs in O(n(1 + ℓ /s2)/p + (s2 + ℓ) logs (n)) time assuming p tensor core units. Equivalently, the algorithm performs O(n/s2) multiplications of square matrices of size s.