iree.turbine.aot¶

aot¶

Toolkit for ahead-of-time (AOT) compilation and export of PyTorch programs.

class iree.turbine.aot.AbstractTensor(*size: int | None, dtype: dtype = torch.float32)¶

Represents a tensor of known rank and dtype.

create_intrinsic(ir_value: Value) → Intrinsic¶: Creates a proxy object that can flow through a procedural trace.

dtype¶

get_ir_type(builder: ModuleBuilder) → Type¶: Gets the corresponding IR type.

size¶

class iree.turbine.aot.CompiledModule(*, context: ~iree.compiler._mlir_libs._site_initialize.<locals>.Context | None = None, module_op: ~iree.compiler._mlir_libs._mlir.ir.Operation | None = None, import_to: ~iree.turbine.aot.compiled_module.ImportPhase | None | str = 'import')¶

Base class for all staged modules.

classmethod create_from_dict(name: str, dct: dict, *, export_name: str | None = None, options: ModuleBuilderOptions | None = None) → CompiledModuleMeta¶

Creates a CompiledModule subclass with an explicit dictionary of members.

This is the unsugared form of:

``` class Foo(CompiledModule, export_name=”bar”):

def member(): …

```

static expand_custom_ops(inst: CompiledModule)¶: Performs custom torch.operator expansion for known custom ops.

static get_class_info(cls: CompiledModuleMeta) → CompiledModuleClassInfo¶

static get_info(inst: CompiledModule) → CompiledModuleInstanceInfo¶

static get_mlir_module(inst: CompiledModule) → Operation¶

static get_module_builder(inst: CompiledModule) → Operation¶

class jittable(wrapped_f, *, decompose_ops: List[Any] | None = None, decomposition_table: Dict[Any, Callable[[...], Any]] | None = None, dynamic_shapes: Dict[str, Any] = {}, function_name: str | None = None, passes: Sequence[str] = ('functorch_functionalize',))¶

Decorator which takes a PyTorch function and makes it callable from tracing.

It will be internally JIT-ed and exported into the module as needed.

decomposition_table¶

dynamic_shapes¶

function_name¶

resolve_call(proc_trace: IrTrace, *py_args, dynamic_shapes: Dict[str, Any] = {}, **py_kwargs)¶

wrapped_f¶

static run_import(inst: CompiledModule, import_to: ImportPhase | str | None = 'import')¶

static run_pass_pipeline(inst: CompiledModule, pipeline: str, enable_ir_printing: bool = False)¶

Runs an arbitrary pass pipeline against the current IR.

Parameters:

pipeline – The text format pass pipeline as supported by PassManager.parse.
enable_ir_printing – Enables print-after-all to stderr.

static save_mlir(inst: CompiledModule, path: Path | str)¶

Saves a snapshot of the MLIR module in this CompiledModule to a file.

This is a convenience wrapper around the facilities of the underlying API and does not expose all features.

Parameters:: path – The file path to write to. If the extension is “.mlirbc”, it will be written as bytecode.

static signature_info(*, arg_device: dict[int, DeviceAffinity] | None = None) → Callable¶: Annotate an export target function. This annotation is only required when additional information needs to be provided.

class iree.turbine.aot.DeviceAffinity(ordinal: int, queues: list | None = None)¶: This is used to provide device affinities to exported function arguments.

class iree.turbine.aot.DeviceTensorTrait(ordinal: int, queues: list | None = None)¶

Represents a ‘trait’ that can be applied to a Tensor to signal that it is to be loaded to a speific device at execution time.

static get(from_tensor: Tensor) → DeviceTensorTrait | None¶

ordinal: int¶

queues: list | None = None¶

set(to_tensor: Tensor)¶

class iree.turbine.aot.ExportOutput(session: Session, compiled_module: CompiledModule, *, importer_uses_session: bool = False)¶

Wrapper around a CompiledModule produced by export.

Compiles the exported program to an executable binary.

Parameters:

save_to – Where to save the compiled binary. Can be one of: None: outputs to a memory buffer and return the API Output. (str, Path): Outputs to a file Output: Raw compiler API Output object to save to.
target_backends – A comma-delimitted string of IREE target backends or a sequence of strings. If None does not specify any target backend. Then the user must set other appropriate compiler flags e.g. export_output.session.set_flags(”–iree-hal-target-device=llvm-cpu”)

