'design_spaces' Submodule

• Introduction
•  BaseDesignSpace
  •  _format_dot_edge_label
  •  _format_dot_edge_style
  •  _format_dot_node_label
  •  _format_dot_node_style
  •  _get_attr
  •  _set_attr
  •  add_node
  •  append_child
  •  get_children
  •  get_id
  •  get_name
  •  get_parent
  •  print_dot
  •  set_children
  •  set_id
  •  set_name
  •  set_parent
  •  verify
•  GraphPath
  •  __getitem__
  •  __iter__
  •  __len__
  •  __reversed__
  •  isdisjoint
•  KernelFusionDesignSpace
  •  greedy_fuse_kernels
  •  print_dot
•  LinalgTilingDesignSpace
  •  get_linalg_op_intensity_aware_unroll_sizes
  •  get_linalg_op_naive_permutation
  •  get_linalg_op_naive_tile_sizes
  •  get_linalg_op_naive_unroll_sizes
  •  get_linalg_op_naive_vec_sizes
  •  intensity_aware_exploration
  •  naive_exploration
  •  print_dot
•  StreamSizingDesignSpace
  •  apply_edge_stream_depths
  •  infer_edge_stream_patterns
  •  solve_edge_stream_depths

Introduction

from streamtensor import design_spaces

Design spaces for the StreamTensor compiler.

This module provides classes for representing design spaces for the StreamTensor compiler. A design space is a directed graph where each node represents a an operation of the IR at a certain level of abstraction.

BaseDesignSpace(module, entry='forward')

Bases: MultiDiGraph

Base class for design spaces.

This class provides a base implementation for design spaces. It is a directed graph where each node represents an operation of the IR at a certain level of abstraction.

Attributes:

Name	Type	Description
module	Module	The MLIR module.
top	FuncOp	The top function in the module.

Parameters:

Name	Type	Description	Default
module	Module	The MLIR module.	required
entry	str	The name of the top function in the module.	'forward'

Raises:

Type	Description
ValueError	If the top function is not found.

_format_dot_edge_label(prev, next, key, params_to_print=None)

Formats the label of an edge for DOT.

This method can be overridden by subclasses to customize the label of an edge.

Parameters:

Name	Type	Description	Default
prev	OpView	The previous node.	required
next	OpView	The next node.	required
key	int	The key of the edge.	required
params_to_print	Optional[List[str]]	The parameters to print.	None

Returns:

Type	Description
str	The formatted label of the edge.

_format_dot_edge_style(prev, next, key)

Formats the style of an edge for DOT.

This method can be overridden by subclasses to customize the style of an edge.

Parameters:

Name	Type	Description	Default
prev	OpView	The previous node.	required
next	OpView	The next node.	required
key	int	The key of the edge.	required

Returns:

Type	Description
Dict[str, str]	A dictionary of style attributes.

_format_dot_node_label(node, params_to_print=None)

Formats the label of a node for DOT.

This method can be overridden by subclasses to customize the label of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
params_to_print	Optional[List[str]]	The parameters to print.	None

Returns:

Type	Description
str	The formatted label of the node.

_format_dot_node_style(node)

Formats the style of a node for DOT.

This method can be overridden by subclasses to customize the style of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required

Returns:

Type	Description
Dict[str, str]	A dictionary of style attributes.

_get_attr(node, attr_name)

Gets an attribute of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
attr_name	str	The name of the attribute.	required

Returns:

Type	Description
Any	The attribute of the node.

Raises:

Type	Description
ValueError	If the attribute is not found.

_set_attr(node, attr_name, attr)

Sets an attribute of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
attr_name	str	The name of the attribute.	required
attr	Any	The attribute of the node.	required

add_node(node, name, id=None, parent=None, children=None, **attr)

Adds a node to the design space.

Parameters:

Name	Type	Description	Default
node	OpView	The node to add.	required
name	str	The name of the node.	required
id	Optional[int]	The ID of the node.	None
parent	Optional[OpView]	The parent of the node.	None
children	Optional[List[OpView]]	The children of the node.	None
attr	Dict[str, Any]	Additional attributes of the node.	{}

append_child(parent, child)

Appends a child to a node.

Parameters:

Name	Type	Description	Default
parent	OpView	The parent node.	required
child	OpView	The child node.	required

get_children(node)

Gets the children of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required

Returns:

Type	Description
Optional[List[OpView]]	The children of the node.

get_id(node)

Gets the ID of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required

Returns:

Type	Description
int	The ID of the node.

get_name(node)

Gets the name of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required

Returns:

Type	Description
str	The name of the node.

