Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(llvm): Tail loop emission #1749

Merged
merged 7 commits into from
Dec 11, 2024
+767 −9
Merged

feat(llvm): Tail loop emission #1749

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
595e63f
tests: Add tail loop emission insta tests
croyzor Dec 10, 2024
86190ee
typos
croyzor Dec 10, 2024
c3ddd38
WIP: Tail loop emission
croyzor Dec 10, 2024
56942ce
tests pass
doug-q Dec 11, 2024
87c48ed
fmt
doug-q Dec 11, 2024
7bcb9eb
add tail loop execution test
ss2165 Dec 11, 2024
2d6ac0c
remove comments and print
ss2165 Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion hugr-llvm/src/emit/func.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,7 @@ impl<'c, 'a, H: HugrView> EmitFuncContext<'c, 'a, H> {
Ok(r)
}

/// Returns a [RowMailBox] mapped to thie ouput wires of `node`. When emitting a node
/// Returns a [RowMailBox] mapped to the output wires of `node`. When emitting a node
/// output values are written to this mailbox.
pub fn node_outs_rmb<'hugr, OT: 'hugr>(
&mut self,
Expand Down
15 changes: 14 additions & 1 deletion hugr-llvm/src/emit/func/mailbox.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::{borrow::Cow, rc::Rc};

use anyhow::Result;
use anyhow::{bail, Result};
use delegate::delegate;
use inkwell::{
builder::Builder,
Expand Down Expand Up @@ -148,6 +148,19 @@ impl<'c> RowMailBox<'c> {
builder: &Builder<'c>,
vs: impl IntoIterator<Item = BasicValueEnum<'c>>,
) -> Result<()> {
let vs = vs.into_iter().collect_vec();
#[cfg(debug_assertions)]
{
let actual_types = vs.clone().into_iter().map(|x| x.get_type()).collect_vec();
let expected_types = self.get_types().collect_vec();
if actual_types != expected_types {
bail!(
"RowMailbox::write: Expected types {:?}, got {:?}",
expected_types,
actual_types
);
}
}
zip_eq(self.0.iter(), vs).try_for_each(|(mb, v)| mb.write(builder, v))
}

Expand Down
59 changes: 56 additions & 3 deletions hugr-llvm/src/emit/ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use anyhow::{anyhow, bail, Result};
use hugr_core::ops::{
constant::Sum, Call, CallIndirect, Case, Conditional, Const, ExtensionOp, Input, LoadConstant,
LoadFunction, OpTag, OpTrait, OpType, Output, Tag, Value, CFG,
LoadFunction, OpTag, OpTrait, OpType, Output, Tag, TailLoop, Value, CFG,
};
use hugr_core::{
hugr::views::SiblingGraph,
Expand Down Expand Up @@ -58,7 +58,7 @@ where
.ok_or(anyhow!("DataflowParentEmitter: Output taken twice"))
}

pub fn emit_children(mut self, context: &mut EmitFuncContext<'c, '_, H>) -> Result<()> {
pub fn emit_children(&mut self, context: &mut EmitFuncContext<'c, '_, H>) -> Result<()> {
use petgraph::visit::Topo;
let node = self.node;
if !OpTag::DataflowParent.is_superset(node.tag()) {
Expand Down Expand Up @@ -306,6 +306,59 @@ fn emit_cfg<'c, H: HugrView>(
cfg::CfgEmitter::new(context, args)?.emit_children(context)
}

fn emit_tail_loop<'c, H: HugrView>(
context: &mut EmitFuncContext<'c, '_, H>,
args: EmitOpArgs<'c, '_, TailLoop, H>,
) -> Result<()> {
let node = args.node();

// Make a block to jump to when we `Break`
let out_bb = context.new_basic_block("loop_out", None);
// A block for the body of the loop
let body_bb = context.new_basic_block("loop_body", Some(out_bb));

let (body_i_node, body_o_node) = node.get_io().unwrap();
let body_i_rmb = context.node_outs_rmb(body_i_node)?;
let body_o_rmb = context.node_ins_rmb(body_o_node)?;

body_i_rmb.write(context.builder(), args.inputs)?;
context.builder().build_unconditional_branch(body_bb)?;

let control_llvm_sum_type = {
let sum_ty = SumType::new([node.just_inputs.clone(), node.just_outputs.clone()]);
context.llvm_sum_type(sum_ty)?
};

context.build_positioned(body_bb, move |context| {
let inputs = body_i_rmb.read_vec(context.builder(), [])?;
emit_dataflow_parent(
context,
EmitOpArgs {
node,
inputs,
outputs: body_o_rmb.promise(),
},
)?;
let dataflow_outputs = body_o_rmb.read_vec(context.builder(), [])?;
let control_val = LLVMSumValue::try_new(dataflow_outputs[0], control_llvm_sum_type)?;
let mut outputs = Some(args.outputs);

control_val.build_destructure(context.builder(), |builder, tag, mut values| {
values.extend(dataflow_outputs[1..].iter().copied());
if tag == 0 {
body_i_rmb.write(builder, values)?;
builder.build_unconditional_branch(body_bb)?;
} else {
outputs.take().unwrap().finish(builder, values)?;
builder.build_unconditional_branch(out_bb)?;
}
Ok(())
})
})?;
context.builder().position_at_end(out_bb);
Ok(())
}

fn emit_optype<'c, H: HugrView>(
context: &mut EmitFuncContext<'c, '_, H>,
args: EmitOpArgs<'c, '_, OpType, H>,
Expand All @@ -330,7 +383,7 @@ fn emit_optype<'c, H: HugrView>(
context.push_todo_func(node.into_ot(fd));
Ok(())
}

OpType::TailLoop(x) => emit_tail_loop(context, args.into_ot(x)),
_ => Err(anyhow!("Invalid child for Dataflow Parent: {node}")),
}
}
Expand Down
61 changes: 61 additions & 0 deletions hugr-llvm/src/emit/snapshots/[email protected]
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
---
source: hugr-llvm/src/emit/test.rs
expression: mod_str
---
; ModuleID = 'test_context'
source_filename = "test_context"