Returns:

None unless if save_to=None, in which case, we return the backing compiler API Ouptut object. It can be queried for its backing memory via its map_memory() method.

import_to(import_to: ImportPhase | str)¶

Compiles the modules to a mnemonic import phase.

This is a no-op if already compiled to this phase.

property mlir_module: Operation¶: Gets the MLIR module resulting from the last compilation phase.

print_readable(large_elements_limit: int = 50)¶: Prints a human readable version of the current compilation IR.

save_mlir(file_path: str | Path)¶

Saves the current compilation IR to a path on disk.

Parameters:: file_path – Path to save the file. If it has a “.mlirbc” extension, it will be saved as bytecode. Otherwise as text.

verify()¶: Runs the verifier on the module, raising an exception on failure.

class iree.turbine.aot.ExternalTensorTrait(external_scope: str, external_name: str)¶

Represents a ‘trait’ that can be applied to a Tensor to signal that it is to be loaded by name from an external archive at AOT execution time.

external_name: str¶

external_scope: str¶

static get(from_tensor: Tensor) → ExternalTensorTrait | None¶

set(to_tensor: Tensor)¶

class iree.turbine.aot.FxPrograms¶

Represents a named set of ExportedPrograms.

This facility works around a design flaw in Torch where they conflated ExportedPrograms as representing a single entry-point while also having each instance persist its own state_dict and constants. How many times, in how many frameworks, do we have to fight this design flaw? Apparently once more.

This base class represents the set of programs, either loaded from storage or built live. The tricky part it is managing is to do all of this while aliasing state and captured constants. Having those be physically shared is an essential optimization.

In order to manage saving/loading of the set of things, we manually splice the state_dict and constants dict such that while saving, we only persist the first encountered instance of any reference. Any subsequent instances are replaced with a SharedStateTensor, which on load can be re-associated.

As this is primarily targeted at being able to decouple FX tracing from further manipulation (which for reasons unknown, is competing with the race of entropy to the heat death of the universe in terms of performance), we don’t take a lot of pains to be optimized for distribution or storage of the resulting artifacts.

In the future, this same technique could be employed to elide parameters that we know we are going to resolve symbolically later, keeping them from being loaded and consuming memory during model export and compilation.

We have faith that in the fullness of time, the design flaws in Torch that require this kind of thing to exist will be resolved, and we then won’t need this hack.

static load(path: str | PathLike) → FxPrograms¶

save(path: str | PathLike) → int¶

Saves the set of exported programs to a descriptor file.

Returns the number of tensors deduped (for debugging/testing).

class iree.turbine.aot.FxProgramsBuilder(root_module: Module)¶

Builds a new set of exported programs that are all variations of the same root nn.Module.

This can be used to construct multi-entrypoint sets of ExportedPrograms in a way that alias information is preserved for lifted tensors.

Usage:

``` class MyModule(nn.Module):

…

fxb = FxProgramBuilder(MyModule())

@fxb.export_program(args=example_args) def entrypoint(m, x, y):

return m.forward(x, y)

fxb.save(“/some/path.json”) ```

export_program(f=None, *, args=None, kwargs=None, dynamic_shapes=None, strict: bool = True, name: str | None = None, arg_device: dict[int, DeviceAffinity] | None = None)¶

class iree.turbine.aot.ParameterArchive(file_path: str | Path | None = None, *, mmap: bool = True, readable: bool = True, writable: bool = False)¶

Allows access to a parameter archive as CPU tensors.

TODO: Add more helpers for reading tensors once we get upstream versions that have that integrated.

property index: ParameterIndex¶

items() → List[Tuple[str, ParameterArchiveEntry]]¶

Returns the items in the archive.

Note that there can be duplicates if the archive was constructed that way.

load(file_path: str | Path, *, mmap: bool = True, readable: bool = True, writable: bool = False)¶: Loads index entries from a file adding them to the in-memory archive.

class iree.turbine.aot.ParameterArchiveBuilder¶

Helper for building parameter archives from live modules.

add_blob(key: str, blob)¶

Adds a raw blob to the index.

The blob must be interpretable as a buffer.

add_module(module: Module, *, prefix: str = '')¶: Adds all parameters and persistent buffers from a module hierarchy.

