Skip to content

'passes' Submodule

Introduction

from streamtensor import passes

Passes for the StreamTensor compiler.

This module provides a set of passes for the StreamTensor compiler. Each pass is a function that takes a module and applies a transformation to it. The module is modified in place.

apply_affine_loop_normalize(module)

Apply the affine-loop-normalize pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_affine_scalrep(module)

Apply the affine-scalrep pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_affine_simplify_structures(module)

Apply the affine-simplify-structures pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_comprehensive_bufferize(module)

Apply the streamtensor-comprehensive-bufferize pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_dataflow_to_func(module, create_kernel_wrapper=True)

Apply the streamtensor-convert-dataflow-to-func pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
create_kernel_wrapper bool

Whether to skip the conversion of kernel tasks.

 True

apply_convert_itensor_empty_to_instance(module)

Apply the streamtensor-convert-itensor-empty-to-instance pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_linalg_to_loops(module)

Apply the convert-linalg-to-loops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_memref_to_pointer(module)

Apply the streamtensor-convert-memref-to-pointer pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_tensor_empty_to_instance(module)

Apply the streamtensor-convert-tensor-empty-to-instance pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_tensor_primitive_to_linalg(module)

Apply the streamtensor-convert-tensor-primitive-to-linalg pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_tensor_to_kernel(module)

Apply the streamtensor-convert-tensor-to-kernel pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_convert_tensor_to_linalg(module)

Apply the convert-tensor-to-linalg pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_decompose_tensor_ops(module)

Apply the streamtensor-decompose-tensor-ops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_eliminate_empty_tensors(module)

Apply the eliminate-empty-tensors pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_ensure_itensor_single_use(module)

Apply the streamtensor-ensure-itensor-single-use pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_fold_memref_alias_ops(module)

Apply the fold-memref-alias-ops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_fuse_linalg_fill(module)

Apply the streamtensor-fuse-linalg-fill pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_generate_connectivity(module, num_hbm_port=32)

Apply the streamtensor-generate-connectivity pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
num_hbm_port int

The number of HBM ports.

 32

apply_generate_directive(module, num_uram, num_bram, size_lutram_in_bit)

Apply the streamtensor-generate-directive pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
num_uram int

The number of URAMs.

 required
num_bram int

The number of BRAMs.

 required
size_lutram_in_bit int

The size of LUTRAM in bit.

 required

apply_generate_runtime_host_func(module, host_func_name, timeout_in_ms=1000)

Apply the streamtensor-generate-runtime-host-func pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
host_func_name str

The name of the host function.

 required
timeout_in_ms int

The timeout in milliseconds.

 1000

apply_linalg_fake_quantize(module, quan_bits=8, const_quan_bits=8)

Apply the streamtensor-linalg-fake-quantize pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
quan_bits int

The number of quantization bits.

 8
const_quan_bits int

The number of quantization bits for constants.

 8

apply_linalg_fold_unit_extent_dims(module)

Apply the linalg-fold-unit-extent-dims pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_linalg_fuse_elementwise_ops(module)

Apply the linalg-fuse-elementwise-ops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_linalg_generalize_named_ops(module)

Apply the linalg-generalize-named-ops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_linalg_inline_scalar_operands(module)

Apply the linalg-inline-scalar-operands pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_lower_itensor_to_stream(module)

Apply the streamtensor-lower-itensor-to-stream pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_lower_pack_unpack(module, constant_folding=False)

Apply the streamtensor-lower-pack-unpack pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
constant_folding bool

Whether to perform constant folding.

 False

apply_lower_quantize_ops(module)

Apply the streamtensor-lower-quantize-ops pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_materialize_directive(module)

Apply the streamtensor-materialize-directive pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_materialize_kernel(module)

Apply the streamtensor-materialize-kernel pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_merge_task_chain(module)

Apply the streamtensor-merge-task-chain pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_pack_kernel_interface(module)

Apply the streamtensor-pack-kernel-interface pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_raise_extract_slice_to_chunk(module)

Apply the streamtensor-raise-extract-slice-to-chunk pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_raise_scf_to_affine(module)

Apply the streamtensor-raise-scf-to-affine pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_simplify_task_structure(module)

Apply the streamtensor-simplify-task-structure pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_strip_annotations(module, annotation_name)

Apply the streamtensor-strip-annotations pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
annotation_name str

The name of the annotation to strip.

 required

apply_transform_sequence(module, sequence, delete_sequence=True)

Apply the streamtensor-transform-interpreter pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
sequence NamedSequenceOp

The named transform sequence to apply.

 required
delete_sequence bool

Whether to delete the sequence after applying the pass.

 True

apply_vectorize_itensor(module)

Apply the streamtensor-vectorize-itensor pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required

apply_widen_kernel_interface(module, max_widen_bitwidth=512)

Apply the streamtensor-widen-kernel-interface pass to the module.

Parameters:

Name Type Description Default
module Module

The module to apply the pass to.

 required
max_widen_bitwidth int

The maximum bitwidth to widen to.

 512