'utils' Submodule
- Introduction
- annotate
- convert_tiled_linalg_to_kernel
- extract_loop_properties
- find_connectivity
- find_func
- foreach_transform
- fuse_kernels
- i64_attr
- i64_param
- interchange
- match
- match_linalg_init
- match_linalg_input
- match_linalg_result
- match_linalg_with_conditions
- si64_attr
- str_attr
- str_param
- tile
- tile_reduction
- transform_sequence
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 |
convert_tiled_linalg_to_kernel(tiled_linalg_op_handle, parallel_loops_handles, reduction_loops_handles, inputs_vec_sizes, outputs_vec_sizes, 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 |
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 |
kernel_name
|
str
|
The name of the kernel. |
required |
Returns:
Type | Description |
---|---|
DataflowConvertTiledLinalgToKernelOp
|
The |
extract_loop_properties(node)
¶
Extract the loop properties from the given linalg.generic
op.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
node
|
GenericOp
|
The |
required |
Returns:
Type | Description |
---|---|
List[Tuple[int, Optional[str]]]
|
The loop properties of the |
List[Tuple[int, Optional[str]]]
|
tuples containing the loop range and the iterator type ( |
List[Tuple[int, Optional[str]]]
|
|
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 |
required |
Returns:
Type | Description |
---|---|
Optional[ConnectivityOp]
|
The |
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 |
required |
Returns:
Type | Description |
---|---|
Optional[FuncOp]
|
The |
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 |
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 |
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 |
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 |
required |
interchange
|
List[int]
|
The interchange (permutation) map. |
required |
Returns:
Type | Description |
---|---|
InterchangeOp
|
The |
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 |
required |
op_name
|
str
|
The name of the |
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 |
required |
op_name
|
str
|
The name of the |
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 |
required |
op_name
|
str
|
The name of the |
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 |
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 |
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 |
required |
sizes
|
List[int]
|
The tile sizes. |
required |
interchange
|
Optional[List[int]]
|
The interchange (permutation) of the tile. |
None
|
Returns:
Type | Description |
---|---|
TileUsingForOp
|
The |
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 |
required |
sizes
|
List[int]
|
The tile sizes. |
required |
Returns:
Type | Description |
---|---|
TileReductionUsingForOp
|
The |
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_main'
|
result_types
|
List[Type]
|
The return types of the decorated function. |
[]
|
Returns:
Type | Description |
---|---|
Callable[[Callable[..., List[Value]]], Callable[..., NamedSequenceOp]]
|
The constructed |