add_tensor(name: str, tensor: Tensor)¶: Adds an named tensor to the archive.

property index: ParameterIndex¶

save(file_path: str | Path)¶: Saves the archive.

class iree.turbine.aot.ParameterArchiveEntry(raw: ParameterIndexEntry)¶

Wraps a raw ParameterIndexEntry with additional helpers.

as_flat_tensor() → Tensor¶

Accesses the contents as a uint8 flat tensor.

If it is a splat, then the tensor will be a view of the splat pattern.

Raises a ValueError on unsupported entries.

as_tensor() → Tensor¶

Returns a tensor viewed with appropriate shape/dtype from metadata.

Raises a ValueError if unsupported.

property key: str¶

iree.turbine.aot.abstractify(tree)¶

iree.turbine.aot.current_aot_decompositions() → Dict[OperatorBase, Callable]¶: Gets the current decomposition table for AOT.

Generic export of supported entities.

See a more specific overload for accepted forms.

This function behaves differently based on the type of the mdl argument:

nn.Module: The module is traced with torch.export.export passing it args, kwargs, and dynamic_shapes.
CompiledModule: The module is imported to IR. Additional arguments are illegal in this case.
torch.export.ExportedProgram: A pre-exported program can be passed and it will be used to construct a single-entrypoint module.

Parameters:

mdl – The nn.Module to export.
*example_args – Example tensors.
args – Example arguments to torch.export (if present, then *example_args must be empty.
kwargs – Example keyword arguments.
dynamic_shapes – Dynamic shape specs to pass to torch.export.
arg_device – device affinities for the exported function arguments. On what devices should the program expect its arguments. It is a mapping of argument index to device affinity of the flattened arguments.

Returns:

An ExportOutput object that wraps the compilation and provides easy access.

Exports buffers from an nn.Module.

These are exposed to procedural programs as a dictionary of param/values.

abstractify_tree()¶

items()¶: Yields tuples of name/value exports.

schema() → TreeSpec¶: A schema used to unflatten for access from Python.

Exports a single global into a CompiledModule.

abstractify() → AbstractTypedef¶

items()¶: Yields tuples of name/value exports.

schema() → TreeSpec¶: A schema used to unflatten for access from Python.

Exports a tree of globals into a CompiledModule.

abstractify() → AbstractTypedef¶

items()¶: Yields tuples of name/value exports.

schema() → TreeSpec¶: A schema used to unflatten for access from Python.

Exports parameters from an nn.Module.

These are exposed to procedural programs as a dictionary of param/values.

abstractify_tree()¶

items()¶: Yields tuples of name/value exports.

schema() → TreeSpec¶: A schema used to unflatten for access from Python.

iree.turbine.aot.extend_aot_decompositions(*, from_current: bool = True, add_ops: Sequence[OperatorBase | OpOverloadPacket] | None = None, remove_ops: Sequence[OperatorBase | OpOverloadPacket] | None = None)¶: Context manager which extends the list of decompositions used for AOT.

iree.turbine.aot.externalize_module_parameters(module: Module, *, external_scope: str = '', prefix: str = '')¶: Externalizes parameters and persistent buffers in a module by name.

class iree.turbine.aot.jittable(wrapped_f, *, decompose_ops: List[Any] | None = None, decomposition_table: Dict[Any, Callable[[...], Any]] | None = None, dynamic_shapes: Dict[str, Any] = {}, function_name: str | None = None, passes: Sequence[str] = ('functorch_functionalize',))¶

Decorator which takes a PyTorch function and makes it callable from tracing.

It will be internally JIT-ed and exported into the module as needed.

decomposition_table¶

dynamic_shapes¶

function_name¶

resolve_call(proc_trace: IrTrace, *py_args, dynamic_shapes: Dict[str, Any] = {}, **py_kwargs)¶

wrapped_f¶

iree.turbine.aot.save_module_parameters(file_path: str | Path, module: Module, *, prefix: str = '')¶

One shot save of parameters and persistent buffers on a module.

More options are available by using a ParameterArchiveBuilder.

passes¶

iree.turbine.aot.passes.functorch_functionalize(gm_callable: Any, *args) → GraphModule¶

support¶

class iree.turbine.aot.support.procedural.AbstractIntrinsic¶

Base class for descriptor types that can be converted to Python proxies.

create_intrinsic(value: Value) → Intrinsic¶: Creates a proxy object that can flow through a procedural trace.

get_ir_type(builder: ModuleBuilder) → Type¶: Gets the corresponding IR type.

class iree.turbine.aot.support.procedural.AbstractScalar(label: str, type_producer: Callable[[], Type])¶

Represents a scalar value of some type.

create_intrinsic(ir_value: Value) → Intrinsic¶: Creates a proxy object that can flow through a procedural trace.

get_ir_type(builder: ModuleBuilder) → Type¶: Gets the corresponding IR type.

label¶

type_producer¶

class iree.turbine.aot.support.procedural.AbstractTensor(*size: int | None, dtype: dtype = torch.float32)¶

Represents a tensor of known rank and dtype.

create_intrinsic(ir_value: Value) → Intrinsic¶: Creates a proxy object that can flow through a procedural trace.

dtype¶

get_ir_type(builder: ModuleBuilder) → Type¶: Gets the corresponding IR type.

size¶

class iree.turbine.aot.support.procedural.AbstractTypedef¶

Base class for instances which declare some form of public arg/result type definition.

get_ir_type(builder: ModuleBuilder) → Type¶

class iree.turbine.aot.support.procedural.Abstractifiable¶

Indicates that a type knows how to abstractify itself.

abstractify() → AbstractTypedef¶

class iree.turbine.aot.support.procedural.Any(*args, **kwargs)¶

Special type indicating an unconstrained type.

Any is compatible with every type.
Any assumed to have all methods.
All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class iree.turbine.aot.support.procedural.CallableIntrinsic¶

Intrinsic subclass that supports calls.

This is separate so as to make error handling better (i.e. does not support calls) for intrinsics that are not callable.

class iree.turbine.aot.support.procedural.DeviceAffinity(ordinal: int, queues: list | None = None)¶: This is used to provide device affinities to exported function arguments.

class iree.turbine.aot.support.procedural.DictAttr(*args, **kwargs)¶

attr_name = 'builtin.dictionary'¶

get = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

property type¶: (self) -> iree.compiler._mlir_libs._mlir.ir.Type

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.EmptyType¶

class iree.turbine.aot.support.procedural.F32Type(*args, **kwargs)¶

get = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

type_name = 'builtin.f32'¶

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.F64Type(*args, **kwargs)¶

get = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

type_name = 'builtin.f64'¶

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.FunctionBuilder(*, module_builder: ModuleBuilder, func_op: FuncOp)¶

Helpers for building function bodies.

context¶

emit_return(*ir_values: Value)¶

func_op¶

ip¶

loc¶

module_builder¶

return_types: Sequence[Type] | None¶

class iree.turbine.aot.support.procedural.GlobalAttributes(mutable: bool = False, external: bool | None = None, external_scope: str | None = None, name_mapper: Callable[[str], str | None] | None = None, noinline: bool = False, uninitialized: bool | None = None)¶

Settings for how to initialize the global.

external¶

external_scope¶

infer_external_from_tensor(t: Tensor) → Tuple[bool, str | None, str | None]¶: If externality is not specified, infers it from the tensor.

map_name(name: str) → str¶

mutable¶

name_mapper¶

noinline¶

uninitialized¶

class iree.turbine.aot.support.procedural.GlobalsDef(attrs: GlobalAttributes)¶

Base class for all exporting descriptors.

items() → Generator[Tuple[str, Any], None, None]¶: Yields tuples of name/value exports.

schema() → TreeSpec¶: A schema used to unflatten for access from Python.

track(module_builder: ModuleBuilder, export_namespace: str) → Any¶: Track the given pack of globals, returning a struct that can be used to access them.

class iree.turbine.aot.support.procedural.IREEEmitter¶

tensor_dim(source: IrTensor, index: int, *, dtype: dtype | None = None) → IrScalar¶: Gets the dimension size of a tensor at a static position.

tensor_empty(*dims: int | Value, dtype: dtype = torch.float32) → IrTensor¶

Constructs a tensor with uninitialized values.

TODO: Support an IREE/raw element type in addition to the torch dtype. See: https://github.com/nod-ai/SHARK-ModelDev/issues/130

tensor_reshape(source: IrTensor, *result_dims: int | Value) → IrTensor¶

tensor_slice(source: IrTensor, *indices: Sequence[int | Value | Tuple[int | Value, int | Value]]) → IrTensor¶

Extracts a slice of a tensor.

The given indices must match the rank of the source and each index is interpreted as (start_index[, length]), where the length is taken to be 1 if only a single value is given for an index.

tensor_splat(*dims: int | Value, value: IrScalar | Value, dtype: dtype) → IrTensor¶

tensor_trace(key: str, *ts: IrTensor)¶

tensor_update(target: IrTensor, update: IrTensor, *start_indices: IrScalar | Value | int) → IrTensor¶: Applies an update to a target at start_indices and returns the mutated target.

class iree.turbine.aot.support.procedural.IndexType(*args, **kwargs)¶

get = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

type_name = 'builtin.index'¶

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.IntegerType(*args, **kwargs)¶

SIGNED = 1¶

SIGNLESS = 0¶

class Signedness(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)¶

SIGNED = 1¶

SIGNLESS = 0¶

UNSIGNED = 2¶

UNSIGNED = 2¶

get = <nanobind.nb_func object>¶

get_signed = <nanobind.nb_func object>¶

get_signless = <nanobind.nb_func object>¶

get_unsigned = <nanobind.nb_func object>¶

property is_signed¶: Returns whether this is a signed integer

property is_signless¶: Returns whether this is a signless integer

property is_unsigned¶: Returns whether this is an unsigned integer

property signedness¶: (self) -> iree.compiler._mlir_libs._mlir.ir.IntegerType.Signedness

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

type_name = 'builtin.integer'¶

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

property width¶: Returns the width of the integer type

class iree.turbine.aot.support.procedural.Intrinsic¶

Objects which interact natively with the tracing system implement this.

property ir_value: Value¶

property ir_values: Sequence[Value]¶

resolve_assignment(proc_trace: IrTrace, ir_values: Sequence[Value])¶

resolve_call(proc_trace: IrTrace, *args, **kwargs)¶

resolve_ir_values(proc_trace: IrTrace) → Sequence[Value]¶

class iree.turbine.aot.support.procedural.IrGlobalScalar(export_name: str, info: GlobalsDef, *, symbol_name: str, global_op: Operation, global_type: Type)¶

An IrScalar that is loaded from a global and associated with its aggregate.

export_name¶

global_op: Operation¶

global_type: IrType¶

info¶

resolve_assignment(proc_trace: IrTrace, ir_values: Sequence[Value])¶

resolve_ir_values(trace: IrTrace) → Sequence[Value]¶

set(other)¶

symbol_name: str¶

class iree.turbine.aot.support.procedural.IrGlobalTensor(export_name: str, info: GlobalsDef, *, symbol_name: str, global_op: Operation, global_type: Type, dtype: dtype)¶

An IrScalar that is loaded from a global and associated with its aggregate.

export_name¶

global_op: Operation¶

info¶

resolve_assignment(proc_trace: IrTrace, ir_values: Sequence[Value])¶

resolve_ir_values(trace: IrTrace) → Sequence[Value]¶

symbol_name: str¶

class iree.turbine.aot.support.procedural.IrImmediateScalar(ir_value: Value)¶

Represents an IR scalar value.

resolve_ir_values(proc_trace: IrTrace) → Sequence[Value]¶

class iree.turbine.aot.support.procedural.IrImmediateTensor(ir_value: Value, dtype: dtype)¶

Represents a Value in the IR under construction during procedural tracing.

resolve_ir_values(proc_trace: IrTrace) → Sequence[Value]¶

class iree.turbine.aot.support.procedural.IrScalar(ir_type: Type)¶

An intrinsic that represents a scalar value.

Subclasses are responsible for providing either value or load semantics.

ir_type¶

set(other)¶

class iree.turbine.aot.support.procedural.IrTensor(ir_type: Type, dtype: dtype)¶

An intrinsic that represents a tensor value.

Carries additional metadata needed to resolve dimensions and original PyTorch attributes.

dtype¶

get_dim_value(index: int, *, constant_cache: Dict[int, Value] | None = None, resolved_ir_value: Value | None = None) → Value¶

Gets a dimension as an Index value.

Requires that an InsertionPoint and Location are on the context stack.

This will cache the dim value, returning the cached value later if requested.

ir_type¶

property rank: int¶

set_dynamic_dim_values(values: Sequence[Value])¶: Sets all dynamic dim values.

class iree.turbine.aot.support.procedural.IrTrace(*, module_builder: ModuleBuilder, func_op: FuncOp)¶

Gets callbacks for tracing events.

finalize()¶: Called when the trace is finished (popped off the stack).

handle_assignment(scope, target, updated_value)¶

handle_call(target: Intrinsic, args, kwargs)¶

iree.turbine.aot.support.procedural.IrType¶: alias of Type

class iree.turbine.aot.support.procedural.LiveGlobalCollectionProxy(raw_collection)¶: Proxy object around a collection which knows how to redirect setitem.

class iree.turbine.aot.support.procedural.Location¶

property attr¶: Get the underlying LocationAttr.

property callee¶: Gets the callee location from a CallSiteLoc.

property caller¶: Gets the caller location from a CallSiteLoc.

callsite = <nanobind.nb_func object>¶

property child_loc¶: Gets the child location from a NameLoc.

property context¶: Context that owns the Location.

current = None¶

emit_error¶

Emits an error diagnostic at this location.

Parameters:: message – The error message to emit.

property end_col¶: Gets the end column number from a FileLineColLoc.

property end_line¶: Gets the end line number from a FileLineColLoc.

file = <nanobind.nb_func object>¶

property filename¶: Gets the filename from a FileLineColLoc.

from_attr = <nanobind.nb_func object>¶

fused = <nanobind.nb_func object>¶

is_a_callsite¶: Returns True if this location is a CallSiteLoc.

is_a_file¶: Returns True if this location is a FileLineColLoc.

is_a_fused¶: Returns True if this location is a FusedLoc.

is_a_name¶: Returns True if this location is a NameLoc.

property locations¶: Gets the list of locations from a FusedLoc.

name = <nanobind.nb_func object>¶

property name_str¶: Gets the name string from a NameLoc.

property start_col¶: Gets the start column number from a FileLineColLoc.

property start_line¶: Gets the start line number from a FileLineColLoc.

unknown = <nanobind.nb_func object>¶

class iree.turbine.aot.support.procedural.MaterializedGlobal¶

Tags an Ir* that is duck-typed as a global.

global_op: Operation¶

global_type: Type¶

ir_type: Type¶

symbol_name: str¶

class iree.turbine.aot.support.procedural.ModuleBuilder(module_op: Operation, *, options: ModuleBuilderOptions | None = None)¶

Wrapper around module and IR accounting for a module being built.

body¶

cache¶

context¶

create_func_op(symbol_name: str, argument_types: Sequence[Type], is_public: bool = True, add_entry_block: bool = True, argument_attributes: ArrayAttr | list[DictAttr] | None = None) → Tuple[str, FuncOp]¶

create_tensor_global(symbol_name: str, t: Tensor, *, attrs: GlobalAttributes, logical_name: str | None = None) → Tuple[str, Operation, Type]¶

create_typed_global(symbol_name: str, global_type: Type, *, attrs: GlobalAttributes, logical_name: str | None = None) → Tuple[str, Operation]¶

finalize_construct()¶

fx_py_attr_tracker¶

global_ref_tracker¶

handle_mlir_error(op: Operation, e: MLIRError, message: str)¶

ip¶

last_global_op: Operation | None¶

module_op¶

native_type_converter¶

options¶

symbol_table¶

torch_dtype_to_iree_type(dtype: dtype) → Type¶

unique_auto_symbol(requested_name: str) → str¶

class iree.turbine.aot.support.procedural.Operation¶

property block¶: Returns the block containing this operation.

create = <nanobind.nb_func object>¶

property operation¶: Returns self (the operation).

property opview¶: Returns an OpView of this operation.

Note

If the operation has a registered and loaded dialect then this OpView will be concrete wrapper class.

parse = <nanobind.nb_func object>¶

replace_uses_of_with¶: Replaces uses of the ‘of’ value with the ‘with’ value inside the operation.

property successors¶: Returns the list of Operation successors.

class iree.turbine.aot.support.procedural.ProcedureTrace(*, module_builder: ModuleBuilder, func_op: FuncOp, proxy_posargs, proxy_kwargs)¶

Captures execution of a Python func into IR.

static define_func(module_builder: ModuleBuilder, *, symbol_name: str, posargs: Sequence, kwargs: dict, loc: Location, arg_device: dict[int, DeviceAffinity] | None = None) → ProcedureTrace¶

handle_assignment(scope, target, updated_value)¶

handle_call(target: Intrinsic, args, kwargs)¶: Implements calls to jittable functions.

proxy_kwargs¶

proxy_posargs¶

trace_py_func(py_f: Callable)¶

exception iree.turbine.aot.support.procedural.ProcedureTraceError(message: str)¶

class iree.turbine.aot.support.procedural.RankedTensorType(*args, **kwargs)¶

property encoding¶: (self) -> iree.compiler._mlir_libs._mlir.ir.Attribute | None

get = <nanobind.nb_func object>¶

get_unchecked = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

type_name = 'builtin.tensor'¶

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.ShapedType(*args, **kwargs)¶

property element_type¶: Returns the element type of the shaped type.

get_dim_size¶: Returns the dim-th dimension of the given ranked shaped type.

get_dynamic_size = <nanobind.nb_func object>¶

get_dynamic_stride_or_offset = <nanobind.nb_func object>¶

property has_rank¶: Returns whether the given shaped type is ranked.

property has_static_shape¶: Returns whether the given shaped type has a static shape.

is_dynamic_dim¶: Returns whether the dim-th dimension of the given shaped type is dynamic.

is_dynamic_size = <nanobind.nb_func object>¶

is_dynamic_stride_or_offset¶: Returns whether the given value is used as a placeholder for dynamic strides and offsets in shaped types.

is_static_dim¶: Returns whether the dim-th dimension of the given shaped type is static.

is_static_size = <nanobind.nb_func object>¶

is_static_stride_or_offset¶: Returns whether the given shaped type stride or offset value is statically-sized.

property rank¶: Returns the rank of the given ranked shaped type.

property shape¶: Returns the shape of the ranked shaped type as a list of integers.

property static_typeid¶

object, /) -> iree.compiler._mlir_libs._mlir.ir.TypeID

