mlir.dialects._loop_transform_ops_gen

Attributes

_ods_ir

Classes

ApplyForLoopCanonicalizationPatternsOp	Collects patterns for canonicalizing operations inside SCF loop bodies.
ApplySCFStructuralConversionPatternsOp	Collects patterns for performing structural conversions of SCF operations.
ApplySCFToControlFlowPatternsOp	Collects patterns that lower structured control flow ops to unstructured
ForallToForOp	Converts the scf.forall operation pointed to by the given handle into a
ForallToParallelOp	Converts the scf.forall operation pointed to by the given handle into an
HoistLoopInvariantSubsetsOp	This transform hoists loop-invariant subset ops out of the targeted
LoopCoalesceOp	Given a perfect loop nest identified by the outermost loop,
LoopFuseSiblingOp	Fuses the target loop into the source loop assuming they are
LoopOutlineOp	Moves the loop into a separate function with the specified name and replaces
LoopPeelOp	Rewrite the given loop with a main loop and a partial (first or last) loop.
LoopPipelineOp	Transforms the given loops one by one to achieve software pipelining for
LoopPromoteIfOneIterationOp	Promotes the given target loop op if it has a single iteration. I.e., the
LoopUnrollAndJamOp	Unrolls & jams each loop associated with the given handle to have up to the given
LoopUnrollOp	Unrolls each loop associated with the given handle to have up to the given
ParallelForToNestedForOps	Converts the scf.parallel operation pointed to by the given handle into a
TakeAssumedBranchOp	Given an scf.if conditional, inject user-defined information that it is

Functions