get_parent(node)

Gets the parent of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required

Returns:

Type	Description
Optional[OpView]	The parent of the node.

print_dot(file_name, filter=lambda op: True)

Prints the design space as a DOT file.

Parameters:

Name	Type	Description	Default
file_name	str	The name of the DOT file.	required
filter	Callable[[OpView], bool]	A filter function to select nodes to print.	lambda op: True

set_children(node, children)

Sets the children of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
children	Optional[List[OpView]]	The children of the node.	required

set_id(node, id)

Sets the ID of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
id	int	The ID of the node.	required

set_name(node, name)

Sets the name of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
name	str	The name of the node.	required

set_parent(node, parent)

Sets the parent of a node.

Parameters:

Name	Type	Description	Default
node	OpView	The node.	required
parent	Optional[OpView]	The parent of the node.	required

verify()

Verifies the design space.

Raises:

Type	Description
ValueError	If the parent-child relationship is incorrect.

GraphPath(nodes, keys)

A path in a graph.

This class represents a path in a graph. It is a sequence of nodes and keys of the edges between the nodes.

Attributes:

Name	Type	Description
nodes	List[OpView]	The nodes in the path.
internal_nodes_set	set	The set of internal nodes in the path.
keys	List[int]	The keys of the edges in the path.

Parameters:

Name	Type	Description	Default
nodes	List[OpView]	The nodes in the path.	required
keys	List[int]	The keys of the edges in the path.	required

__getitem__(index)

Gets the edge at the specified index.

Parameters:

Name	Type	Description	Default
index	int	The index of the edge.	required

Returns:

Type	Description
Tuple[OpView, OpView, int]	A tuple of the source node, target node, and key.

__iter__()

Iterates over the path.

Returns:

Type	Description
Iterator[Tuple[OpView, OpView, int]]	An iterator of tuples of the source node, target node, and key.

__len__()

Gets the length of the path.

Returns:

Type	Description
int	The length of the path.

__reversed__()

Iterates over the path in reverse.

Returns:

Type	Description
Iterator[Tuple[OpView, OpView, int]]	An iterator of tuples of the target node, source node, and key in reverse order.

isdisjoint(other)

Checks if the path is disjoint with another path.

Parameters:

Name	Type	Description	Default
other	GraphPath	The other path.	required

Returns:

Type	Description
bool	Whether the path is disjoint with the other path.

KernelFusionDesignSpace(module, entry='forward')

Bases: BaseDesignSpace

Design space for kernel fusion.

This class provides a design space for kernel fusion. It is a directed graph where each node represents a kernel op. The exploration results are held in the fused_set_list attribute of this class.

Attributes:

Name	Type	Description
topo_sorted_nodes	List[OpView]	The topologically sorted nodes.
fused_set_list	List[set]	The list of sets of fused nodes.

Parameters:

Name	Type	Description	Default
module	Module	The MLIR module.	required
entry	str	The name of the top function in the module.	'forward'

Raises:

Type	Description
ValueError	If the top function is not found.
ValueError	If the tensor type is not supported.

greedy_fuse_kernels(max_cost=4000)

Performs greedy kernel fusion.

Parameters:

Name	Type	Description	Default
max_cost	int	The maximum fusion cost.	4000

print_dot(file_name)

Prints the design space as a DOT file.

Parameters:

Name	Type	Description	Default
file_name	str	The name of the DOT file.	required

LinalgTilingDesignSpace(module, entry='forward', op_ii_map=None)

Bases: BaseDesignSpace

Design space for linalg tiling.

This class provides a design space for linalg tiling. It is a directed graph where each node represents a linalg op.

After exploration, the following parameters will be added to each node: - parallel_tile_sizes: The tile sizes for parallel loops. - reduction_tile_sizes: The tile sizes for reduction loops. - unroll_sizes: The unroll sizes for loops. - inputs_vec_sizes: The vector sizes for input tensors. - outputs_vec_sizes: The vector sizes for output tensors. - permutation: The permutation of loops.

Attributes:

Name	Type	Description
op_ii_map	Dict[str, int]	A map from operation names to initiation intervals.

Parameters:

Name	Type	Description	Default
module	Module	The MLIR module.	required
entry	str	The name of the top function in the module.	'forward'
op_ii_map	Optional[Dict[str, int]]	A map from operation names to initiation intervals.	None

Raises:

Type	Description
ValueError	If the top function is not found.

get_linalg_op_intensity_aware_unroll_sizes(node, overall_unroll_size, parallel_tile_sizes) staticmethod

