Categorical Foundations for CuTe Layouts

TL;DR

This paper introduces a categorical framework for understanding and formalizing CuTe layouts in NVIDIA's CUTLASS library, providing mathematical foundations and a Python implementation that aligns with existing GPU tensor data manipulation methods.

Contribution

It develops a novel categorical model for CuTe layouts, characterizes the layouts arising from this model, and offers a Python implementation aligned with CUTLASS.

Findings

Categorical framework accurately models CuTe layouts.

Complete characterization of layout classes from the framework.

Python implementation demonstrates practical alignment with CUTLASS.

Abstract

NVIDIA's CUTLASS library provides a robust and expressive set of methods for describing and manipulating multi-dimensional tensor data on the GPU. These methods are conceptually grounded in the abstract notion of a CuTe layout and a rich algebra of such layouts, including operations such as composition, logical product, and logical division. In this paper, we present a categorical framework for understanding this layout algebra by focusing on a naturally occurring class of tractable layouts. To this end, we define two categories Tuple and Nest whose morphisms give rise to layouts. We define a suite of operations on morphisms in these categories and prove their compatibility with the corresponding layout operations. Moreover, we give a complete characterization of the layouts which arise from our construction. Finally, we provide a Python implementation of our categorical constructions, along with tests that demonstrate alignment with CUTLASS behavior. This implementation can be found at our git repository https://github.com/ColfaxResearch/layout-categories.