apply_patterns_scf_for_loop_canonicalization(...)
apply_conversion_patterns_scf_structural_conversions(...)
apply_conversion_patterns_scf_scf_to_control_flow(...)
loop_forall_to_for(→ Union[_ods_ir, _ods_ir, ...)
loop_forall_to_parallel(→ Union[_ods_ir, _ods_ir, ...)
loop_hoist_loop_invariant_subsets(...)
loop_coalesce(→ _ods_ir)
loop_fuse_sibling(→ _ods_ir)
loop_outline(→ _ods_ir)
loop_peel(→ _ods_ir)
loop_pipeline(→ _ods_ir)
loop_promote_if_one_iteration(...)
loop_unroll_and_jam(→ LoopUnrollAndJamOp)
loop_unroll(→ LoopUnrollOp)
loop_parallel_for_to_nested_fors(→ Union[_ods_ir, ...)
scf_take_assumed_branch(→ TakeAssumedBranchOp)

Module Contents

mlir.dialects._loop_transform_ops_gen._ods_ir

class mlir.dialects._loop_transform_ops_gen.ApplyForLoopCanonicalizationPatternsOp(*, loc=None, ip=None)

Bases: _ods_ir

Collects patterns for canonicalizing operations inside SCF loop bodies. At the moment, only affine.min/max computations with iteration variables, loop bounds and loop steps are canonicalized.

OPERATION_NAME = 'transform.apply_patterns.scf.for_loop_canonicalization'

_ODS_REGIONS = (0, True)

mlir.dialects._loop_transform_ops_gen.apply_patterns_scf_for_loop_canonicalization(*, loc=None, ip=None) → ApplyForLoopCanonicalizationPatternsOp

class mlir.dialects._loop_transform_ops_gen.ApplySCFStructuralConversionPatternsOp(*, loc=None, ip=None)

Bases: _ods_ir

Collects patterns for performing structural conversions of SCF operations.

OPERATION_NAME = 'transform.apply_conversion_patterns.scf.structural_conversions'

_ODS_REGIONS = (0, True)

mlir.dialects._loop_transform_ops_gen.apply_conversion_patterns_scf_structural_conversions(*, loc=None, ip=None) → ApplySCFStructuralConversionPatternsOp

class mlir.dialects._loop_transform_ops_gen.ApplySCFToControlFlowPatternsOp(*, loc=None, ip=None)

Bases: _ods_ir

Collects patterns that lower structured control flow ops to unstructured control flow.

OPERATION_NAME = 'transform.apply_conversion_patterns.scf.scf_to_control_flow'

_ODS_REGIONS = (0, True)

mlir.dialects._loop_transform_ops_gen.apply_conversion_patterns_scf_scf_to_control_flow(*, loc=None, ip=None) → ApplySCFToControlFlowPatternsOp

class mlir.dialects._loop_transform_ops_gen.ForallToForOp(transformed, target, *, loc=None, ip=None)

Bases: _ods_ir

Converts the scf.forall operation pointed to by the given handle into a set of nested scf.for operations. Each new operation corresponds to one induction variable of the original “multifor” loop.

The operand handle must be associated with exactly one payload operation.

Loops with shared outputs are currently not supported.

Return Modes

Consumes the operand handle. Produces a silenceable failure if the operand is not associated with a single scf.forall payload operation. Returns as many handles as the given forall op has induction variables that are associated with the generated scf.for loops. Produces a silenceable failure if another number of resulting handles is requested.

OPERATION_NAME = 'transform.loop.forall_to_for'

_ODS_REGIONS = (0, True)

target() → _ods_ir

transformed() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_forall_to_for(transformed, target, *, loc=None, ip=None) → _ods_ir | _ods_ir | ForallToForOp

class mlir.dialects._loop_transform_ops_gen.ForallToParallelOp(transformed, target, *, loc=None, ip=None)

Bases: _ods_ir

Converts the scf.forall operation pointed to by the given handle into an scf.parallel operation.

The operand handle must be associated with exactly one payload operation.

Loops with outputs are not supported.

Return Modes

Consumes the operand handle. Produces a silenceable failure if the operand is not associated with a single scf.forall payload operation. Returns a handle to the new scf.parallel operation. Produces a silenceable failure if another number of resulting handles is requested.

OPERATION_NAME = 'transform.loop.forall_to_parallel'

_ODS_REGIONS = (0, True)

target() → _ods_ir

transformed() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_forall_to_parallel(transformed, target, *, loc=None, ip=None) → _ods_ir | _ods_ir | ForallToParallelOp

class mlir.dialects._loop_transform_ops_gen.HoistLoopInvariantSubsetsOp(target, *, loc=None, ip=None)

Bases: _ods_ir

This transform hoists loop-invariant subset ops out of the targeted loop-like op. It looks for matching subset extraction/insertion op pairs and hoists them. The loop body operates on a newly introduced region iter_arg.

Subset ops are hoisted only from the targeted op. If subset ops should be hoisted from an entire loop nest, this transformation must be applied to each loop-like op of the loop nest, starting with the innermost loop and ending with the outermost loop.

Example:

%r = scf.for ... iter_args(%t = %a) -> (tensor<?xf32>) {
  %0 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
  %1 = "test.foo"(%0) : (tensor<5xf32>) -> (tensor<5xf32>)
  %2 = tensor.insert_slice %1 into %t[0][5][1]
      : tensor<5xf32> into tensor<?xf32>
  scf.yield %2 : tensor<?xf32>
}

Is transformed to:

%0 = tensor.extract_slice %a[0][5][1] : tensor<?xf32> to tensor<5xf32>
%new_loop:2 = scf.for ... iter_args(%t = %a, %h = %0) -> (tensor<?xf32>) {
  %1 = "test.foo"(%h) : (tensor<5xf32>) -> (tensor<5xf32>)
  scf.yield %t, %2 : tensor<?xf32>, tensor<5xf32>
}
%r = tensor.insert_slice %new_loop#1 into %new_loop#0
    : tensor<5xf32> into tensor<?xf32>

Subset ops are hoisted only if there are no conflicting subset ops. E.g., if there were a second overlapping extraction in the above example, no ops could be hoisted safely.

This transform reads the target handle and modifies the payload. This transform does not invalidate any handles, but loop-like ops are replaced with new loop-like ops when a subset op is hoisted. The transform rewriter updates all handles accordingly.

OPERATION_NAME = 'transform.loop.hoist_loop_invariant_subsets'

_ODS_REGIONS = (0, True)

target() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_hoist_loop_invariant_subsets(target, *, loc=None, ip=None) → HoistLoopInvariantSubsetsOp

class mlir.dialects._loop_transform_ops_gen.LoopCoalesceOp(transformed, target, *, loc=None, ip=None)

Bases: _ods_ir

Given a perfect loop nest identified by the outermost loop, perform loop coalescing in a bottom-up one-by-one manner.

Return modes

The return handle points to the coalesced loop if coalescing happens, or the given input loop if coalescing does not happen.

OPERATION_NAME = 'transform.loop.coalesce'

_ODS_REGIONS = (0, True)

target() → _ods_ir

transformed() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_coalesce(transformed, target, *, loc=None, ip=None) → _ods_ir

class mlir.dialects._loop_transform_ops_gen.LoopFuseSiblingOp(fused_loop, target, source, *, loc=None, ip=None)

Bases: _ods_ir

Fuses the target loop into the source loop assuming they are independent of each other. In the fused loop, the arguments, body and results of target are placed before those of source.

For fusion of two scf.for loops, the bounds and step size must match. For fusion of two scf.forall loops, the bounds and the mapping must match. Otherwise a silencable failure is produced.

The target and source handles must refer to exactly one operation, otherwise a definite failure is produced. It is the responsibility of the user to ensure that the target and source loops are independent of each other – this op will only perform rudimentary legality checks.

Return modes

This operation consumes the target and source handles and produces the fused_loop handle, which points to the fused loop.

OPERATION_NAME = 'transform.loop.fuse_sibling'

_ODS_REGIONS = (0, True)

target() → _ods_ir

source() → _ods_ir

fused_loop() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_fuse_sibling(fused_loop, target, source, *, loc=None, ip=None) → _ods_ir

class mlir.dialects._loop_transform_ops_gen.LoopOutlineOp(function, call, target, func_name, *, loc=None, ip=None)

Bases: _ods_ir

Moves the loop into a separate function with the specified name and replaces the loop in the Payload IR with a call to that function. Takes care of forwarding values that are used in the loop as function arguments. If the operand is associated with more than one loop, each loop will be outlined into a separate function. The provided name is used as a base for forming actual function names following SymbolTable auto-renaming scheme to avoid duplicate symbols. Expects that all ops in the Payload IR have a SymbolTable ancestor (typically true because of the top-level module).

Return Modes

Returns a handle to the list of outlined functions and a handle to the corresponding function call operations in the same order as the operand handle.

Produces a definite failure if outlining failed for any of the targets.

OPERATION_NAME = 'transform.loop.outline'

_ODS_REGIONS = (0, True)

target() → _ods_ir

func_name() → _ods_ir

function() → _ods_ir

call() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_outline(function, call, target, func_name, *, loc=None, ip=None) → _ods_ir

class mlir.dialects._loop_transform_ops_gen.LoopPeelOp(peeled_loop, remainder_loop, target, *, peel_front=None, fail_if_already_divisible=None, loc=None, ip=None)

Bases: _ods_ir

Rewrite the given loop with a main loop and a partial (first or last) loop. When the peelFront option is set to true, the first iteration is peeled off. Otherwise, updates the given loop so that its step evenly divides its range and puts the remaining iteration into a separate loop or a conditional.

In the absence of sufficient static information, this op may peel a loop, even if the step always divides the range evenly at runtime.

Return modes

This operation ignores non-scf::ForOp ops and drops them in the return. The op returns two loops, the peeled loop which has trip count divisible by the step, and the remainder loop.

When peelFront is true, the first result (remainder loop) executes all but the first iteration of the target loop. The second result (peeled loop) corresponds to the first iteration of the loop which can be canonicalized away in the following optimizations.

When peelFront is false, the first result (peeled loop) is the portion of the target loop with the highest upper bound that is divisible by the step. The second result (remainder loop) contains the remaining iterations.

Note that even though the Payload IR modification may be performed in-place, this operation consumes the operand handle and produces a new one.

Return Modes

Produces a definite failure if peeling fails.

OPERATION_NAME = 'transform.loop.peel'

_ODS_REGIONS = (0, True)

target() → _ods_ir

peel_front() → _ods_ir

fail_if_already_divisible() → _ods_ir

peeled_loop() → _ods_ir

remainder_loop() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_peel(peeled_loop, remainder_loop, target, *, peel_front=None, fail_if_already_divisible=None, loc=None, ip=None) → _ods_ir

class mlir.dialects._loop_transform_ops_gen.LoopPipelineOp(transformed, target, *, iteration_interval=None, read_latency=None, loc=None, ip=None)

Bases: _ods_ir

Transforms the given loops one by one to achieve software pipelining for each of them. That is, performs some amount of reads from memory before the loop rather than inside the loop, the same amount of writes into memory after the loop, and updates each iteration to read the data for a following iteration rather than the current one.

The amount is specified by the attributes.

The values read and about to be stored are transferred as loop iteration arguments. Currently supports memref and vector transfer operations as memory reads/writes.

Return modes

This operation ignores non-scf::For ops and drops them in the return. If all the operations referred to by the target PDLOperation pipeline properly, the transform succeeds. Otherwise the transform produces a silenceable failure. The return handle points to only the subset of successfully produced pipelined loops, which can be empty.

OPERATION_NAME = 'transform.loop.pipeline'

_ODS_REGIONS = (0, True)

target() → _ods_ir

iteration_interval() → _ods_ir

read_latency() → _ods_ir

transformed() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_pipeline(transformed, target, *, iteration_interval=None, read_latency=None, loc=None, ip=None) → _ods_ir

class mlir.dialects._loop_transform_ops_gen.LoopPromoteIfOneIterationOp(target, *, loc=None, ip=None)

Bases: _ods_ir

Promotes the given target loop op if it has a single iteration. I.e., the loop op is removed and only the body remains.

Return modes

This transform fails if the target is mapped to ops that are loops. Ops are considered loops if they implement the LoopLikeOpInterface. Otherwise, this transform always succeeds. The transform consumes the target handle and modifies the payload.

OPERATION_NAME = 'transform.loop.promote_if_one_iteration'

_ODS_REGIONS = (0, True)

target() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_promote_if_one_iteration(target, *, loc=None, ip=None) → LoopPromoteIfOneIterationOp

class mlir.dialects._loop_transform_ops_gen.LoopUnrollAndJamOp(target, factor, *, loc=None, ip=None)

Bases: _ods_ir

Unrolls & jams each loop associated with the given handle to have up to the given number of loop body copies per iteration. If the unroll factor is larger than the loop trip count, the latter is used as the unroll factor instead.

Return modes

This operation ignores non-scf.for, non-affine.for ops and drops them in the return. If all the operations referred to by the target operand unroll properly, the transform succeeds. Otherwise the transform produces a silenceable failure.

Does not return handles as the operation may result in the loop being removed after a full unrolling.

OPERATION_NAME = 'transform.loop.unroll_and_jam'

_ODS_REGIONS = (0, True)

target() → _ods_ir

factor() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_unroll_and_jam(target, factor, *, loc=None, ip=None) → LoopUnrollAndJamOp

class mlir.dialects._loop_transform_ops_gen.LoopUnrollOp(target, factor, *, loc=None, ip=None)

Bases: _ods_ir

Unrolls each loop associated with the given handle to have up to the given number of loop body copies per iteration. If the unroll factor is larger than the loop trip count, the latter is used as the unroll factor instead.

Return modes

This operation ignores non-scf.for, non-affine.for ops and drops them in the return. If all the operations referred to by the target operand unroll properly, the transform succeeds. Otherwise the transform produces a silenceable failure.

Does not return handles as the operation may result in the loop being removed after a full unrolling.

OPERATION_NAME = 'transform.loop.unroll'

_ODS_REGIONS = (0, True)

target() → _ods_ir

factor() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_unroll(target, factor, *, loc=None, ip=None) → LoopUnrollOp

class mlir.dialects._loop_transform_ops_gen.ParallelForToNestedForOps(transformed, target, *, loc=None, ip=None)

Bases: _ods_ir

Converts the scf.parallel operation pointed to by the given handle into a set of nested scf.for operations. Each new operation corresponds to one dimension of the original parallel loop.

The operand handle must be associated with exactly one payload operation.

Loops with shared outputs are currently not supported.

Return Modes

Consumes the operand handle. Produces a silenceable failure if the operand is not associated with a single scf.parallel payload operation. Returns as many handles as the given parallel op has dimensions that are associated with the generated scf.for loops. Produces a silenceable failure if another number of resulting handles is requested.

OPERATION_NAME = 'transform.loop.parallel_for_to_nested_fors'

_ODS_REGIONS = (0, True)

target() → _ods_ir

transformed() → _ods_ir

mlir.dialects._loop_transform_ops_gen.loop_parallel_for_to_nested_fors(transformed, target, *, loc=None, ip=None) → _ods_ir | _ods_ir | ParallelForToNestedForOps

class mlir.dialects._loop_transform_ops_gen.TakeAssumedBranchOp(target, *, take_else_branch=None, loc=None, ip=None)

Bases: _ods_ir

Given an scf.if conditional, inject user-defined information that it is always safe to execute only the if or else branch.

This is achieved by just replacing the scf.if by the content of one of its branches.

This is particularly useful for user-controlled rewriting of conditionals that exist solely to guard against out-of-bounds behavior.

At the moment, no assume or assert operation is emitted as it is not always desirable. In the future, this may be controlled by a dedicated attribute.

Return modes

The transform only consumes its operand and does not produce any result. The transform definitely fails if take_else_branch is specified and the else region is empty.

OPERATION_NAME = 'transform.scf.take_assumed_branch'

_ODS_REGIONS = (0, True)

target() → _ods_ir

take_else_branch() → bool

mlir.dialects._loop_transform_ops_gen.scf_take_assumed_branch(target, *, take_else_branch=None, loc=None, ip=None) → TakeAssumedBranchOp