Skip to main content

Python class

KernelLibrary

KernelLibrary

class max.graph.KernelLibrary(context, paths=[])

Manages custom kernel libraries and operations for a graph.

A kernel library provides access to custom operations and kernels that can be loaded from various sources including Mojo binary packages (.mojopkg) and Mojo source directories. The library handles verification and registration of custom operations within the MLIR context.

Parameters:

  • context (mlir.Context )
  • paths (list [ Path ] )

add_path()

add_path(path)

Adds a kernel library path to the analysis.

Parameters:

path (Path ) – The Path to the kernel library to be added to the current analysis.

Return type:

None

library_paths()

library_paths()

Returns the list of kernel library paths.

Returns:

A list of Path objects representing the currently loaded kernel library paths.

Return type:

list[Path]

load_paths()

load_paths(context, custom_extensions)

Loads custom operations from provided library paths.

Performs additional “smart” library loading logic for custom operation libraries in additional formats. The loading logic supports the following formats:

  • Compiled Mojo binary packages with .mojopkg extension
  • Mojo source directory with custom operations

The loaded libraries are added to the current kernel library.

Parameters:

  • context (Context ) – The MLIR context for loading MLIR operations.
  • custom_extensions (Iterable [ Path ] ) – The file paths to the custom operation libraries.

Return type:

None

verify_custom_op()

verify_custom_op(custom_op)

Verifies that a custom operation is valid within the current context.

Parameters:

custom_op (Operation ) – The mlir.Operation to be verified against the current kernel library analysis.

Return type:

None