Puzzle 17: 1D Convolution Op

Bridging to Python with MAX Graph

We’re now entering Part IV of our GPU puzzle journey: Interfacing with Python via MAX Graph Custom Ops.

In previous puzzles, we’ve learned how to write efficient GPU kernels in Mojo. Now we’ll explore how to:

  • Package these kernels as custom operations that can be called from Python
  • Integrate with the MAX Graph system for accelerated machine learning
  • Bridge the gap between high-level Python APIs and low-level GPU code

This integration allows us to leverage the performance of Mojo GPU kernels while working in familiar Python environments.

Overview

In Puzzle 13, we implemented a 1D convolution kernel that runs efficiently on the GPU. Now we’ll take this kernel and transform it into a custom operation that can be called directly from Python using MAX Graph.

The 1D convolution kernel we’ll be working with is already implemented:

comptime TPB = 15
comptime BLOCKS_PER_GRID = (2, 1)


def conv1d_kernel[
    input_size: Int,
    conv_size: Int,
    OutLayout: TensorLayout,
    InLayout: TensorLayout,
    ConvLayout: TensorLayout,
    dtype: DType = DType.float32,
](
    output: TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin],
    input: TileTensor[mut=True, dtype, InLayout, MutAnyOrigin],
    kernel: TileTensor[mut=True, dtype, ConvLayout, MutAnyOrigin],
):
    var global_i = block_dim.x * block_idx.x + thread_idx.x
    var local_i = thread_idx.x
    # Convert generic TileTensors to LayoutTensor for indexing (flat_rank proof required)
    var input_lt = input.to_layout_tensor()
    var kernel_lt = kernel.to_layout_tensor()
    var output_lt = output.to_layout_tensor()
    # first: need to account for padding
    var shared_a = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[TPB + conv_size - 1]())
    var shared_b = stack_allocation[
        dtype=dtype, address_space=AddressSpace.SHARED
    ](row_major[conv_size]())
    if global_i < input_size:
        shared_a[local_i] = rebind[Scalar[dtype]](input_lt[global_i])

    # second: load elements needed for convolution at block boundary
    if local_i < conv_size - 1:
        # indices from next block
        var next_idx = global_i + TPB
        if next_idx < input_size:
            shared_a[TPB + local_i] = rebind[Scalar[dtype]](input_lt[next_idx])
        else:
            # Initialize out-of-bounds elements to 0 to avoid reading from uninitialized memory
            # which is an undefined behavior
            shared_a[TPB + local_i] = 0

    if local_i < conv_size:
        shared_b[local_i] = rebind[Scalar[dtype]](kernel_lt[local_i])

    barrier()

    if global_i < input_size:
        var local_sum: Scalar[dtype] = 0

        comptime for j in range(conv_size):
            if local_i + j < TPB + conv_size - 1:
                local_sum += shared_a[local_i + j] * shared_b[j]

        output_lt.store[1](Index(global_i), local_sum)

The key aspects of this puzzle include:

  1. Custom op registration: Understanding how to expose Mojo functions to Python via the @extensibility.register decorator
  2. Packaging custom ops: Learning how to package Mojo code for use with MAX Graph
  3. Python integration: Calling custom operations from Python through MAX Graph
  4. Cross-language data flow: Managing data types and memory between Python and GPU

This custom operation will:

  • Accept NumPy arrays as input from Python
  • Transfer this data to the GPU
  • Execute our optimized convolution kernel
  • Return the results back to Python

When you complete this puzzle, you’ll have created a seamless bridge between Python’s rich ecosystem and Mojo’s powerful GPU performance.

Code to complete

To complete this puzzle, you only need to fill in the call to conv1d_kernel in conv1d.mojo:

import extensibility

from extensibility import InputTensor, OutputTensor
from max.gpu.host import DeviceBuffer


