tilelang.ascend.transform.z3_scheduler¶

Z3-based scheduler for auto-scheduling.

This module provides a Python implementation of the Z3 scheduler that can be called from C++ via TVM FFI.

Attributes¶

Z3_AVAILABLE

Functions¶

`z3_schedule_python`(latencies, iis, resource_flags, ...)	Schedule a straight-line task list.
`z3_schedule_ffi`(latencies, iis, resource_flags, ...[, ...])	FFI wrapper for z3_schedule_python.
`z3_schedule_loop_python`(num_stages, latencies, iis, ...)	Schedule a loop body with modulo, stage, and storage constraints.
`z3_schedule_loop_ffi`(num_stages, latencies, iis, ...)	FFI wrapper for z3_schedule_loop_python.

Module Contents¶

tilelang.ascend.transform.z3_scheduler.Z3_AVAILABLE = True¶

tilelang.ascend.transform.z3_scheduler.z3_schedule_python(latencies, iis, resource_flags, data_deps, resource_deps, owner_exclusion_deps=None, pipe_order_deps=None, verbose=False)¶

Schedule a straight-line task list.

Parameters:

latencies (list[int]) – Execution latency of each task in cycles.
iis (list[int]) – Initiation interval of each task in cycles.
resource_flags (list[int]) – Resource pipe mask for each task (bitmask of ResourcePipe values): MTE1=1, MTE2=2, MTE3=4, Cube=8, Vector=16, Fixpipe=32, Scalar=64.
data_deps (list[tuple[int, int, int]]) – Directed (i, j, latency) constraints requiring task j to start at least latency cycles after task i.
resource_deps (list[tuple[int, int]]) – Unordered task pairs that cannot issue concurrently on one resource.
owner_exclusion_deps (list[tuple[int, int, int, int]] | None) – Undirected shared-storage owner conflicts (i, j, latency_i_j, latency_j_i). Either owner may run first, but their conflicting storage lifetimes cannot overlap.
pipe_order_deps (list[tuple[int, int]] | None) – Source-ordered task pairs that share a hardware pipe.
verbose (bool) – Print solver inputs and the selected schedule.

Returns:

start_times: Start time for each task sorted_indices: Task indices sorted by start time

Return type:

tuple[list[int], list[int]]

tilelang.ascend.transform.z3_scheduler.z3_schedule_ffi(latencies, iis, resource_flags, data_deps, resource_deps, owner_exclusion_deps=None, pipe_order_deps=None)¶

FFI wrapper for z3_schedule_python.

This function accepts TVM containers and converts them to Python lists.

tilelang.ascend.transform.z3_scheduler.z3_schedule_loop_python(num_stages, latencies, iis, resource_flags, data_deps, resource_deps, owner_exclusion_deps, pipe_order_deps, buffer_sizes, memory_groups, stage_order_deps=None, recalculate_buffer_versions=True, enable_offset=False, manual_stages=None, manual_schedule=False, seed=42, verbose=False)¶

Schedule a loop body with modulo, stage, and storage constraints.

Each start time is represented as k * II + r. Directed dependencies constrain absolute start times, while resource and owner exclusions choose a non-overlapping modulo order. The solver finds the minimum feasible II by binary search.

Parameters:

num_stages (int) – Maximum number of physical versions available to automatic buffers.
latencies (list[int]) – Execution latency of each task in cycles.
iis (list[int]) – Initiation interval of each task in cycles.
resource_flags (list[int]) – Resource pipe mask for each task (bitmask of ResourcePipe values): MTE1=1, MTE2=2, MTE3=4, Cube=8, Vector=16, Fixpipe=32, Scalar=64.
data_deps (list[tuple[int, int, int, int]]) – Directed (i, j, distance, latency) constraints. Negative distances encode automatic buffer-version variables.
resource_deps (list[tuple[int, int]]) – Unordered task pairs that cannot issue concurrently on one resource.
owner_exclusion_deps (list[tuple[int, int, int, int]] | None) – Undirected shared-storage owner conflicts (i, j, latency_i_j, latency_j_i). The chosen modulo order also constrains the reverse wraparound hand-off.
pipe_order_deps (list[tuple[int, int]] | None) – Source-ordered task pairs sharing a hardware pipe.
buffer_sizes (list[int]) – Bytes occupied by one version of each automatic buffer.
memory_groups (list[list[int]]) – Capacity followed by member buffer indices for each memory scope.
stage_order_deps (list[tuple[int, int]] | None) – Stage-order constraints (u, w) requiring k_u <= k_w (same or earlier iteration-stage).
recalculate_buffer_versions (bool) – Recompute the minimum version counts from the selected schedule.
enable_offset (bool) – When False, constrain each resource so that, among the tasks using that resource (pipe bit in resource_flags), max(start_time) - min(start_time) < II. This keeps all tasks on a resource within a single II window (no offsetting a resource’s tasks across pipeline stages). When True, no such constraint is added. Defaults to False (constraint applied).
manual_stages (list[int] | None) – Frontend stages for tasks in source order. Used only when manual_schedule is True; omitted entries are not allowed.
manual_schedule (bool) – Preserve source issue order independently on every hardware pipe and constrain every task’s Z3 stage to the corresponding manual_stages entry.
seed (int | None) – Z3 random seed; None leaves the solver default unchanged.
verbose (bool) – Print solver inputs, search progress, and the selected schedule.

Returns:

start_times: Start time for each task buffer_versions: Number of versions for each buffer minimal_II: The minimal initiation interval found

Return type:

tuple[list[int], list[int], int]

tilelang.ascend.transform.z3_scheduler.z3_schedule_loop_ffi(num_stages, latencies, iis, resource_flags, data_deps, resource_deps, owner_exclusion_deps, pipe_order_deps, buffer_sizes, memory_groups, stage_order_deps=None, enable_offset=False, manual_stages=None, manual_schedule=False)¶

FFI wrapper for z3_schedule_loop_python.

This function accepts TVM containers and converts them to Python lists. Dependency arrays are ordered as directed data dependencies, resource exclusions, shared-storage owner exclusions, and source pipe ordering. memory_groups contains [capacity, idx0, idx1, ...] entries; stage_order_deps contains (u, w) constraints requiring k_u <= k_w. The remaining arguments carry the loop’s offset and manual-stage policy from C++.