Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Chapter 4: Lowering to the LLVM Dialect

In Chapter 3 we produced IR in the Kaleidoscope dialect — operations like kaleidoscope.if, kaleidoscope.binop, and kaleidoscope.while that still carry high-level, structured semantics. This chapter converts that IR into the LLVM dialect, where all control flow is flat (basic blocks + branches) and every op corresponds closely to an LLVM instruction.

The implementation for this chapter lives in examples/kaleidoscope/to_llvm.rs.

Design: Dialect Conversion

pliron ships a dialect conversion framework in pliron::irbuild::dialect_conversion. It is deliberately simpler than MLIR’s equivalent:

  • No unrealized conversion casts. Types that do not need conversion are kept as-is.
  • Definitions before uses. The framework guarantees that when an op’s rewrite callback fires, all operands have already been converted.
  • Each op rewrites itself. Rather than a monolithic pattern-match switch, each op implements (in this case) the ToLLVMDialect interface from pliron_llvm.

The three moving parts are:

PieceResponsibility
DialectConversion traitDecides which ops to convert and calls the per-op rewrite
DialectConversionRewriterWraps IRRewriter and records all mutations
apply_dialect_conversionDrives the worklist, ensures def-before-use order

The DialectConversion trait

#![allow(unused)]
fn main() {
pub trait DialectConversion {
    fn can_convert_op(&self, ctx: &Context, op: Ptr<Operation>) -> bool;
    fn rewrite(
        &mut self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        op: Ptr<Operation>,
        operands_info: &OperandsInfo,
    ) -> Result<()>;
    // optional: can_convert_type, convert_type
}
}

can_convert_op is the filter. rewrite is called for every op that passes the filter, with the insertion point already positioned before the op.

OperandsInfo

The operands_info parameter gives each rewrite access to the history of type changes its operands went through during the conversion pass. For the Kaleidoscope lowering we do not convert types (everything stays i64), so this is left unused in most rewrites.

apply_dialect_conversion

#![allow(unused)]
fn main() {
pub fn apply_dialect_conversion<C: DialectConversion>(
    ctx: &mut Context,
    conversion: &mut C,
    op: Ptr<Operation>,
) -> Result<()>
}

The algorithm walks the IR tree rooted at op, collecting convertible ops into a worklist. It then repeatedly pops from the worklist, and ensures to process operand-defining ops first. After each rewrite, recorded mutations are inspected: erased ops are dropped from the worklist, newly inserted ops are added (if required), and block-argument types are updated if any successor references changed.

The conversion driver: KalToLLVM

Our driver is a zero-field struct that implements DialectConversion:

#![allow(unused)]
fn main() {
pub struct KalToLLVM;

impl DialectConversion for KalToLLVM {
    fn can_convert_op(&self, ctx: &Context, op: Ptr<Operation>) -> bool {
        op_impls::<dyn ToLLVMDialect>(&*Operation::get_op_dyn(op, ctx))
            || Operation::get_op::<builtin::ops::FuncOp>(op, ctx).is_some()
    }

    fn rewrite(
        &mut self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        op: Ptr<Operation>,
        operands_info: &OperandsInfo,
    ) -> Result<()> {
        if let Some(func_op) = Operation::get_op::<builtin::ops::FuncOp>(op, ctx) {
            // Convert from builtin.func to llvm.func by updating the function type and argument types.
            return lower_func_op_to_llvm(&func_op, ctx, rewriter);
        }
        let op_dyn = Operation::get_op_dyn(op, ctx);
        let to_llvm_op = op_cast::<dyn ToLLVMDialect>(&*op_dyn)
            .expect("Matched Op must implement ToLLVMDialect");
        to_llvm_op.rewrite(ctx, rewriter, operands_info)
    }
}
}

op_impls::<dyn ToLLVMDialect> checks whether the runtime type of the op object implements the ToLLVMDialect interface. Most conversions go through that path: op_cast downcasts dyn Op to dyn ToLLVMDialect and dispatches to the op’s own rewrite method.

There is one explicit special case: builtin.func is matched and rewritten to llvm.func in the driver itself. We do this because builtin.func is not a Kaleidoscope op and does not implement ToLLVMDialect, but we still need to convert function signatures and block-argument types before lowering function bodies.

