accelforge.tracegen package#

Submodules#

accelforge.tracegen.accesstrace module#

Trace which tensor elements a LoopTree mapping touches at each timestep.

The entry point is trace_accesses(), which walks a Mapping and returns an AccessTrace: a record of (timestep, tensor element) pairs, plus the timestep spans over which each memory level holds one tile. Feeding that to accelforge.plotting.accesstrace.plot_access_trace() gives the classic “iteration space vs. data space” picture, from which reuse, tiling, and fusion behavior are directly readable.

class accelforge.tracegen.accesstrace.AccessTrace[source]#

Bases: object

A full record of which tensor elements a mapping accesses at each timestep.

A timestep is one iteration of the innermost temporal loop, i.e. one step of the sequential schedule that the LoopTree describes. Spatial loops do not advance the timestep, so spatially-parallel accesses share an x-coordinate.

__init__(traces=<factory>, n_timesteps=0, tensor_shapes=<factory>, tensor_ranks=<factory>, einsum_timespans=<factory>, tile_windows=<factory>, truncated=False)#

einsum_timespans: dict[str, tuple[int, int]]#: {einsum: (first_timestep, last_timestep + 1)} for each Einsum.

property einsums: list[str]#: The names of all traced Einsums, in the order they are first executed.

for_tensor(tensor)[source]#

Return every TensorAccessTrace that touches tensor.

Return type:: list[TensorAccessTrace]

property memory_levels: list[str]#: The names of the storage components the mapping holds tensors in, outermost first.

n_timesteps: int = 0#: The total number of timesteps in the mapping.

tensor_ranks: dict[str, tuple[str, ...]]#: The rank names of each traced tensor, ordered to match tensor_shapes.

tensor_shapes: dict[str, tuple[int, ...]]#: The shape of each traced tensor.

property tensors: list[str]#: The names of all traced tensors, in the order they are first accessed.

tile_lifetime(memory_level, tensor)[source]#

The timestep spans over which one tile of tensor stays resident in memory_level. Empty if that level never holds the tensor.

Return type:: list[tuple[int, int]]

tile_windows: dict[str, dict[str, list[tuple[int, int]]]]#

the timestep spans over which one tile of the tensor stays resident in that memory level. A memory level is the component of a Storage node, and a new span starts whenever a loop outside that node advances.

Type:: {memory_level: {tensor: [(first_timestep, last_timestep + 1), ...]}}

to_dataframe()[source]#: Return a tidy pandas.DataFrame with one row per access and columns einsum, tensor, is_output, timestep, and element.

traces: list[TensorAccessTrace]#: One entry per (Einsum, tensor) pair that the mapping touches.

truncated: bool = False#: Whether the trace was cut short by the max_timesteps argument.

class accelforge.tracegen.accesstrace.TensorAccessTrace[source]#

Bases: object

The accesses that one Einsum makes to one tensor, as a flat list of (timestep, element) pairs.

All arrays share a length: entry i of every array describes the same access. Duplicate accesses within a single timestep are removed, so an element that is read many times in one timestep appears once.

__init__(tensor, einsum, is_output, timestep, element, coordinates, ranks, tensor_shape)#

coordinates: ndarray#: (n_accesses, n_ranks) integer array of per-rank tensor coordinates.

einsum: str#: The name of the Einsum making the accesses.

element: ndarray#: Integer array of flattened (row-major) tensor element indices.

is_output: bool#: Whether this Einsum writes the tensor (True) or reads it (False).

property n_accesses: int#: The number of distinct (timestep, element) accesses in this trace.

ranks: tuple[str, ...]#: The names of the tensor’s ranks, in the order used by coordinates.

tensor: str#: The name of the tensor being accessed.

tensor_shape: tuple[int, ...]#: The size of each of the tensor’s ranks, in the order used by coordinates.

property tensor_size: int#: The number of elements in the tensor.

timestep: ndarray#: Integer array of timesteps at which each access happens.

accelforge.tracegen.accesstrace.trace_accesses(spec_or_mapping, workload=None, tensors=None, einsums=None, max_timesteps=None, max_points=DEFAULT_MAX_POINTS)[source]#

Trace which tensor elements a LoopTree mapping accesses at each timestep.

Parameters:

spec_or_mapping – Either a Spec (whose mapping and workload are used), a Mappings result (whose single mapping is used), or a Mapping. If a bare Mapping is given, workload is required.
workload (Workload | None) – The workload the mapping targets. Required only when spec_or_mapping is a bare Mapping.
tensors (Optional[Iterable[str]]) – If given, only trace these tensors.
einsums (Optional[Iterable[str]]) – If given, only trace these Einsums. Timesteps are still numbered as if all Einsums ran, so traces stay aligned on the x-axis.
max_timesteps (int | None) – If given, trace only the first this-many timesteps. Useful for peeking at the start of a schedule too large to enumerate in full.
max_points (int) – Safety limit on the number of iteration-space points enumerated for a single Einsum. Raises ValueError if exceeded; raise the limit or use max_timesteps to see a prefix instead.

