Introduction to TileTensor

Let’s take a quick break from solving puzzles to preview a powerful abstraction that will make our GPU programming journey more enjoyable: 🥁… the TileTensor.

💡 This is a motivational overview of TileTensor’s capabilities. Don’t worry about understanding everything now - we’ll explore each feature in depth as we progress through the puzzles.

The challenge: Growing complexity

Let’s look at the challenges we’ve faced so far:

# Puzzle 1: Simple indexing
output[unsafe_offset=i] = a[unsafe_offset=i] + 10.0

# Puzzle 2: Multiple array management
output[unsafe_offset=i] = a[unsafe_offset=i] + b[unsafe_offset=i]

# Puzzle 3: Bounds checking
if i < size:
    output[unsafe_offset=i] = a[unsafe_offset=i] + 10.0

As dimensions grow, code becomes more complex:

# Traditional 2D indexing for row-major 2D matrix
var idx = row * WIDTH + col
if row < height and col < width:
    output[unsafe_offset=idx] = a[unsafe_offset=idx] + 10.0

The solution: A peek at TileTensor

TileTensor will help us tackle these challenges with elegant solutions. Here’s a glimpse of what’s coming:

Natural indexing: Use tensor[i, j] instead of manual offset calculations
Flexible memory layouts: Support for row-major, column-major, and tiled organizations
Performance optimization: Efficient memory access patterns for GPU

A taste of what’s ahead

Let’s look at a few examples of what TileTensor can do. Don’t worry about understanding all the details now - we’ll cover each feature thoroughly in upcoming puzzles.

Basic usage example

from layout import TileTensor
from layout.tile_layout import row_major

# Define layout
comptime HEIGHT = 2
comptime WIDTH = 3
comptime layout = row_major[HEIGHT, WIDTH]()
comptime LayoutType = type_of(layout)

# Create tensor
var tensor = TileTensor(buffer, layout)

# Access elements naturally
tensor[0, 0] = 1.0  # First element
tensor[1, 2] = 2.0  # Last element

To learn more about Layout and TileTensor, see these guides in the MAX documentation:

Quick example

Let’s put everything together with a simple example that demonstrates the basics of TileTensor:

# ===----------------------------------------------------------------------=== #
# Copyright (c) 2026, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
from max.gpu.host import DeviceContext
from layout import TileTensor
from layout.tile_layout import row_major

comptime HEIGHT = 2
comptime WIDTH = 3
comptime dtype = DType.float32
comptime layout = row_major[HEIGHT, WIDTH]()
comptime LayoutType = type_of(layout)


def kernel(
    tensor: TileTensor[mut=True, dtype, LayoutType, MutAnyOrigin],
):
    print("Before:")
    print(tensor)
    tensor[0, 0] += 1
    print("After:")
    print(tensor)


def main() raises:
    var ctx = DeviceContext()

    var a = ctx.enqueue_create_buffer[dtype](HEIGHT * WIDTH)
    a.enqueue_fill(0)
    var tensor = TileTensor(a, layout)
    # Note: since `tensor` is a device tensor we can't print it without the kernel wrapper
    ctx.enqueue_function[kernel](tensor, grid_dim=1, block_dim=1)

    ctx.synchronize()

When we run this code with:

pixi run tile_tensor_intro

pixi run -e amd tile_tensor_intro

pixi run -e apple tile_tensor_intro

uv run poe tile_tensor_intro

Before:
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
After:
[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]

A TileTensor whose dimensions are all known at compile time prints as nested rows on a single line. Give it a runtime dimension and it prints flat instead, element by element.

Let’s break down what’s happening:

We create a 2 x 3 tensor with row-major layout
Initially, all elements are zero
Using natural indexing, we modify a single element
The change is reflected in our output

This simple example demonstrates key TileTensor benefits:

Clean syntax for tensor creation and access
Automatic memory layout handling
Natural multi-dimensional indexing

While this example is straightforward, the same patterns will scale to complex GPU operations in upcoming puzzles. You’ll see how these basic concepts extend to:

Multi-threaded GPU operations
Shared memory optimizations
Complex tiling strategies
Hardware-accelerated computations

Ready to start your GPU programming journey with TileTensor? Let’s dive into the puzzles!

💡 Tip: Keep this example in mind as we progress - we’ll build upon these fundamental concepts to create increasingly sophisticated GPU programs.