Entry point

#![allow(unused)]
fn main() {
/// Lower a [`ModuleOp`] containing Kaleidoscope dialect ops in place.
///
/// Uses the [`DialectConversion`] infrastructure: each Kaleidoscope op
/// implements [`ToLLVMDialect`] and knows how to lower itself to LLVM ops.
pub fn lower_module(ctx: &mut Context, module: ModuleOp) -> Result<IRStatus> {
    apply_dialect_conversion(ctx, &mut KalToLLVM, module.get_operation())
}
}

The module is modified in place: builtin.module and builtin.func are kept as-is because they do not implement ToLLVMDialect. Only the Kaleidoscope ops inside the function bodies are converted.

Lowering simple ops

kaleidoscope.constant -> llvm.constant

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalConstantOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let val_attr = self.value_attr(ctx);
        let llvm_const = LlvmConstantOp::new(ctx, Box::new(val_attr));
        let new_result = llvm_const.get_result(ctx);
        rewriter.insert_op(ctx, &llvm_const);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![new_result]);
        Ok(())
    }
}
}

rewriter.insert_op inserts the new op at the current insertion point (before the op being replaced). replace_operation_with_values replaces every use of the original result with the new result and then erases the original op.

kaleidoscope.decl -> llvm.alloca

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalDeclOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i32_ty = IntegerType::get(ctx, 32, Signedness::Signless);
        let elem_ty = self.variable_type(ctx);
        let size_attr = IntegerAttr::new(i32_ty, APInt::from_i32(1, bw(32)));
        let size_const = LlvmConstantOp::new(ctx, Box::new(size_attr));
        let size_val = size_const.get_result(ctx);
        rewriter.insert_op(ctx, &size_const);
        let alloca = AllocaOp::new(ctx, elem_ty, size_val);
        let alloca_ptr = alloca.get_result(ctx);
        rewriter.insert_op(ctx, &alloca);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![alloca_ptr]);
        Ok(())
    }
}
}

DeclOp allocates a variable slot. The LLVM equivalent is alloca elem_type, i32 1. The alloca instruction requires an i32 count, so a separate llvm.constant 1 : i32 is inserted first.

kaleidoscope.load -> llvm.load

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalLoadOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);
        // `operands_info` carries the already-converted slot operand.
        let slot = operands_info
            .lookup_most_recent_type(self.slot(ctx))
            .map_or(self.slot(ctx), |_| self.slot(ctx));
        // Operand was already updated in-place by the framework.
        let load = LlvmLoadOp::new(ctx, slot, i64_ty.into());
        let result = load.get_result(ctx);
        rewriter.insert_op(ctx, &load);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![result]);
        Ok(())
    }
}
}

kaleidoscope.store -> llvm.store

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalStoreOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let slot = self.slot(ctx);
        let value = self.stored_value(ctx);
        let store = LlvmStoreOp::new(ctx, value, slot);
        rewriter.insert_op(ctx, &store);
        rewriter.erase_operation(ctx, self.get_operation());
        Ok(())
    }
}
}

StoreOp has no result, so we call erase_operation directly instead of replace_operation_with_values.

kaleidoscope.binop -> llvm.add / sub / mul / icmp + sext

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for BinOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);
        let lhs = self.lhs(ctx);
        let rhs = self.rhs(ctx);
        let kind = self.kind(ctx);

        let result = match kind {
            BinOpKind::Add => {
                let op = AddOp::new_with_overflow_flag(
                    ctx,
                    lhs,
                    rhs,
                    IntegerOverflowFlagsAttr::default(),
                );
                let r = op.get_result(ctx);
                rewriter.insert_op(ctx, &op);
                r
            }
            BinOpKind::Sub => {
                let op = SubOp::new_with_overflow_flag(
                    ctx,
                    lhs,
                    rhs,
                    IntegerOverflowFlagsAttr::default(),
                );
                let r = op.get_result(ctx);
                rewriter.insert_op(ctx, &op);
                r
            }
            BinOpKind::Mul => {
                let op = MulOp::new_with_overflow_flag(
                    ctx,
                    lhs,
                    rhs,
                    IntegerOverflowFlagsAttr::default(),
                );
                let r = op.get_result(ctx);
                rewriter.insert_op(ctx, &op);
                r
            }
            _ => {
                // Comparison: ICmpOp yields i1; sign-extend to i64.
                let pred = binop_kind_to_icmp_pred(kind);
                let icmp = ICmpOp::new(ctx, pred, lhs, rhs);
                let cmp_i1 = icmp.get_result(ctx);
                rewriter.insert_op(ctx, &icmp);
                let sext = SExtOp::new(ctx, cmp_i1, i64_ty.into());
                let r = sext.get_result(ctx);
                rewriter.insert_op(ctx, &sext);
                r
            }
        };
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![result]);
        Ok(())
    }
}
}