Returns:

The traced accesses, and the tile-residency spans of every memory level the mapping stores tensors in. Pass to accelforge.plotting.accesstrace.plot_access_trace() to visualize.

Return type:

AccessTrace

accelforge.tracegen.tracemapping module#

class accelforge.tracegen.tracemapping.TemporalTrace[source]#

Bases: object

rank_variable: str#

accelforge.tracegen.tracemapping.trace_iterations(mapping, spec)[source]#

Module contents#

class accelforge.tracegen.AccessTrace[source]#

Bases: object

A full record of which tensor elements a mapping accesses at each timestep.

__init__(traces=<factory>, n_timesteps=0, tensor_shapes=<factory>, tensor_ranks=<factory>, einsum_timespans=<factory>, tile_windows=<factory>, truncated=False)#

einsum_timespans: dict[str, tuple[int, int]]#: {einsum: (first_timestep, last_timestep + 1)} for each Einsum.

property einsums: list[str]#: The names of all traced Einsums, in the order they are first executed.

for_tensor(tensor)[source]#

Return every TensorAccessTrace that touches tensor.

Return type:: list[TensorAccessTrace]

property memory_levels: list[str]#: The names of the storage components the mapping holds tensors in, outermost first.

n_timesteps: int = 0#: The total number of timesteps in the mapping.

tensor_ranks: dict[str, tuple[str, ...]]#: The rank names of each traced tensor, ordered to match tensor_shapes.

tensor_shapes: dict[str, tuple[int, ...]]#: The shape of each traced tensor.

property tensors: list[str]#: The names of all traced tensors, in the order they are first accessed.

tile_lifetime(memory_level, tensor)[source]#

The timestep spans over which one tile of tensor stays resident in memory_level. Empty if that level never holds the tensor.

Return type:: list[tuple[int, int]]

tile_windows: dict[str, dict[str, list[tuple[int, int]]]]#

Type:: {memory_level: {tensor: [(first_timestep, last_timestep + 1), ...]}}

to_dataframe()[source]#: Return a tidy pandas.DataFrame with one row per access and columns einsum, tensor, is_output, timestep, and element.

traces: list[TensorAccessTrace]#: One entry per (Einsum, tensor) pair that the mapping touches.

truncated: bool = False#: Whether the trace was cut short by the max_timesteps argument.

class accelforge.tracegen.TensorAccessTrace[source]#

Bases: object

The accesses that one Einsum makes to one tensor, as a flat list of (timestep, element) pairs.

__init__(tensor, einsum, is_output, timestep, element, coordinates, ranks, tensor_shape)#

coordinates: ndarray#: (n_accesses, n_ranks) integer array of per-rank tensor coordinates.

einsum: str#: The name of the Einsum making the accesses.

element: ndarray#: Integer array of flattened (row-major) tensor element indices.

is_output: bool#: Whether this Einsum writes the tensor (True) or reads it (False).

property n_accesses: int#: The number of distinct (timestep, element) accesses in this trace.

ranks: tuple[str, ...]#: The names of the tensor’s ranks, in the order used by coordinates.

tensor: str#: The name of the tensor being accessed.

tensor_shape: tuple[int, ...]#: The size of each of the tensor’s ranks, in the order used by coordinates.

property tensor_size: int#: The number of elements in the tensor.

timestep: ndarray#: Integer array of timesteps at which each access happens.

accelforge.tracegen.trace_accesses(spec_or_mapping, workload=None, tensors=None, einsums=None, max_timesteps=None, max_points=DEFAULT_MAX_POINTS)[source]#

Trace which tensor elements a LoopTree mapping accesses at each timestep.

Parameters:

spec_or_mapping – Either a Spec (whose mapping and workload are used), a Mappings result (whose single mapping is used), or a Mapping. If a bare Mapping is given, workload is required.
workload (Workload | None) – The workload the mapping targets. Required only when spec_or_mapping is a bare Mapping.
tensors (Optional[Iterable[str]]) – If given, only trace these tensors.
einsums (Optional[Iterable[str]]) – If given, only trace these Einsums. Timesteps are still numbered as if all Einsums ran, so traces stay aligned on the x-axis.
max_timesteps (int | None) – If given, trace only the first this-many timesteps. Useful for peeking at the start of a schedule too large to enumerate in full.
max_points (int) – Safety limit on the number of iteration-space points enumerated for a single Einsum. Raises ValueError if exceeded; raise the limit or use max_timesteps to see a prefix instead.

Returns:

The traced accesses, and the tile-residency spans of every memory level the mapping stores tensors in. Pass to accelforge.plotting.accesstrace.plot_access_trace() to visualize.

Return type:

AccessTrace