define i64 @_hl.main.1() {
alloca_block:
br label %entry_block

entry_block: ; preds = %alloca_block
br label %loop_body

loop_body: ; preds = %12, %entry_block
%"9_0.0" = phi i64 [ 3, %entry_block ], [ %14, %12 ]
%"9_1.0" = phi i64 [ 7, %entry_block ], [ %0, %12 ]
%0 = mul i64 %"9_1.0", 2
%1 = icmp eq i64 %"9_0.0", 0
%2 = select i1 %1, { i32, {}, {} } { i32 1, {} poison, {} undef }, { i32, {}, {} } { i32 0, {} undef, {} poison }
%3 = extractvalue { i32, {}, {} } %2, 0
switch i32 %3, label %4 [
i32 1, label %6
]

4: ; preds = %loop_body
%5 = extractvalue { i32, {}, {} } %2, 1
br label %cond_17_case_0

6: ; preds = %loop_body
%7 = extractvalue { i32, {}, {} } %2, 2
br label %cond_17_case_1

loop_out: ; preds = %15
ret i64 %0

cond_17_case_0: ; preds = %4
%8 = sub i64 %"9_0.0", 1
%9 = insertvalue { i64 } undef, i64 %8, 0
%10 = insertvalue { i32, { i64 }, {} } { i32 0, { i64 } poison, {} poison }, { i64 } %9, 1
br label %cond_exit_17

cond_17_case_1: ; preds = %6
br label %cond_exit_17

cond_exit_17: ; preds = %cond_17_case_1, %cond_17_case_0
%"011.0" = phi { i32, { i64 }, {} } [ %10, %cond_17_case_0 ], [ { i32 1, { i64 } poison, {} undef }, %cond_17_case_1 ]
%11 = extractvalue { i32, { i64 }, {} } %"011.0", 0
switch i32 %11, label %12 [
i32 1, label %15
]

12: ; preds = %cond_exit_17
%13 = extractvalue { i32, { i64 }, {} } %"011.0", 1
%14 = extractvalue { i64 } %13, 0
br label %loop_body

15: ; preds = %cond_exit_17
%16 = extractvalue { i32, { i64 }, {} } %"011.0", 2
br label %loop_out
}
129 changes: 129 additions & 0 deletions ...lvm/src/emit/snapshots/hugr_llvm__emit__test__test_fns__tail_loop@[email protected]
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
---
source: hugr-llvm/src/emit/test.rs
expression: mod_str
---
; ModuleID = 'test_context'
source_filename = "test_context"