Arithmetic ops are straightforward one-for-one translations. Comparison ops are two-step: llvm.icmp produces an i1 (1-bit boolean), which is then sign-extended to i64 by llvm.sext so the result type is consistent with the rest of the Kaleidoscope value universe.

kaleidoscope.call -> llvm.call

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalCallOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);
        let callee_attr = self
            .get_attr_callee(ctx)
            .expect("CallOp must have callee attribute")
            .clone();
        let callee_ident = pliron::identifier::Identifier::from(callee_attr);
        let n_args = self.get_operation().deref(ctx).get_num_operands();
        let args: Vec<Value> = (0..n_args)
            .map(|i| self.get_operation().deref(ctx).get_operand(i))
            .collect();
        let arg_types: Vec<Ptr<TypeObj>> = (0..n_args).map(|_| i64_ty.into()).collect();
        let llvm_func_ty = FuncType::get(ctx, i64_ty.into(), arg_types, false);
        let llvm_call = LlvmCallOp::new(
            ctx,
            CallOpCallable::Direct(callee_ident),
            llvm_func_ty,
            args,
        );
        let result = llvm_call.get_result(ctx);
        rewriter.insert_op(ctx, &llvm_call);
        rewriter.replace_operation_with_values(ctx, self.get_operation(), vec![result]);
        Ok(())
    }
}
}

llvm.call requires a FuncType at the call site. Because every Kaleidoscope function takes and returns i64, the callee type is always (i64, …) -> i64 regardless of which function is being called.

kaleidoscope.return -> llvm.return

#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalReturnOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let val = self.value(ctx);
        let ret = LlvmReturnOp::new(ctx, Some(val));
        rewriter.insert_op(ctx, &ret);
        rewriter.erase_operation(ctx, self.get_operation());
        Ok(())
    }
}
}

Lowering control flow

Control flow ops are the most interesting because they carry nested regions that must be flattened into LLVM’s flat CFG. pliron’s DialectConversionRewriter provides the necessary primitives:

MethodEffect
split_block(ctx, block, BeforeOperation(op))Moves ops from op onwards into a fresh successor block; returns the new block
inline_region(ctx, region, AfterBlock(block))Moves all blocks from region into the parent function, inserting after block
create_block(ctx, AfterBlock(b), label, arg_types)Creates a new empty block and inserts it
set_insertion_point(AtBlockEnd(b))Moves the insertion cursor to the end of b

kaleidoscope.if -> conditional CFG

Before conversion:

^entry:
  ... (IfOp with then-region and else-region) ...
  ... (rest of the function) ...

After conversion:

^entry:          ; everything up to IfOp + icmp + cond_br
^then_block:     ; inlined from then-region, ends with br ^merge
^else_block:     ; inlined from else-region, ends with br ^merge
^merge:          ; rest of the function (from split_block)
#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalIfOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);
        let cond = self.condition(ctx);
        let then_region = self.then_region(ctx);
        let else_region = self.else_region(ctx);

        // The entry block of each region is the single block in the region.
        let then_entry = then_region
            .deref(ctx)
            .get_head()
            .expect("IfOp then_region must have a block");
        let else_entry = else_region
            .deref(ctx)
            .get_head()
            .expect("IfOp else_region must have a block");

        let then_term = then_entry
            .deref(ctx)
            .get_terminator(ctx)
            .expect("then block must have a terminator");
        let else_term = else_entry
            .deref(ctx)
            .get_terminator(ctx)
            .expect("else block must have a terminator");

        // Convert the i64 condition to i1 via `icmp ne cond, 0`.
        let zero_attr = IntegerAttr::new(i64_ty, APInt::from_i64(0, bw(64)));
        let zero_const = LlvmConstantOp::new(ctx, Box::new(zero_attr));
        let zero_val = zero_const.get_result(ctx);
        rewriter.insert_op(ctx, &zero_const);
        let cmp = ICmpOp::new(ctx, ICmpPredicateAttr::NE, cond, zero_val);
        let cmp_i1 = cmp.get_result(ctx);
        rewriter.insert_op(ctx, &cmp);

        // Split the current block at the IfOp position to create the merge block.
        let pre_if_block = self
            .get_operation()
            .deref(ctx)
            .get_parent_block()
            .expect("IfOp must be in a block");
        let merge_block = rewriter.split_block(
            ctx,
            pre_if_block,
            OpInsertionPoint::BeforeOperation(self.get_operation()),
            Some("if_merge".try_into().unwrap()),
        );

        // Emit conditional branch in pre_if_block.
        rewriter.set_insertion_point(OpInsertionPoint::AtBlockEnd(pre_if_block));
        let cond_br = CondBrOp::new(ctx, cmp_i1, then_entry, vec![], else_entry, vec![]);
        rewriter.insert_op(ctx, &cond_br);

        // Replace YieldOp in then-branch with branch to merge.
        if Operation::is_op::<KalYieldOp>(then_term, ctx) {
            rewriter.set_insertion_point(OpInsertionPoint::BeforeOperation(then_term));
            let then_br = BrOp::new(ctx, merge_block, vec![]);
            rewriter.insert_op(ctx, &then_br);
            rewriter.erase_operation(ctx, then_term);
        }

        // Replace YieldOp in else-branch with branch to merge.
        if Operation::is_op::<KalYieldOp>(else_term, ctx) {
            rewriter.set_insertion_point(OpInsertionPoint::BeforeOperation(else_term));
            let else_br = BrOp::new(ctx, merge_block, vec![]);
            rewriter.insert_op(ctx, &else_br);
            rewriter.erase_operation(ctx, else_term);
        }

        // Inline both regions after the pre_if_block.
        rewriter.inline_region(
            ctx,
            then_region,
            BlockInsertionPoint::AfterBlock(pre_if_block),
        );
        rewriter.inline_region(
            ctx,
            else_region,
            BlockInsertionPoint::AfterBlock(then_entry),
        );

        // The IfOp itself has no results, so just erase it.
        rewriter.erase_operation(ctx, self.get_operation());
        Ok(())
    }
}
}

The key steps:

  1. split_block cuts the current block at the IfOp, creating merge_block with everything that came after.
  2. icmp ne cond, 0 converts the i64 condition to i1.
  3. cond_br cmp_i1, then_entry, else_entry terminates pre_if_block.
  4. The YieldOp terminators in both branches are replaced with br merge_block.
  5. inline_region moves the then- and else-blocks into the enclosing function.
  6. The IfOp shell (which now has empty regions) is erased.

kaleidoscope.while -> loop CFG

Before conversion:

^entry:
  ... (WhileOp with body-region containing cond updates) ...
  ... (rest of the function) ...

After conversion:

^entry:          ; everything up to WhileOp + br ^header
^while_header:   ; load cond, icmp ne cond, 0, cond_br ^body / ^exit
^body_block:     ; inlined from body-region, ends with br ^header (back-edge)
^exit:           ; rest of the function (from split_block)
#![allow(unused)]
fn main() {
#[op_interface_impl]
impl ToLLVMDialect for KalWhileOp {
    fn rewrite(
        &self,
        ctx: &mut Context,
        rewriter: &mut DialectConversionRewriter,
        _operands_info: &OperandsInfo,
    ) -> Result<()> {
        let i64_ty = IntegerType::get(ctx, 64, Signedness::Signless);
        let cond_ptr = self.cond_ptr(ctx);
        let body_region = self.body_region(ctx);
        let body_entry = body_region
            .deref(ctx)
            .get_head()
            .expect("WhileOp body_region must have a block");

        // Erase the YieldOp at the end of the body.
        let while_term = body_entry
            .deref(ctx)
            .get_terminator(ctx)
            .expect("body block must have a terminator");

        // The block containing the WhileOp is split to create the exit block.
        let pre_while_block = self
            .get_operation()
            .deref(ctx)
            .get_parent_block()
            .expect("WhileOp must be in a block");
        let exit_block = rewriter.split_block(
            ctx,
            pre_while_block,
            OpInsertionPoint::BeforeOperation(self.get_operation()),
            Some("while_exit".try_into().unwrap()),
        );

        // Create the header block.
        let header_block = rewriter.create_block(
            ctx,
            BlockInsertionPoint::AfterBlock(pre_while_block),
            Some("while_header".try_into().unwrap()),
            vec![],
        );

        // Emit an unconditional branch into the header from pre_while_block.
        rewriter.set_insertion_point(OpInsertionPoint::AtBlockEnd(pre_while_block));
        let br_to_header = BrOp::new(ctx, header_block, vec![]);
        rewriter.insert_op(ctx, &br_to_header);

        // Header: load condition, compare to zero, cond_br to body or exit.
        rewriter.set_insertion_point(OpInsertionPoint::AtBlockEnd(header_block));
        let cond_load = LlvmLoadOp::new(ctx, cond_ptr, i64_ty.into());
        let cond_i64 = cond_load.get_result(ctx);
        rewriter.insert_op(ctx, &cond_load);
        let zero_attr = IntegerAttr::new(i64_ty, APInt::from_i64(0, bw(64)));
        let zero_const = LlvmConstantOp::new(ctx, Box::new(zero_attr));
        let zero_val = zero_const.get_result(ctx);
        rewriter.insert_op(ctx, &zero_const);
        let cmp = ICmpOp::new(ctx, ICmpPredicateAttr::NE, cond_i64, zero_val);
        let cmp_i1 = cmp.get_result(ctx);
        rewriter.insert_op(ctx, &cmp);
        let cond_br = CondBrOp::new(ctx, cmp_i1, body_entry, vec![], exit_block, vec![]);
        rewriter.insert_op(ctx, &cond_br);

        // Replace YieldOp at end of body with back-edge to header.
        if Operation::is_op::<KalYieldOp>(while_term, ctx) {
            rewriter.set_insertion_point(OpInsertionPoint::BeforeOperation(while_term));
            let back_edge = BrOp::new(ctx, header_block, vec![]);
            rewriter.insert_op(ctx, &back_edge);
            rewriter.erase_operation(ctx, while_term);
        }

        // Inline the body region after the header block.
        rewriter.inline_region(
            ctx,
            body_region,
            BlockInsertionPoint::AfterBlock(header_block),
        );

        // Erase the WhileOp.
        rewriter.erase_operation(ctx, self.get_operation());
        Ok(())
    }
}
}

WhileOp uses the memory-backed condition pattern introduced in Chapter 3: the loop condition variable is a DeclOp slot in the outer block. The header loads that slot on each iteration and branches accordingly. The body region updates the slot at the end of every iteration, then its YieldOp is replaced with a back-edge branch to the header.

Tests

The test helper parses a Kaleidoscope source string, lowers it to the Kaleidoscope dialect, applies the LLVM lowering in place, then prints and returns the resulting IR:

#![allow(unused)]
fn main() {
    fn lower_to_llvm(src: &str) -> String {
        let funcs = parse_program(src).expect("parse error");
        let ctx = &mut Context::new();
        let module = ModuleOp::new(ctx, "test".try_into().expect("valid module name"));
        for func in &funcs {
            let func_op = lower_function(ctx, func).expect("kaleidoscope lowering failed");
            module.append_operation(ctx, func_op.get_operation(), 0);
        }
        let module_op_ptr = module.get_operation();
        lower_module(ctx, module).expect("LLVM lowering failed");
        verify_operation(module_op_ptr, ctx)
            .expect("module verification failed after LLVM lowering");
        format!("{}", module_op_ptr.disp(ctx))
    }
}

Try it out:

cargo test --example kaleidoscope -- --show-output fibonacci_to_llvm
cargo test --example kaleidoscope -- --show-output factorial_to_llvm
cargo test --example kaleidoscope -- --show-output if_else_to_llvm

Next step

Chapter 5 feeds the LLVM-dialect module into pliron-llvm’s JIT backend to compile and execute the Kaleidoscope program.