A Parallel Scan Algorithm in the Tensor Core Unit Model

TL;DR

This paper introduces a parallel scan algorithm tailored for the Tensor Core Unit model, leveraging matrix multiplication as a fundamental operation to optimize prefix sum computations.

Contribution

It proposes a novel parallel scan algorithm within the TCU model, analyzing its depth and runtime based on matrix multiplication operations.

Findings

Algorithm achieves depth at most 2*log_s(n)

Runs in O(n(1 + l/s^2)/p + (s^2 + l) log_s(n)) time

Performs O(n/s^2) matrix multiplications

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 $O(n(1 + \ell /s^2)/p + (s^2 + \ell) \log_s (n))$ time assuming $p$ tensor core units. Equivalently, the algorithm performs $O(n/s^2)$ multiplications of square matrices of size s.