define i64 @_hl.main.1() {
alloca_block:
%"0" = alloca i64, align 8
%"7_0" = alloca i64, align 8
%"5_0" = alloca i64, align 8
%"8_0" = alloca i64, align 8
%"9_0" = alloca i64, align 8
%"9_1" = alloca i64, align 8
%"17_0" = alloca { i32, { i64 }, {} }, align 8
%"16_0" = alloca i64, align 8
%"15_0" = alloca i64, align 8
%"12_0" = alloca i64, align 8
%"13_0" = alloca { i32, {}, {} }, align 8
%"011" = alloca { i32, { i64 }, {} }, align 8
%"013" = alloca i64, align 8
%"23_0" = alloca i64, align 8
%"20_0" = alloca i64, align 8
%"24_0" = alloca i64, align 8
%"25_0" = alloca { i32, { i64 }, {} }, align 8
%"019" = alloca i64, align 8
%"29_0" = alloca { i32, { i64 }, {} }, align 8
%"27_0" = alloca i64, align 8
br label %entry_block

entry_block: ; preds = %alloca_block
store i64 7, i64* %"7_0", align 4
store i64 3, i64* %"5_0", align 4
%"5_01" = load i64, i64* %"5_0", align 4
%"7_02" = load i64, i64* %"7_0", align 4
store i64 %"5_01", i64* %"9_0", align 4
store i64 %"7_02", i64* %"9_1", align 4
br label %loop_body

loop_body: ; preds = %12, %entry_block
%"9_03" = load i64, i64* %"9_0", align 4
%"9_14" = load i64, i64* %"9_1", align 4
store i64 2, i64* %"15_0", align 4
store i64 0, i64* %"12_0", align 4
store i64 %"9_03", i64* %"9_0", align 4
store i64 %"9_14", i64* %"9_1", align 4
%"9_15" = load i64, i64* %"9_1", align 4
%"15_06" = load i64, i64* %"15_0", align 4
%0 = mul i64 %"9_15", %"15_06"
store i64 %0, i64* %"16_0", align 4
%"9_07" = load i64, i64* %"9_0", align 4
%"12_08" = load i64, i64* %"12_0", align 4
%1 = icmp eq i64 %"9_07", %"12_08"
%2 = select i1 %1, { i32, {}, {} } { i32 1, {} poison, {} undef }, { i32, {}, {} } { i32 0, {} undef, {} poison }
store { i32, {}, {} } %2, { i32, {}, {} }* %"13_0", align 4
%"13_09" = load { i32, {}, {} }, { i32, {}, {} }* %"13_0", align 4
%"9_010" = load i64, i64* %"9_0", align 4
%3 = extractvalue { i32, {}, {} } %"13_09", 0
switch i32 %3, label %4 [
i32 1, label %6
]

4: ; preds = %loop_body
%5 = extractvalue { i32, {}, {} } %"13_09", 1
store i64 %"9_010", i64* %"013", align 4
br label %cond_17_case_0

6: ; preds = %loop_body
%7 = extractvalue { i32, {}, {} } %"13_09", 2
store i64 %"9_010", i64* %"019", align 4
br label %cond_17_case_1

loop_out: ; preds = %15
%"8_026" = load i64, i64* %"8_0", align 4
store i64 %"8_026", i64* %"0", align 4
%"027" = load i64, i64* %"0", align 4
ret i64 %"027"

cond_17_case_0: ; preds = %4
%"014" = load i64, i64* %"013", align 4
store i64 1, i64* %"23_0", align 4
store i64 %"014", i64* %"20_0", align 4
%"20_015" = load i64, i64* %"20_0", align 4
%"23_016" = load i64, i64* %"23_0", align 4
%8 = sub i64 %"20_015", %"23_016"
store i64 %8, i64* %"24_0", align 4
%"24_017" = load i64, i64* %"24_0", align 4
%9 = insertvalue { i64 } undef, i64 %"24_017", 0
%10 = insertvalue { i32, { i64 }, {} } { i32 0, { i64 } poison, {} poison }, { i64 } %9, 1
store { i32, { i64 }, {} } %10, { i32, { i64 }, {} }* %"25_0", align 4
%"25_018" = load { i32, { i64 }, {} }, { i32, { i64 }, {} }* %"25_0", align 4
store { i32, { i64 }, {} } %"25_018", { i32, { i64 }, {} }* %"011", align 4
br label %cond_exit_17

cond_17_case_1: ; preds = %6
%"020" = load i64, i64* %"019", align 4
store { i32, { i64 }, {} } { i32 1, { i64 } poison, {} undef }, { i32, { i64 }, {} }* %"29_0", align 4
%"29_021" = load { i32, { i64 }, {} }, { i32, { i64 }, {} }* %"29_0", align 4
store { i32, { i64 }, {} } %"29_021", { i32, { i64 }, {} }* %"011", align 4
store i64 %"020", i64* %"27_0", align 4
br label %cond_exit_17

cond_exit_17: ; preds = %cond_17_case_1, %cond_17_case_0
%"012" = load { i32, { i64 }, {} }, { i32, { i64 }, {} }* %"011", align 4
store { i32, { i64 }, {} } %"012", { i32, { i64 }, {} }* %"17_0", align 4
%"17_022" = load { i32, { i64 }, {} }, { i32, { i64 }, {} }* %"17_0", align 4
%"16_023" = load i64, i64* %"16_0", align 4
store { i32, { i64 }, {} } %"17_022", { i32, { i64 }, {} }* %"17_0", align 4
store i64 %"16_023", i64* %"16_0", align 4
%"17_024" = load { i32, { i64 }, {} }, { i32, { i64 }, {} }* %"17_0", align 4
%"16_025" = load i64, i64* %"16_0", align 4
%11 = extractvalue { i32, { i64 }, {} } %"17_024", 0
switch i32 %11, label %12 [
i32 1, label %15
]

12: ; preds = %cond_exit_17
%13 = extractvalue { i32, { i64 }, {} } %"17_024", 1
%14 = extractvalue { i64 } %13, 0
store i64 %14, i64* %"9_0", align 4
store i64 %"16_025", i64* %"9_1", align 4
br label %loop_body

15: ; preds = %cond_exit_17
%16 = extractvalue { i32, { i64 }, {} } %"17_024", 2
store i64 %"16_025", i64* %"8_0", align 4
br label %loop_out
}
Loading
Loading