The Control Flow Dialect

The control flow dialect provides the most generic control flow semantics via cf.Branch and cf.ConditionalBranch.

cf.Branch

the cf.Branch statement is used to mark how basic block branches to another basic block without condition. This represents an edge on the control flow graph (CFG).

^1(%2):
  │   %y = py.constant.constant 1 : !py.int
  │        cf.br ^3(%y)
  ^2(%3):
  │ %y_1 = py.constant.constant 2 : !py.int
  │        cf.br ^3(%y_1)
  ^3(%y_2):
  │        func.return %y_2

Definition the cf.Branch statement takes a successor block and its argument. The cf.Branch is a terminator thus it should always be the last statement of a block.

cf.ConditionalBranch

The cf.ConditionalBranch statement represents a conditional branching statement that looks like following (the cf.cond_br statement):

^0(%main_self, %x):
│   %0 = py.constant.constant 1 : !py.int
│   %1 = py.cmp.gt(lhs=%x, rhs=%0) : !py.bool
│        cf.cond_br %1 goto ^1(%1) else ^2(%1)

Definition, cf.ConditionalBranch takes a boolean condition cond of type ir.SSAValue and:

• then successor and its argument
• else successor and its argument

this statement is also a terminator, which means it must be the last statement of a block.

Combining together - lowering from Python

Now combining these two statemente together, we can represent most of the Python control flows, e.g if-else and for-loops. These two statement basically just provides a basic way describing the edges on a control flow graph (CFG) by assuming the node only has one or two outgoing edges.

As an example, the following Python program:

from kirin.prelude import basic_no_opt

@basic_no_opt
def main(x):
    if x > 1:
        y = 1
    else:
        y = 2
    return y

will be lowered to the following SSA form in cf dialect:

func.func main(!Any) -> !Any {
  ^0(%main_self, %x):
  │   %0 = py.constant.constant 1 : !py.int
  │   %1 = py.cmp.gt(lhs=%x, rhs=%0) : !py.bool
  │        cf.cond_br %1 goto ^1(%1) else ^2(%1)
  ^1(%2):
  │   %y = py.constant.constant 1 : !py.int
  │        cf.br ^3(%y)
  ^2(%3):
  │ %y_1 = py.constant.constant 2 : !py.int
  │        cf.br ^3(%y_1)
  ^3(%y_2):
  │        func.return %y_2
} // func.func main

And similarly, we can lower a for-loop into the cf dialect:

@basic_no_opt
def main(x):
    for i in range(5):
        x = x + i
    return x

will be lowered into the following SSA form:

func.func main(!Any) -> !Any {
  ^0(%main_self, %x_1):
  │ %0 = py.constant.constant 0 : !py.int
  │ %1 = py.constant.constant 5 : !py.int
  │ %2 = py.constant.constant 1 : !py.int
  │ %3 = py.range.range(start=%0, stop=%1, step=%2) : !py.range
  │ %4 = py.iterable.iter(value=%3) : !Any
  │ %5 = py.constant.constant None : !py.NoneType
  │ %6 = py.iterable.next(iter=%4) : !Any
  │ %7 = py.cmp.is(lhs=%6, rhs=%5) : !py.bool
  │      cf.cond_br %7 goto ^2(%x_1) else ^1(%6, %x_1)
  ^1(%i, %x_2):
  │ %x = py.binop.add(%x_2, %i) : ~T
  │ %8 = py.iterable.next(iter=%4) : !Any
  │ %9 = py.cmp.is(lhs=%8, rhs=%5) : !py.bool
  │      cf.cond_br %9 goto ^2(%x) else ^1(%8, %x)
  ^2(%x_3):
  │      func.return %x_3
} // func.func main

However, as you may already notice, lowering from Python directly to cf dialect will lose some of the high-level information such as the control flow is actually a for-loop. This information can be useful when one wants to perform some optimization. This is why we are taking the same route as MLIR with a structural IR (via ir.Regions). For the interested readers, please proceed to Structural Control Flow for further reading.

API Reference

Branch(*, successor: ir.Block)

arguments: tuple[SSAValue, ...]

name = 'br'

successor: Block = block()

traits = frozenset({IsTerminator()})

print_impl(printer: Printer) -> None

18
19
20
21
22
23
24
25
26
27
28
29

def print_impl(self, printer: Printer) -> None:
    with printer.rich(style="keyword"):
        printer.print_name(self)

    printer.plain_print(" ")
    printer.plain_print(printer.state.block_id[self.successor])
    printer.print_seq(
        self.arguments,
        delim=", ",
        prefix="(",
        suffix=")",
    )

verify() -> None

15
16

def verify(self) -> None:
    return

ConditionalBranch(
    cond: ir.SSAValue,
    *,
    then_successor: ir.Block,
    else_successor: ir.Block
)

cond: SSAValue = argument(Bool)

else_arguments: tuple[SSAValue, ...]

else_successor: Block = block()

name = 'cond_br'

then_arguments: tuple[SSAValue, ...]

then_successor: Block = block()

traits = frozenset({IsTerminator()})

print_impl(printer: Printer) -> None

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

def print_impl(self, printer: Printer) -> None:
    with printer.rich(style="keyword"):
        printer.print_name(self)

    printer.plain_print(" ")
    printer.print(self.cond)

    with printer.rich(style="keyword"):
        printer.plain_print(" goto ")

    printer.plain_print(printer.state.block_id[self.then_successor])
    printer.plain_print("(")
    printer.print_seq(self.then_arguments, delim=", ")
    printer.plain_print(")")

    with printer.rich(style="keyword"):
        printer.plain_print(" else ")

    printer.plain_print(printer.state.block_id[self.else_successor])
    printer.plain_print("(")
    printer.print_seq(self.else_arguments, delim=", ")
    printer.plain_print(")")

verify() -> None

67
68

def verify(self) -> None:
    return