Type:: (arg

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

class iree.turbine.aot.support.procedural.StringAttr(*args, **kwargs)¶

attr_name = 'builtin.string'¶

get = <nanobind.nb_func object>¶

get_typed = <nanobind.nb_func object>¶

static_typeid = <iree.compiler._mlir_libs._mlir.ir.TypeID object>¶

property type¶: (self) -> iree.compiler._mlir_libs._mlir.ir.Type

property typeid¶: (self) -> iree.compiler._mlir_libs._mlir.ir.TypeID

property value¶: Returns the value of the string attribute

property value_bytes¶: Returns the value of the string attribute as bytes

class iree.turbine.aot.support.procedural.TreeAbstractifiable¶

Indicates that a type decomposes into a tree that can be abstractified.

abstractify_tree() → Any¶

class iree.turbine.aot.support.procedural.TreeSpec(type: Any, context: Any, children_specs: list[Self])¶

child(index: int) → Self¶

children() → list[Self]¶

property children_specs: list[Self]¶

property context: Any¶

flatten_up_to(tree: Any) → list[Any]¶

is_leaf() → bool¶

num_children: int¶

num_leaves: int¶

num_nodes: int¶

type: Any¶

unflatten(leaves: Iterable[Any]) → Any¶

class iree.turbine.aot.support.procedural.Value(*args, **kwargs)¶

property context¶: Context in which the value lives.

dump¶: Dumps a debug representation of the object to stderr.

get_name¶

Overloaded function.

get_name(self, use_local_scope: bool = False, use_name_loc_as_prefix: bool = False) -> str

Returns the string form of value as an operand.

Args:
use_local_scope: Whether to use local scope for naming. use_name_loc_as_prefix: Whether to use the location attribute (NameLoc) as prefix.

Returns:
The value’s name as it appears in IR (e.g., %0, %arg0).
get_name(self, state: iree.compiler._mlir_libs._mlir.ir.AsmState) -> str

Returns the string form of value as an operand (i.e., the ValueID).

property location¶: Returns the source location of the value.

maybe_downcast¶: Downcasts the Value to a more specific kind if possible.

property owner¶: Returns the owner of the value (Operation for results, Block for arguments).

replace_all_uses_except¶: Replace all uses of this value with the with value, except for those in exceptions. exceptions can be either a single operation or a list of operations.

replace_all_uses_with¶: Replace all uses of value with the new value, updating anything in the IR that uses self to use the other value instead.

set_type¶: Sets the type of the value.

property type¶: Returns the type of the value.

property uses¶: Returns an iterator over uses of this value.

iree.turbine.aot.support.procedural.abstractify(tree)¶

iree.turbine.aot.support.procedural.abstractify_single_value(value) → AbstractTypedef¶

iree.turbine.aot.support.procedural.attributes_from_argument_device_affinities(affinities: dict[int, ~iree.turbine.aot.tensor_traits.DeviceAffinity] | None, arguments_count: int, context: ~iree.compiler._mlir_libs._site_initialize.<locals>.Context) → list[dict[str, Attribute]]¶: Get as attributes for function op arguments.

iree.turbine.aot.support.procedural.build_tensor_dim_value(t: Value, dim: int, constant_cache: dict[int, Value] | None = None) → Value¶

iree.turbine.aot.support.procedural.cast(typ, val)¶

Cast a value to a type.

This returns the value unchanged. To the type checker this signals that the return value has the designated type, but at runtime we intentionally don’t check anything (we want this to be as fast as possible).

iree.turbine.aot.support.procedural.contextmanager(func)¶

@contextmanager decorator.

Typical usage:

@contextmanager def some_generator(<arguments>):

<setup> try:

yield <value>

finally:
<cleanup>

This makes this:

with some_generator(<arguments>) as <variable>:
<body>

equivalent to this:

<setup> try:

<variable> = <value> <body>

finally:
<cleanup>

iree.turbine.aot.support.procedural.convert_py_value_to_ir(proc_trace: ProcedureTrace, py_value: Any) → Sequence[Value]¶: Given procedurally traced python values, type check and convert to IR.

iree.turbine.aot.support.procedural.current_ir_trace() → IrTrace¶

iree.turbine.aot.support.procedural.new_ir_trace_scope(ir_trace: IrTrace)¶

iree.turbine.aot.support.procedural.tree_flatten(tree: Any, is_leaf: Callable[[Any], bool] | None = None) → tuple[list[Any], TreeSpec]¶: Flattens a pytree into a list of values and a TreeSpec that can be used to reconstruct the pytree.

iree.turbine.aot.support.procedural.tree_map(func: Callable[[...], Any], tree: Any, *rests: Any, is_leaf: Callable[[Any], bool] | None = None) → Any¶

Map a multi-input function over pytree args to produce a new pytree.

build actions¶

iree.turbine.aot.build_actions.turbine_generate(generator: Callable, *args, name: str, out_of_process: bool = True, **kwargs)¶

Invokes a user-defined generator callable as an action, performing turbine import and storing the resulting artifacts as outputs.

Because torch-based generation is usually quite slow and a bottleneck, this action takes pains to use the out of process action pool, allowing multiple generation activities to take place concurrently. Since this requires interacting with the pickle infrastructure, it puts some constraints on usage:

generator must be a pickleable callable. In practice, this means that it must be a named function at module scope (without decorator) or a named class at module scope with a __call__ method.
args and kwargs must be pickleable. In practice, this means primitive values.

Arguments to the generator are taken from the positional and unmatched keyword arguments passed to turbine_generate.

The generator makes artifacts available as outputs by returning corresponding Python instances (which must be declared as typing parameters for the remoting to work):

ExportOutput: The result of calling aot.export(…) will result in save_mlir() being called on it while still in the subprocess to write to a file names {name}.mlir if there is one return or {name}_{n}.mlir if multiple.

By default, import is run in a subprocess pool. It can be run in the main process by passing out_of_process=False.

See testing/example_builder.py for an example.

class iree.turbine.aot.build_actions.RemoteGenerator(generation_thunk, thunk_args, thunk_kwargs, return_info: list[tuple[ReturnMarshaller, Path]])¶

class iree.turbine.aot.build_actions.TurbineBuilderAction(thunk, thunk_args, thunk_kwargs, concurrency, **kwargs)¶