@extensibility.register("conv1d")
struct Conv1DCustomOp:
    @staticmethod
    def execute[
        # The kind of device this will be run on: "cpu" or "gpu"
        target: StaticString,
        input_size: Int,
        conv_size: Int,
        dtype: DType = DType.float32,
    ](
        output: OutputTensor[dtype=dtype, rank=1, static_spec=_],
        input: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        kernel: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        # the context is needed for some GPU calls
        ctx: DeviceContext,
    ) raises:
        comptime out_layout_val = row_major[input_size]()
        comptime OutLayout = type_of(out_layout_val)
        comptime conv_layout_val = row_major[conv_size]()
        comptime ConvLayout = type_of(conv_layout_val)

        var output_tensor = TileTensor[
            mut=True, dtype, OutLayout, MutAnyOrigin
        ](output.unsafe_ptr(), out_layout_val)
        var input_tensor = TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin](
            input.unsafe_ptr(), out_layout_val
        )
        var kernel_tensor = TileTensor[
            mut=True, dtype, ConvLayout, MutAnyOrigin
        ](kernel.unsafe_ptr(), conv_layout_val)

        comptime if target == "gpu":
            var gpu_ctx = ctx
            # making sure the output tensor is zeroed out before the kernel is called
            gpu_ctx.enqueue_memset(
                DeviceBuffer[output_tensor.dtype](
                    gpu_ctx,
                    output.unsafe_ptr(),
                    input_size,
                    owning=False,
                ),
                0,
            )

            # FILL ME IN with 2 lines calling our conv1d_kernel

        elif target == "cpu":
            # we can fallback to CPU
            pass
        else:
            raise Error("Unsupported target: " + target)

View full file: problems/p17/op/conv1d.mojo

You can run the puzzle with:

pixi run p17

pixi run -e amd p17

pixi run -e apple p17

uv run poe p17

When successful, you should see output similar to:

Input array: [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14.]
Convolution kernel: [0. 1. 2. 3.]
Expected result (NumPy calculation): [14. 20. 26. 32. 38. 44. 50. 56. 62. 68. 74. 80. 41. 14.  0.]
Compiling 1D convolution graph...
Executing 1D convolution...
1D Convolution result (custom Mojo kernel): [14. 20. 26. 32. 38. 44. 50. 56. 62. 68. 74. 80. 41. 14.  0.]
Verification passed: Custom kernel results match NumPy calculation

This indicates that your custom MAX Graph operation correctly implements the 1D convolution algorithm.

Solution

To solve this puzzle, we need to integrate our 1D convolution kernel with the MAX Graph system. The key is to properly call our kernel from the execute method in the Conv1DCustomOp struct.

The solution is:

            comptime kernel = conv1d_kernel[
                input_size, conv_size, OutLayout, OutLayout, ConvLayout
            ]
            gpu_ctx.enqueue_function[kernel](
                output_tensor,
                input_tensor,
                kernel_tensor,
                grid_dim=BLOCKS_PER_GRID,
                block_dim=(TPB, 1),
            )
These two statements do several important things:
  1. Calls enqueue_function on the GPU context (gpu_ctx is of type DeviceContext) to schedule our kernel execution
  2. Binds the layout and size information as compile-time parameters through the comptime kernel = conv1d_kernel[...] binding
  3. Provides the output, input, and kernel tensors as runtime arguments
  4. Configures the execution grid with the appropriate dimensions

Let’s break down how this works in the larger context:

Python-Mojo integration flow

  1. Python side (problems/p17/p17.py):

    • Creates NumPy arrays for input and kernel
    • Calls conv_1d() function which wraps our operation in MAX Graph
    • Converts NumPy arrays to MAX driver Buffers with Buffer.from_numpy(input).to(device)
    • Loads the custom operation package with custom_extensions=[mojo_kernels]

  2. Graph building:

  3. Custom op registration:

    • The @extensibility.register("conv1d") decorator exposes our operation to MAX Graph. See @extensibility.register
    • The execute method parameters define the interface (inputs, outputs, context)
    • Input/output tensors are converted to TileTensors for use in our kernel
    • Device context manages GPU memory allocation and kernel execution

  4. Kernel execution:

    • When model.execute(...) is called, our conv1d_kernel receives the data
    • GPU thread configuration is set with grid_dim and block_dim
    • Results are transferred back to CPU with result.to(CPU())
    • NumPy verification compares our results with the expected output

