A Parallel Scan Algorithm in the Tensor Core Unit Model

Abstract

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 $(s^2,\ell )$-TCU model, we show that for inputs of size n, the algorithm has depth at most $2\lfloor \log _s(n)\rfloor$ and runs in $\mathcal {O}(n(1+\ell /s^2) / p + (s^2 + \ell ) \log _s (n))$ assuming p tensor core units. Equivalently, the algorithm performs $\mathcal {O}(n/s^2)$ multiplications of square matrices of size s.

A Appendix

We provide a functional end-to-end (but not high-performance) Python implementation of Algorithm 1 using NumPy (v1.24.1) [9]. The implementation demonstrates the memory layout operations required to orchestrate the batched matrix multiplications of Algorithm 1.

1.1 A.1 Correctness of Algorithm 1

In this section, we prove the correctness of Algorithm 1. We reformulate Algorithm 1 using recursion as stated in Algorithm 2. The recursive formulation will enable us to prove the correctness using strong induction. Indeed, we prove by induction that MatMulScanRecursive is correct for all inputs that are powers of s, given an arbitrary $s\ge 2$.

In particular, it suffices to show that the Recurse method with input $\boldsymbol{z}$ and s has the following precondition/postcondition relation: given the precondition that on all consecutive chunks of size s of $\boldsymbol{z}$, i.e., $(0,1,\dots , s-1), (s,s+1,\dots , 2s-1), \dots$, the “local” prefix sums on each chunk is precomputed, Recurse returns the prefix sum of $\boldsymbol{z}$ (postcondition). Indeed, by the definition of MatMulScanRecursive in Line 2 the “local” prefix sums of size s are computed and, in Line 3, the Recurse method is called with the precondition to be true.

Base Case. For inputs of size less than s, the termination criterion of Line 6 is met, therefore the postcondition follows directly from the precondition since the input size is less than s.

Inductive Step. The inductive hypothesis is that Recurse is correct for input sizes strictly less than n. We will show that Recurse is correct for inputs of size n. Indeed, given an input $\boldsymbol{z}$ where all its “local” prefix sums are precomputed, we prove that Recurse with input $\boldsymbol{z}$ returns the prefix sum of $\boldsymbol{z}$. Now, Line 8 computes the “local” prefix sums on the s-strided subvector $\boldsymbol{x}[\text {start}::\text {step}]$. The prefix sum on $\boldsymbol{x}[\text {start}::\text {step}]$ is computed on Line 9 by the inductive hypothesis. Then, Line 11 broadcasts and add the correct prefix sum values of the s-strided subvector of $\boldsymbol{z}$ to the corresponding s following indices of each subvector. Hence, the postcondition of Recurse holds.