Gets the unroll sizes for a linalg operation using an intensity-aware strategy.

Parameters:

Name	Type	Description	Default
node	GenericOp	The linalg operation.	required
overall_unroll_size	int	The overall unroll size.	required
parallel_tile_sizes	List[int]	The parallel tile sizes.	required

Returns:

Type	Description
List[int]	The unroll sizes.

get_linalg_op_naive_permutation(node) staticmethod

Gets the permutation for a linalg operation using a naive strategy.

Parameters:

Name	Type	Description	Default
node	GenericOp	The linalg operation.	required

Returns:

Type	Description
List[int]	The permutation.

get_linalg_op_naive_tile_sizes(node, default_tile_size=16) staticmethod

Gets the tile sizes for a linalg operation using a naive strategy.

Parameters:

Name	Type	Description	Default
node	GenericOp	The linalg operation.	required
default_tile_size	int	The default tile size.	16

Returns:

Type	Description
Tuple[List[int], List[int]]	A tuple of the parallel and reduction tile sizes.

get_linalg_op_naive_unroll_sizes(node, default_unroll_size=2) staticmethod

Gets the unroll sizes for a linalg operation using a naive strategy.

Parameters:

Name	Type	Description	Default
node	GenericOp	The linalg operation.	required
default_unroll_size	int	The default unroll size.	2

Returns:

Type	Description
List[int]	The unroll sizes.

get_linalg_op_naive_vec_sizes(node, unroll_sizes) staticmethod

Gets the vector sizes for a linalg operation using a naive strategy.

Parameters:

Name	Type	Description	Default
node	GenericOp	The linalg operation.	required
unroll_sizes	List[int]	The unroll sizes.	required

Returns:

Type	Description
Tuple[List[List[int]], List[List[int]]]	A tuple of the input and output vector sizes.

intensity_aware_exploration(default_tile_size=16, overall_unroll_size=16)

Performs intensity-aware exploration of the design space.

Parameters:

Name	Type	Description	Default
default_tile_size	int	The default tile size for linalg tiling.	16
overall_unroll_size	int	The overall unroll size for linalg tiling.	16

naive_exploration(default_tile_size=16, default_unroll_size=2)

Performs naive exploration of the design space.

Parameters:

Name	Type	Description	Default
default_tile_size	int	The default tile size for linalg tiling.	16
default_unroll_size	int	The default unroll size for linalg tiling.	2

print_dot(file_name)

Prints the design space as a DOT file.

Parameters:

Name	Type	Description	Default
file_name	str	The name of the DOT file.	required

StreamSizingDesignSpace(module, entry='forward', resolve_producer_consumer=False)

Bases: BaseDesignSpace

Design space for stream sizing.

This class provides a design space for stream sizing. It is a directed graph where each node represents a task op. To explore the design space, the resolve_producer_consumer attribute should be set to False, which means that the producer and consumer are the top-level tasks in the task graph instead of the actual leaf producer and consumer tasks.

Attributes:

Name	Type	Description
resolve_producer_consumer	bool	Whether to resolve producer-consumer relationships.

Parameters:

Name	Type	Description	Default
module	Module	The MLIR module.	required
entry	str	The name of the top function in the module.	'forward'
resolve_producer_consumer	bool	Whether to resolve producer-consumer relationships for printing DOT.	False

Raises:

Type	Description
AssertionError	Non-task ops are found in the task graph.

apply_edge_stream_depths()

Applies edge stream depths to the design space.

infer_edge_stream_patterns()

Infers edge stream patterns for the design space.

The edge stream pattern is a tuple of (token, source_slope, target_slope, source_delay). The token is the number of tokens that are transferred in the stream channel. The source slope is producing rate of the stream producer. The target slope is consuming rate of the stream consumer. The source delay is the delay from the first input token to the first output token of the stream producer.

Raises:

Type	Description
AssertionError	The producer-consumer relationship is resolved.

solve_edge_stream_depths(pipeline_rewinding=True, stream_size_map=None, stream_full_size=False)

Solves edge stream depths for the design space.

The edge stream depth is the max number of tokens that can be stored in the stream channel. Essentially, this method first solve the delay of each edge through linear programming, and then calculates the depth of each edge based on the delay.

Parameters:

Name	Type	Description	Default
pipeline_rewinding	bool	Whether to enable pipeline rewinding.	True
stream_size_map	Optional[Dict[Tuple[str, str], int]]	The overriding map from edge to stream size.	None
stream_full_size	bool	Set all stream sizes to the maximum (token number).	False

Raises:

Type	Description
Exception	The solver did not find an optimal solution.