CuTe Layouts in Lean 4

A Lean 4 formalization of Categorical Foundations for CuTe Layouts.

What is CuTe?

CuTe (CUDA Templates) is NVIDIA's tensor layout library. This project formalizes its mathematical foundations: the category Tuple whose morphisms correspond to non-degenerate tractable layouts.

Structure

CuTe/
├── Layout/
│   ├── Defs.lean       # ShapeStridePair, FlatLayout, NestedTuple
│   ├── Order.lean      # Lexicographic ordering ⪯
│   ├── Tractable.lean  # Tractability, non-degeneracy predicates
│   └── Operations.lean # coalesce, complement, toCoords
├── Category/
│   ├── Fin0.lean       # Pointed finite sets (Option (Fin n))
│   ├── Tuple.lean      # Category Tuple, TupleMorphism, toLayout
│   ├── Functor.lean    # Realization functor, colex
│   └── Nest.lean       # Nested tuples, hierarchical layouts
└── Morphism/
    └── Encode.lean     # Layout ↔ Morphism bijection

Key Concepts

• Layout: List of (shape, stride) pairs [(s₁,d₁), ..., (sₘ,dₘ)]
• Tractable: Divisibility condition on strides
• TupleMorphism: Dimension-preserving pointed set map
• Correspondence: Tractable layouts ↔ tuple morphisms

Building

lake build

Demo

Run the compute attestation demo to see example workloads and threshold checking:

lake exe cute-demo

This demonstrates:

• FLOPs counting for various ML workloads (matmuls, GPT-2 style inference)
• Threshold attestation against governance limits (10²³ to 10²⁶ FLOPs)
• Conservative handling of unknown operations

Cuda2Lean: CUDA as Lean Data

The Cuda2Lean library provides a deep embedding of CUDA code as Lean data structures, enabling formal verification of CUDA programs against specifications.

Cuda2Lean/
├── AST/           # CUDA syntax as inductive types (Type, Expr, Stmt, Program)
├── Semantics/     # Operational semantics (Value, Env, Eval, Exec)
├── Programs/      # Specific programs (OpBuffer, Counter kernel)
└── Correspondence/# Proofs matching CUDA to CuTe.Compute spec

CUDA Analyzer

Analyze real CUDA files to see how they map to Cuda2Lean representations:

# Clone NVIDIA's CUDA samples (one-time setup)
git clone https://github.com/NVIDIA/cuda-samples data/cuda-samples

# Analyze a random .cu file
lake exe cuda-analyze

# Or specify a directory
lake exe cuda-analyze data/cuda-samples/Samples/3_CUDA_Features/cudaTensorCoreGemm

The analyzer:

• Finds __global__ kernels and extracts parameters
• Detects operation types (matmul, convolution, reduction) via heuristics
• Shows the CudaKernel representation
• Extracts dimension #defines (M, N, K, BLOCK_SIZE, etc.)
• Demonstrates how to create OpEntry for FLOPs attestation

Example output for a tensor core GEMM:

Detected Kernels: compute_gemm (line 197)
Operation Type: matmul
Relevant Defines: M=16, N=16, K=16

Suggested attestation:
  let entry := OpEntry.matmul M N K
  let flops := computeFlops entry  -- 2 * M * N * K

References

• Paper
• Python reference
• CLAUDE.md for implementation details