Key components in detail

  1. Custom Op Structure:

    @extensibility.register("conv1d")
struct Conv1DCustomOp:
    @staticmethod
    def execute[target: StaticString, input_size: Int, conv_size: Int, dtype: DType = DType.float32](
        output: OutputTensor[dtype=dtype, rank=1, static_spec=_],
        input: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        kernel: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        ctx: DeviceContext,
    ) raises:
        # Implementation
    • target indicates the device type (“gpu” or “cpu”)
    • input_size and conv_size are parameters passed from Python
    • Tensor types ensure correct shape and type checking
    • raises is the effect annotation, marking that execute can propagate errors

  2. Tensor Conversion:

    comptime out_layout_val = row_major[input_size]()
comptime OutLayout = type_of(out_layout_val)

var output_tensor = TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin](
    output.unsafe_ptr(), out_layout_val
)
var input_tensor = TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin](
    input.unsafe_ptr(), out_layout_val
)

    The layouts are constructed from the op’s compile-time parameters and the tensors are built from raw pointers—they are not extracted from the OutputTensor/InputTensor arguments.

    • MAX Graph tensors are wrapped as Mojo TileTensors over the same memory
    • This allows our kernel to work with them directly
    • Because the layouts are compile-time values, the kernel’s indexing arithmetic is resolved statically

  3. Device Context Usage:

    var gpu_ctx = ctx
gpu_ctx.enqueue_memset(...)  # Zero output buffer
gpu_ctx.enqueue_function[...](...) # Schedule kernel
    • Device context manages GPU resources
    • Memory operations ensure correct buffer state
    • Function enqueueing schedules our kernel for execution

This solution demonstrates the complete flow from Python data through MAX Graph to GPU execution and back, leveraging Mojo’s powerful type system and parametric functions to create efficient, type-safe, accelerated operations.

Understanding MAX Graph custom ops

Check out the follow tutorials for more details:

Custom op registration

The core of creating a custom operation is the @extensibility.register decorator and the associated structure:

@extensibility.register("conv1d")
struct Conv1DCustomOp:
    @staticmethod
    def execute[...](
        output: OutputTensor[dtype=dtype, rank=1, static_spec=_],
        input: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        kernel: InputTensor[dtype=dtype, rank=output.rank, static_spec=_],
        ctx: DeviceContext,
    ) raises:
        # Implementation here

Key components of the registration:

  • The name passed to the decorator ("conv1d") is what Python code will use to call this operation
  • The struct must have an execute method with the correct signature
  • OutputTensor and InputTensor types define the interface for Python data
  • DeviceContext provides access to the execution environment

Packaging custom ops

Before the custom operation can be used from Python, it needs to be packaged:

mojo package op -o op.mojoc

This command:

  1. Compiles the Mojo code into a deployable package
  2. Creates the necessary metadata for MAX Graph to understand the operation
  3. Produces a binary artifact (op.mojoc) that can be loaded by Python

The package must be placed in a location where MAX Graph can find it, typically in a directory accessible to the Python code.

Python integration

On the Python side, here’s how the custom operation is used:

# Path to the directory containing our Mojo operations
mojo_kernels = Path(__file__).parent / "op"

# Configure our graph with the custom conv1d operation
with Graph(
    "conv_1d_graph",
    input_types=[...],
    custom_extensions=[mojo_kernels],  # Load our custom op package
) as graph:
    # Define inputs to the graph
    input_value, kernel_value = graph.inputs

    # Use our custom operation by name
    output = ops.custom(
        name="conv1d",  # Must match the name in @extensibility.register
        values=[input_value, kernel_value],
        out_types=[...],
        parameters={
            "input_size": input_tensor.shape[0],
            "conv_size": kernel_tensor.shape[0],
            "dtype": dtype,
        },
    )[0].tensor

The key elements are:

  1. Specifying the path to our custom operations with custom_extensions
  2. Calling ops.custom with the registered operation name
  3. Passing input values and parameters that match our operation’s signature