tessl/pypi-iree-base-compiler
IREE Python Compiler API - MLIR-based end-to-end compiler for Machine Learning models
—
Pending
api-embedding.mddocs/
Compiler API Embedding
Low-level Python bindings for IREE's C embedding API, providing direct access to compiler infrastructure for advanced use cases requiring fine-grained control over compilation sessions, invocations, and I/O handling.
Capabilities
Session Management
Manages compiler sessions that wrap MLIR contexts with configuration flags for setting up the compiler environment.
class Session:
"""
Compiler session wrapping an MLIRContext with configuration flags.
Represents the top-level compiler state and configuration. Sessions are
thread-safe for read operations but require external synchronization
for mutations.
"""
def __init__(self, **kwargs): ...
def invocation(self) -> Invocation:
"""
Create a new compiler invocation within this session.
Returns:
Invocation: New invocation object for compilation
"""
def set_flags(self, *flags: str) -> None:
"""
Set compiler flags for this session.
Parameters:
- *flags: Variable number of compiler flag strings
"""
def get_flags(self, non_default_only: bool = False) -> Sequence[str]:
"""
Get current compiler flags.
Parameters:
- non_default_only: Only return non-default flags
Returns:
Sequence[str]: List of current flags
"""
@property
def context(self):
"""
Get the MLIR context for this session.
Returns:
MLIRContext: The underlying MLIR context
"""
def close(self) -> None:
"""
Close the session and release resources.
"""Compilation Invocations
Manages individual compilation invocations that wrap MLIR Modules during the compilation process.
class Invocation:
"""
Compiler invocation wrapping a Module being compiled.
Represents a single compilation unit with its associated source,
transformations, and output. Invocations are not thread-safe.
"""
def parse_source(self, source: Source) -> bool:
"""
Parse source input for this invocation.
Parameters:
- source: Source object containing input to parse
Returns:
bool: True if parsing succeeded
"""
def pipeline(self, pipeline_type: int) -> bool:
"""
Run a compilation pipeline.
Parameters:
- pipeline_type: Pipeline type constant
Returns:
bool: True if pipeline succeeded
"""
def run_pass_pipeline(self, pipeline_str: str) -> bool:
"""
Run a custom pass pipeline.
Parameters:
- pipeline_str: Pass pipeline string
Returns:
bool: True if pipeline succeeded
"""
def output_ir(self, output: Output) -> None:
"""
Output MLIR IR to an output object.
Parameters:
- output: Output object to write to
"""
def output_ir_bytecode(self, output: Output, bytecode_version: int) -> None:
"""
Output MLIR bytecode to an output object.
Parameters:
- output: Output object to write to
- bytecode_version: Bytecode format version
"""
def output_vm_bytecode(self, output: Output) -> None:
"""
Output VM bytecode to an output object.
Parameters:
- output: Output object to write to
"""
def enable_console_diagnostics(self) -> None:
"""
Enable console diagnostic output.
"""
def export_module(self):
"""
Export the compiled module as MLIR Operation.
Returns:
ir.Operation: MLIR Operation containing the module
"""
def close(self) -> None:
"""
Close the invocation and release resources.
"""Source Input
Manages input sources for compilation including MLIR text, bytecode, and other supported formats.
class Source:
"""
Source input for compiler invocations.
Wraps various input formats including MLIR assembly text,
MLIR bytecode, and other supported source formats.
"""
@staticmethod
def open_file(session: Session, file_path: str) -> Source:
"""
Create source from file path.
Parameters:
- session: Session object
- file_path: Path to source file
Returns:
Source: Source object for the file
"""
@staticmethod
def wrap_buffer(session: Session, buffer, *, buffer_name: str = "source.mlir") -> Source:
"""
Create source from memory buffer.
Parameters:
- session: Session object
- buffer: Buffer containing source data
- buffer_name: Name for the buffer source
Returns:
Source: Source object wrapping the buffer
"""
def close(self) -> None:
"""
Close the source and release resources.
"""Output Management
Manages compilation outputs including VM bytecode, MLIR text, and other generated formats.
class Output:
"""
Compilation output containing generated code and metadata.
Provides access to the results of compilation including
bytecode, text representations, and diagnostic information.
"""
@staticmethod
def open_file(file_path: str) -> Output:
"""
Create output that writes to a file.
Parameters:
- file_path: Path to output file
Returns:
Output: Output object for the file
"""
@staticmethod
def open_membuffer() -> Output:
"""
Create output that writes to memory buffer.
Returns:
Output: Output object for memory buffer
"""
def keep(self) -> Output:
"""
Mark output to be kept after close.
Returns:
Output: Self for chaining
"""
def write(self, buffer) -> None:
"""
Write data to output.
Parameters:
- buffer: Data to write
"""
def map_memory(self) -> memoryview:
"""
Map output memory for reading.
Returns:
memoryview: Memory view of output content
"""
def close(self) -> None:
"""
Close the output and release resources.
"""Pipeline Types
Enumeration of compilation pipeline types for the pipeline() method.
class PipelineType(IntEnum):
"""
Compilation pipeline type constants.
"""
IREE_COMPILER_PIPELINE_STD = 0
IREE_COMPILER_PIPELINE_HAL_EXECUTABLE = 1
IREE_COMPILER_PIPELINE_PRECOMPILE = 2Global Initialization
Global initialization function for setting up command-line argument processing.
def _initializeGlobalCL(*cl_args: str) -> None:
"""
Initialize global command-line processing.
Parameters:
- *cl_args: Command-line arguments to initialize with
Must be called once before using the embedding API.
Sets up LLVM command-line argument parsing and global state.
"""Usage Examples
Basic Compilation Workflow
from iree.compiler.api import Session, Source, _initializeGlobalCL
# Initialize global state
_initializeGlobalCL()
# Create session
session = Session()
try:
# Create invocation
invocation = session.invocation()
# Set up source
mlir_code = """
func.func @simple_add(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
%0 = arith.addf %arg0, %arg1 : tensor<4xf32>
return %0 : tensor<4xf32>
}
"""
source = Source.wrap_buffer(session, mlir_code.encode())
invocation.parse_source(source)
# Configure compilation
session.set_flags(
"--iree-hal-target-backends=llvm-cpu",
"--iree-input-type=auto"
)
# Create output and run compilation pipeline
output = Output.open_membuffer()
invocation.pipeline(0) # IREE_COMPILER_PIPELINE_STD
invocation.output_vm_bytecode(output)
# Get results
bytecode = output.map_memory()
# Clean up
output.close()
invocation.close()
source.close()
finally:
session.close()File-based Compilation
from iree.compiler.api import Session, Source
session = Session()
try:
invocation = session.invocation()
# Load from file
source = Source.open_file(session, "model.mlir")
invocation.parse_source(source)
# Set compilation flags
session.set_flags(
"--iree-hal-target-backends=cuda",
"--iree-input-type=stablehlo"
)
# Compile and save
output = Output.open_file("model.vmfb")
invocation.pipeline(0) # IREE_COMPILER_PIPELINE_STD
invocation.output_vm_bytecode(output)
# Clean up
output.close()
invocation.close()
source.close()
finally:
session.close()Advanced Session Configuration
from iree.compiler.api import Session
# Configure session with specific options
session = Session(
enable_debug_info=True,
enable_crash_reproducer=True
)
try:
# Set session-level flags
session.set_flags("--iree-hal-target-backends=vulkan-spirv")
# Multiple invocations can share the same session
for input_file in ["model1.mlir", "model2.mlir"]:
invocation = session.invocation()
source = Source.open_file(session, input_file)
invocation.parse_source(source)
output = Output.open_file(f"{input_file}.vmfb")
invocation.pipeline(0) # IREE_COMPILER_PIPELINE_STD
invocation.output_vm_bytecode(output)
output.close()
invocation.close()
source.close()
finally:
session.close()Install with Tessl CLI
npx tessl i tessl/pypi-iree-base-compiler