Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add InterfacePass to generate pass pipeline #122

Merged
merged 1 commit into from
Feb 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 123 additions & 0 deletions mlir/extras/runtime/passes.py
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,35 @@ def affine_loop_tile(
)
return self

def affine_loop_unroll(
self,
unroll_factor: int = None,
unroll_up_to_factor: bool = None,
unroll_full: bool = None,
unroll_num_reps: int = None,
unroll_full_threshold: int = None,
cleanup_unroll: bool = None,
):
"""Unroll affine loops
Args:
unroll-factor: Use this unroll factor for all loops being unrolled
unroll-up-to-factor: Allow unrolling up to the factor specified
unroll-full: Fully unroll loops
unroll-num-reps: Unroll innermost loops repeatedly this many times
unroll-full-threshold: Unroll all loops with trip count less than or equal to this
cleanup-unroll: Fully unroll the cleanup loop when possible.
"""
self.add_pass(
"affine-loop-unroll",
unroll_factor=unroll_factor,
unroll_up_to_factor=unroll_up_to_factor,
unroll_full=unroll_full,
unroll_num_reps=unroll_num_reps,
unroll_full_threshold=unroll_full_threshold,
cleanup_unroll=cleanup_unroll,
)
return self

def affine_loop_unroll_jam(self, unroll_jam_factor: int = None):
"""Unroll and jam affine loops
Args:
Expand Down Expand Up @@ -1051,6 +1080,21 @@ def control_flow_sink(self):
self.add_pass("control-flow-sink")
return self

def convert_affine_for_to_gpu(
self, gpu_block_dims: int = None, gpu_thread_dims: int = None
):
"""Convert top-level AffineFor Ops to GPU kernels
Args:
gpu-block-dims: Number of GPU block dimensions for mapping
gpu-thread-dims: Number of GPU thread dimensions for mapping
"""
self.add_pass(
"convert-affine-for-to-gpu",
gpu_block_dims=gpu_block_dims,
gpu_thread_dims=gpu_thread_dims,
)
return self

def convert_amdgpu_to_rocdl(self, chipset: str = None):
"""Convert AMDGPU dialect to ROCDL dialect

Expand Down Expand Up @@ -1120,6 +1164,16 @@ def convert_arith_to_spirv(self, emulate_lt_32_bit_scalar_types: bool = None):
)
return self

def convert_arm_sme_to_llvm(self, dump_tile_live_ranges: bool = None):
"""Lower the operations from the ArmSME dialect into the LLVM dialect
Args:
dump-tile-live-ranges: Dump the live ranges of SME tiles (for debugging)
"""
self.add_pass(
"convert-arm-sme-to-llvm", dump_tile_live_ranges=dump_tile_live_ranges
)
return self

def convert_arm_sme_to_scf(self):
"""Lower the operations from the ArmSME dialect into the SCF dialect"""
self.add_pass("convert-arm-sme-to-scf")
Expand Down Expand Up @@ -2289,6 +2343,39 @@ def linalg_block_pack_matmul(
)
return self

def linalg_detensorize(self, aggressive_mode: bool = None):
"""Detensorize linalg ops

Detensoring is the process through which a tensor value is converted to one
or potentially more primitive value(s). During this process, operations with
such detensored operands are also converted to an equivalent form that works
on primitives.

The detensoring process is driven by linalg-on-tensor ops. In particular, a
linalg-on-tensor op is checked to see whether *all* its operands can be
detensored. If so, those operands are converted to their primitive
counterparts and the linalg op is replaced by an equivalent op that takes
those new primitive values as operands. Therefore, detensoring an op can be
divided into 2 main logical phases:

1. Detect/match an op that can be detensored.
2. Detensor the operands of the op and replace it with a primitive
equivalent.

In addition to detensoring individual ops, this pass detensors internal
control flow inside a function. All blocks except for the entry block are
detensored by converting their arguments whenever possible.

This can be run on any FunctionOpInterface op and must not be
run on others. This is because it performs specific legalization of the
blocks that make up the body, which it assumes has is a FunctionOpInterface.

Args:
aggressive-mode: Detensorize all ops that qualify for detensoring along with branch operands and basic-block arguments.
"""
self.add_pass("linalg-detensorize", aggressive_mode=aggressive_mode)
return self

def linalg_fold_unit_extent_dims(self, use_rank_reducing_slices: bool = None):
"""Remove unit-extent dimension in Linalg ops on tensors
Args:
Expand Down Expand Up @@ -4497,6 +4584,42 @@ def tosa_to_arith(
)
return self

def tosa_to_linalg(
self,
disable_tosa_decompositions: bool = None,
aggressive_reduce_constant: bool = None,
):
"""Lower TOSA to LinAlg on tensors

Pass that converts TOSA operations to the equivalent operations using the
tensor operations in LinAlg.

Args:
disable-tosa-decompositions: Disable tosa decompositions pass
aggressive-reduce-constant: Always perform the reduce constant optimization
"""
self.add_pass(
"tosa-to-linalg",
disable_tosa_decompositions=disable_tosa_decompositions,
aggressive_reduce_constant=aggressive_reduce_constant,
)
return self

def tosa_to_linalg_named(self, prefer_conv2d_kernel_layout_hwcf: bool = None):
"""Lower TOSA to LinAlg named operations

Pass that converts TOSA operations to the equivalent operations using the
Linalg named operations.

Args:
prefer-conv2d-kernel-layout-hwcf: Prefer generating linalg.conv_2d_nhwc_hwcf over linalg.conv_2d_nhwc_fhwc
"""
self.add_pass(
"tosa-to-linalg-named",
prefer_conv2d_kernel_layout_hwcf=prefer_conv2d_kernel_layout_hwcf,
)
return self

def tosa_to_mlprogram(self):
"""Lower TOSA to the MLProgram dialect

Expand Down
2 changes: 1 addition & 1 deletion scripts/generate_pass_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ def {pass_name.replace('-', '_')}(self):"""

def gather_passes_from_td_json(j):
passes = []
for pass_ in j["!instanceof"]["Pass"]:
for pass_ in j["!instanceof"]["Pass"] + j["!instanceof"]["InterfacePass"]:
pass_ = j[pass_]
options = []
for o in pass_["options"]:
Expand Down