Skip to main content

/

Docs

Nightly
v24.5
Log in

Mojo function

unswitch

unswitch[: origin.set, //, switched_func: fn[Bool]() raises capturing -> None](dynamic_switch: Bool)

Performs a functional unswitch transformation.

Unswitch is a simple pattern that is similar idea to loop unswitching pass but extended to functional patterns. The pattern facilitates the following code transformation that reduces the number of branches in the generated code

Before:

for i in range(...)
    if i < xxx:
        ...
for i in range(...)
    if i < xxx:
        ...

After:

if i < ...
    for i in range(...)
        ...
else
    for i in range(...)
        if i < xxx:
            ...
if i < ...
    for i in range(...)
        ...
else
    for i in range(...)
        if i < xxx:
            ...

This unswitch function generalizes that pattern with the help of meta parameters and can be used to perform both loop unswitching and other tile predicate lifting like in simd and amx.

TODO: Generalize to support multiple predicates. TODO: Once nested lambdas compose well should make unswitch compose with tile in an easy way.

Parameters:

  • ​switched_func (fn[Bool]() raises capturing -> None): The function containing the inner loop logic that can be unswitched.

Args:

  • ​dynamic_switch (Bool): The dynamic condition that enables the unswitched code path.

unswitch[: origin.set, //, switched_func: fn[Bool]() capturing -> None](dynamic_switch: Bool)

Performs a functional unswitch transformation.

Unswitch is a simple pattern that is similar idea to loop unswitching pass but extended to functional patterns. The pattern facilitates the following code transformation that reduces the number of branches in the generated code

Before:

for i in range(...)
    if i < xxx:
        ...
for i in range(...)
    if i < xxx:
        ...

After:

if i < ...
    for i in range(...)
        ...
else
    for i in range(...)
        if i < xxx:
            ...
if i < ...
    for i in range(...)
        ...
else
    for i in range(...)
        if i < xxx:
            ...

This unswitch function generalizes that pattern with the help of meta parameters and can be used to perform both loop unswitching and other tile predicate lifting like in simd and amx.

TODO: Generalize to support multiple predicates. TODO: Once nested lambdas compose well should make unswitch compose with tile in an easy way.

Parameters:

  • ​switched_func (fn[Bool]() capturing -> None): The function containing the inner loop logic that can be unswitched.

Args:

  • ​dynamic_switch (Bool): The dynamic condition that enables the unswitched code path.

unswitch[: origin.set, //, switched_func: fn[Bool, Bool]() capturing -> None](dynamic_switch_a: Bool, dynamic_switch_b: Bool)

Performs a functional 2-predicates unswitch transformation.

Parameters:

  • ​switched_func (fn[Bool, Bool]() capturing -> None): The function containing the inner loop logic that has 2 predicates which can be unswitched.

Args:

  • ​dynamic_switch_a (Bool): The first dynamic condition that enables the outer unswitched code path.
  • ​dynamic_switch_b (Bool): The second dynamic condition that enables the inner unswitched code path.