Skip to content

'utils' Submodule

Introduction

from streamtensor import utils

Utils of the StreamTensor module.

annotate(target, annotation, param=None)

Annotate the target with the given annotation.

Parameters:

Name Type Description Default
target Value

The handle of the target op to annotate.

 required
annotation str

The annotation to add.

 required
param Optional[Any]

The parameter of the annotation.

 None

Returns:

Type Description
AnnotateOp

The transform.annotate op.

convert_tiled_linalg_to_kernel(tiled_linalg_op_handle, parallel_loops_handles, reduction_loops_handles, inputs_vec_sizes, outputs_vec_sizes, max_vec_bitwidth, kernel_name)

Convert the tiled linalg op to the kernel op.

Parameters:

Name Type Description Default
tiled_linalg_op_handle Value

The handle of the tiled linalg op.

 required
parallel_loops_handles List[Value]

The handles of the parallel loops.

 required
reduction_loops_handles List[Value]

The handles of the reduction loops.

 required
inputs_vec_sizes List[List[int]]

The vector sizes of the inputs.

 required
outputs_vec_sizes List[List[int]]

The vector sizes of the outputs.

 required
max_vec_bitwidth int

The maximum vectorize bitwidth.

 required
kernel_name str

The name of the kernel.

 required

Returns:

Type Description
DataflowConvertTiledLinalgToKernelOp

The dataflow.convert_tiled_linalg_to_kernel op.

export_mlir(module, file_object, strip_dense_elements=True)

Export the given module to the given file object in MLIR format.

Parameters:

Name Type Description Default
module

The module to export.

 required
file_object object

The file object to write the MLIR to.

 required
strip_dense_elements bool

Whether to strip dense elements.

 True

extract_loop_properties(node)

Extract the loop properties from the given linalg.generic op.

Parameters:

Name Type Description Default
node GenericOp

The linalg.generic op.

 required

Returns:

Type Description
List[Tuple[int, Optional[str]]]

The loop properties of the linalg.generic op, which is a list of
List[Tuple[int, Optional[str]]]

tuples containing the loop range and the iterator type (reduction or
List[Tuple[int, Optional[str]]]

parallel) of each loop of the linalg.generic op.

find_connectivity(module, name)

Find the hls.connectivity op with the given name in the module.

Parameters:

Name Type Description Default
module Module

The module to search.

 required
name str

The name of the hls.connectivity op to find.

 required

Returns:

Type Description
Optional[ConnectivityOp]

The hls.connectivity op with the given name if found, otherwise None.

find_func(module, name)

Find the func.func op with the given name in the module.

Parameters:

Name Type Description Default
module Module

The module to search.

 required
name str

The name of the func.func op to find.

 required

Returns:

Type Description
Optional[FuncOp]

The func.func op with the given name if found, otherwise None.

flatten_tuple(inputs)

Flatten a nested tuple containing tensors.

Parameters:

Name Type Description Default
inputs Optional[Any]

A tuple that may contain nested tuples and tensors. For example: (a, (b, c)) where a, b, c are tensors.

 required

Returns:

Type Description
Optional[Tuple[Any, ...]]

A flattened tuple containing all tensors in a flat structure.
Optional[Tuple[Any, ...]]

For example: (a, b, c)

foreach_transform(result_types=[])

A decorator to construct a transform.foreach op containing the ops built by the decorated function.

transform.foreach op is used to transform multiple ops contained by a single handle.

The decorated function must have a BlockArgument as its first argument, which is the handle of the target op to be transformed. Any other arguments will be passed transparently when calling the decorated function.

When calling the decorated function, the first argument should be the handle of the target op to be transformed.

Parameters:

Name Type Description Default
result_types List[Type]

The return types of the decorated function.

 []

Returns:

Type Description
Callable[[Callable[..., None]], Callable[..., ForeachOp]]

The constructed transform.foreach op.

fuse_kernels(kernels, kernel_name)

Fuse the given kernels.

Parameters:

Name Type Description Default
kernels List[Value]

The handles of the kernels to fuse.

 required
kernel_name str

The name of the fused kernel.

 required

Returns:

Type Description
DataflowFuseKernelsOp

The dataflow.fuse_kernels op.

i64_attr(value)

Get the 64-bit integer attribute.

Parameters:

Name Type Description Default
value int

The value of the attribute.

 required

Returns:

Type Description
IntegerAttr

The 64-bit integer attribute.

i64_param(value)

Get the 64-bit integer transform.param_constant op.

Parameters:

Name Type Description Default
value int

The value of the constant parameter.

 required

Returns:

Type Description
ParamConstantOp

The 64-bit integer transform.param_constant op.

interchange(linalg_op_handle, interchange)

