Visualize ControlFlow.ipynb

Kernel: Python 3

from llvmlite import ir

Build a simple function with a for loop that sum the first N integer, where N is the argument.

m = ir.Module()

fnty = ir.FunctionType(ir.IntType(32), [ir.IntType(32)])
print(fnty)

i32 (i32)

fn = ir.Function(m, fnty, "count_number")
fn.args[0].name= 'N'
print(fn)

declare i32 @"count_number"(i32 %"N") 

builder = ir.IRBuilder(fn.append_basic_block('entry'))

Initialize variable ct and out with 0

out = builder.alloca(ir.IntType(32), name='out')
ct = builder.alloca(ir.IntType(32), name='ct')
builder.store(out.type.pointee(0), out)
builder.store(ct.type.pointee(0), ct)
print(fn)

define i32 @"count_number"(i32 %"N") 
{
entry:
  %"out" = alloca i32
  %"ct" = alloca i32
  store i32 0, i32* %"out"
  store i32 0, i32* %"ct"
}

Populate loop header

loophead = fn.append_basic_block('loop.header')
loopbody = fn.append_basic_block('loop.body')
loopend = fn.append_basic_block('loop.end')

builder.branch(loophead)
builder.position_at_end(loophead)

# loop if ct < arg0
arg0 = fn.args[0]
pred = builder.icmp_signed('<', builder.load(ct), arg0)
builder.cbranch(pred, loopbody, loopend)

print(fn)

define i32 @"count_number"(i32 %"N") 
{
entry:
  %"out" = alloca i32
  %"ct" = alloca i32
  store i32 0, i32* %"out"
  store i32 0, i32* %"ct"
  br label %"loop.header"
loop.header:
  %".6" = load i32, i32* %"ct"
  %".7" = icmp slt i32 %".6", %"N"
  br i1 %".7", label %"loop.body", label %"loop.end"
loop.body:
loop.end:
}

Populate loop body

builder.position_at_end(loopbody)

# out += ct
builder.store(builder.add(builder.load(out), builder.load(ct)), out)
# ct += 1
builder.store(builder.add(builder.load(ct), ct.type.pointee(1)), ct)
# jump to loophead
builder.branch(loophead)

print(fn)

define i32 @"count_number"(i32 %"N") 
{
entry:
  %"out" = alloca i32
  %"ct" = alloca i32
  store i32 0, i32* %"out"
  store i32 0, i32* %"ct"
  br label %"loop.header"
loop.header:
  %".6" = load i32, i32* %"ct"
  %".7" = icmp slt i32 %".6", %"N"
  br i1 %".7", label %"loop.body", label %"loop.end"
loop.body:
  %".9" = load i32, i32* %"out"
  %".10" = load i32, i32* %"ct"
  %".11" = add i32 %".9", %".10"
  store i32 %".11", i32* %"out"
  %".13" = load i32, i32* %"ct"
  %".14" = add i32 %".13", 1
  store i32 %".14", i32* %"ct"
  br label %"loop.header"
loop.end:
}

Populate loop end

builder.position_at_end(loopend)

builder.ret(builder.load(out))

print(fn)

define i32 @"count_number"(i32 %"N") 
{
entry:
  %"out" = alloca i32
  %"ct" = alloca i32
  store i32 0, i32* %"out"
  store i32 0, i32* %"ct"
  br label %"loop.header"
loop.header:
  %".6" = load i32, i32* %"ct"
  %".7" = icmp slt i32 %".6", %"N"
  br i1 %".7", label %"loop.body", label %"loop.end"
loop.body:
  %".9" = load i32, i32* %"out"
  %".10" = load i32, i32* %"ct"
  %".11" = add i32 %".9", %".10"
  store i32 %".11", i32* %"out"
  %".13" = load i32, i32* %"ct"
  %".14" = add i32 %".13", 1
  store i32 %".14", i32* %"ct"
  br label %"loop.header"
loop.end:
  %".17" = load i32, i32* %"out"
  ret i32 %".17"
}

Visualize CFG

from llvmlite import binding as llvm

dot = llvm.get_function_cfg(fn)
llvm.view_dot_graph(dot)

Optimize the function

# materialize a LLVM module
mod = llvm.parse_assembly(str(m))

# create optimizer
llvm.initialize_native_target()
target_machine = llvm.Target.from_default_triple().create_target_machine()
pto = llvm.create_pipeline_tuning_options(speed_level=3)
pb = llvm.create_pass_builder(target_machine, pto)
pm = pb.getModulePassManager()

# optimize
pm.run(mod, pb)

print(mod)

; ModuleID = '<string>'
target triple = "unknown-unknown-unknown"

; Function Attrs: norecurse nounwind readnone
define i32 @count_number(i32 %N) #0 {
entry:
  %.71 = icmp sgt i32 %N, 0
  br i1 %.71, label %loop.body.preheader, label %loop.end

loop.body.preheader:                              ; preds = %entry
  %0 = add i32 %N, -1
  %1 = zext i32 %0 to i33
  %2 = add i32 %N, -2
  %3 = zext i32 %2 to i33
  %4 = mul i33 %1, %3
  %5 = lshr i33 %4, 1
  %6 = trunc i33 %5 to i32
  %7 = add i32 %6, %N
  %8 = add i32 %7, -1
  br label %loop.end

loop.end:                                         ; preds = %loop.body.preheader, %entry
  %out.0.lcssa = phi i32 [ 0, %entry ], [ %8, %loop.body.preheader ]
  ret i32 %out.0.lcssa
}

attributes #0 = { norecurse nounwind readnone }

View optimized CFG

dot = llvm.get_function_cfg(mod.get_function(fn.name))
llvm.view_dot_graph(dot)

Notice LLVM has optimized the looping away.