Interchange the loops of the given linalg op.

Parameters:

Name Type Description Default
linalg_op_handle Value

The handle of the linalg op to interchange.

 required
interchange List[int]

The interchange (permutation) map.

 required

Returns:

Type Description
InterchangeOp

The linalg.interchange op.

match(target, op_names=[], op_attrs=None)

Match the target with the given op names and op attributes.

Parameters:

Name Type Description Default
target Value

The handle of the target op to match.

 required
op_names List[str]

The names of the ops to match.

 []
op_attrs Optional[Dict[str, Attribute]]

The attributes of the ops to match.

 None

Returns:

Type Description
Value

The handle of the matched op.

match_linalg_init(linalg_op_handle, op_name, position=0, is_all=False)

Match the init of the given linalg op.

The returned handle may contain multiple matched inits, which may need to be transformed with foreach_transform decorated function.

Parameters:

Name Type Description Default
linalg_op_handle BlockArgument

The handle of the linalg op to match.

 required
op_name str

The name of the linalg op to match.

 required
position int

The position of the init to match.

 0
is_all bool

Whether to match all inits.

 False

Returns:

Type Description
Value

The handle of the matched init.

match_linalg_input(linalg_op_handle, op_name, position=0, is_all=False)

Match the input of the given linalg op.

The returned handle may contain multiple matched inputs, which may need to be transformed with foreach_transform decorated function.

Parameters:

Name Type Description Default
linalg_op_handle BlockArgument

The handle of the linalg op to match.

 required
op_name str

The name of the linalg op to match.

 required
position int

The position of the input to match.

 0
is_all bool

Whether to match all inputs.

 False

Returns:

Type Description
Value

The handle of the matched input.

match_linalg_result(linalg_op_handle, op_name, position=0)

Match the result of the given linalg op.

The returned handle may contain multiple matched results, which may need to be transformed with foreach_transform decorated function.

Parameters:

Name Type Description Default
linalg_op_handle BlockArgument

The handle of the linalg op to match.

 required
op_name str

The name of the linalg op to match.

 required
position int

The position of the result to match.

 0

Returns:

Type Description
Value

The handle of the matched result.

match_linalg_with_conditions()

A decorator to match a linalg op with conditions built by the decorated function.

The decorated function should have a BlockArgument as its first argument and a str as its second argument, which are the handle of the linalg op to match and the name (e.g., "linalg.generic") of the linalg op to match, respectively. Any other arguments will be passed transparently when calling the decorated function.

When calling the decorated function, the first argument should be the handle of the linalg op to match, and the second argument should be the name of the linalg op to match.

Returns:

Type Description
Callable[[Callable[..., Value]], Callable[..., Value]]

The handle of the matched linalg op.

si64_attr(value)

Get the signed 64-bit integer attribute.

Parameters:

Name Type Description Default
value int

The value of the attribute.

 required

Returns:

Type Description
IntegerAttr

The signed 64-bit integer attribute.

str_attr(value)

Get the string attribute.

Parameters:

Name Type Description Default
value str

The value of the attribute.

 required

Returns:

Type Description
StringAttr

The string attribute.

str_param(value)

Get the string transform.param_constant op.

Parameters:

Name Type Description Default
value str

The value of the constant parameter.

 required

Returns:

Type Description
ParamConstantOp

The string transform.param_constant op.

tile(linalg_op_handle, sizes, interchange=None)

Tile the given linalg op.

Parameters:

Name Type Description Default
linalg_op_handle Value

The handle of the linalg op to tile.

 required
sizes List[int]

The tile sizes.

 required
interchange Optional[List[int]]

The interchange (permutation) of the tile.

 None

Returns:

Type Description
TileUsingForOp

The linalg.tile op.

tile_reduction(linalg_op_handle, sizes)

Tile the given linalg op with reduction.

Parameters:

Name Type Description Default
linalg_op_handle Value

The handle of the linalg op to tile.

 required
sizes List[int]

The tile sizes.

 required

Returns:

Type Description
TileReductionUsingForOp

The linalg.tile_reduction op.

transform_sequence(name='__transform_main', result_types=[])

A decorator to construct a transform.sequence op containing the ops built by the decorated function.

The decorated function should have a BlockArgument as its first argument, which is the handle of the target op to be transformed. Any other arguments will be passed transparently when calling the decorated function.

When calling the decorated function, the first argument should be a Module in which the transform.sequence op is constructed.

Parameters:

Name Type Description Default
name str

The name of the transform.named_sequence op.

 '__transform_main'
result_types List[Type]

The return types of the decorated function.

 []

Returns:

Type Description
Callable[[Callable[..., List[Value]]], Callable[..., NamedSequenceOp]]

The constructed transform.named_sequence op.