use cubecl_core::ir as core;
use cubecl_opt::{ControlFlow, NodeIndex};
use rspirv::{
dr::Operand,
spirv::{LoopControl, SelectionControl, Word},
};
use crate::{SpirvCompiler, SpirvTarget, item::Item, variable::Variable};
impl<T: SpirvTarget> SpirvCompiler<T> {
pub fn compile_read_bound(
&mut self,
arr: &Variable,
index: Word,
item: Item,
read: impl FnOnce(&mut Self) -> Word,
) -> Word {
let ty = item.id(self);
let len = match arr.has_buffer_len() {
true => self.buffer_length(arr, None, false),
false => self.length(arr, None, false),
};
let bool = self.type_bool();
let cond = self.u_less_than(bool, None, index, len).unwrap();
let current_block = self.current_block.unwrap();
let current_label = self.end_label(current_block);
let in_bounds = self.id();
let next = self.id();
self.selection_merge(next, SelectionControl::DONT_FLATTEN)
.unwrap();
self.branch_conditional(cond, in_bounds, next, vec![1, 0])
.unwrap();
self.begin_block(Some(in_bounds)).unwrap();
let value = read(self);
self.branch(next).unwrap();
let fallback_value = item.constant(self, 0u32.into());
self.state.end_labels.insert(current_block, next);
self.begin_block(Some(next)).unwrap();
self.phi(
ty,
None,
vec![(value, in_bounds), (fallback_value, current_label)],
)
.unwrap()
}
pub fn compile_write_bound(
&mut self,
arr: &Variable,
index: Word,
write: impl FnOnce(&mut Self),
) {
let len = match arr.has_buffer_len() {
true => self.buffer_length(arr, None, false),
false => self.length(arr, None, false),
};
let bool = self.type_bool();
let cond = self.u_less_than(bool, None, index, len).unwrap();
let current_block = self.current_block.unwrap();
let in_bounds = self.id();
let next = self.id();
self.selection_merge(next, SelectionControl::DONT_FLATTEN)
.unwrap();
self.branch_conditional(cond, in_bounds, next, vec![1, 0])
.unwrap();
self.begin_block(Some(in_bounds)).unwrap();
write(self);
self.branch(next).unwrap();
self.begin_block(Some(next)).unwrap();
self.state.end_labels.insert(current_block, next);
}
pub fn compile_copy_bound(
&mut self,
input: &Variable,
out: &Variable,
in_index: Word,
out_index: Word,
len: Option<u32>,
copy: impl FnOnce(&mut Self),
) {
let in_len = match input.has_buffer_len() {
true => self.buffer_length(input, None, false),
false => self.length(input, None, false),
};
let out_len = match out.has_buffer_len() {
true => self.buffer_length(out, None, false),
false => self.length(out, None, false),
};
let bool = self.type_bool();
let int = self.type_int(32, 0);
let in_index = match len {
Some(len) => self.i_add(int, None, in_index, len).unwrap(),
None => in_index,
};
let out_index = match len {
Some(len) => self.i_add(int, None, out_index, len).unwrap(),
None => out_index,
};
let cond_in = self.u_less_than(bool, None, in_index, in_len).unwrap();
let cond_out = self.u_less_than(bool, None, out_index, out_len).unwrap();
let cond = self.logical_and(bool, None, cond_in, cond_out).unwrap();
let current_block = self.current_block.unwrap();
let in_bounds = self.id();
let next = self.id();
self.selection_merge(next, SelectionControl::DONT_FLATTEN)
.unwrap();
self.branch_conditional(cond, in_bounds, next, vec![1, 0])
.unwrap();
self.begin_block(Some(in_bounds)).unwrap();
copy(self);
self.branch(next).unwrap();
self.begin_block(Some(next)).unwrap();
self.state.end_labels.insert(current_block, next);
}
pub fn compile_control_flow(&mut self, control_flow: ControlFlow) {
match control_flow {
ControlFlow::IfElse {
cond,
then,
or_else,
merge,
} => self.compile_if_else(cond, then, or_else, merge),
ControlFlow::Switch {
value,
default,
branches,
merge,
} => self.compile_switch(value, default, branches, merge),
ControlFlow::Loop {
body,
continue_target,
merge,
} => self.compile_loop(body, continue_target, merge),
ControlFlow::LoopBreak {
break_cond,
body,
continue_target,
merge,
} => self.compile_loop_break(break_cond, body, continue_target, merge),
ControlFlow::Return => {
self.ret().unwrap();
self.current_block = None;
}
ControlFlow::Unreachable => {
self.unreachable().unwrap();
self.current_block = None;
}
ControlFlow::None => {
let opt = self.opt.clone();
let children = opt.successors(self.current_block.unwrap());
assert_eq!(
children.len(),
1,
"None control flow should have only 1 outgoing edge"
);
let label = self.label(children[0]);
self.branch(label).unwrap();
self.compile_block(children[0]);
}
}
}
fn compile_if_else(
&mut self,
cond: core::Variable,
then: NodeIndex,
or_else: NodeIndex,
merge: Option<NodeIndex>,
) {
let cond = self.compile_variable(cond);
let then_label = self.label(then);
let else_label = self.label(or_else);
let cond_id = self.read(&cond);
if let Some(merge) = merge {
let merge_label = self.label(merge);
self.selection_merge(merge_label, SelectionControl::NONE)
.unwrap();
}
self.branch_conditional(cond_id, then_label, else_label, None)
.unwrap();
self.compile_block(then);
self.compile_block(or_else);
if let Some(it) = merge {
self.compile_block(it);
}
}
fn compile_switch(
&mut self,
value: core::Variable,
default: NodeIndex,
branches: Vec<(u32, NodeIndex)>,
merge: Option<NodeIndex>,
) {
let value = self.compile_variable(value);
let value_id = self.read(&value);
let default_label = self.label(default);
let targets = branches
.iter()
.map(|(value, block)| {
let label = self.label(*block);
(Operand::LiteralBit32(*value), label)
})
.collect::<Vec<_>>();
if let Some(merge) = merge {
let merge_label = self.label(merge);
self.selection_merge(merge_label, SelectionControl::NONE)
.unwrap();
}
self.switch(value_id, default_label, targets).unwrap();
self.compile_block(default);
for (_, block) in branches {
self.compile_block(block);
}
if let Some(it) = merge {
self.compile_block(it);
}
}
fn compile_loop(&mut self, body: NodeIndex, continue_target: NodeIndex, merge: NodeIndex) {
let body_label = self.label(body);
let continue_label = self.label(continue_target);
let merge_label = self.label(merge);
self.loop_merge(merge_label, continue_label, LoopControl::NONE, vec![])
.unwrap();
self.branch(body_label).unwrap();
self.compile_block(body);
self.compile_block(continue_target);
self.compile_block(merge);
}
fn compile_loop_break(
&mut self,
break_cond: core::Variable,
body: NodeIndex,
continue_target: NodeIndex,
merge: NodeIndex,
) {
let break_cond = self.compile_variable(break_cond);
let cond_id = self.read(&break_cond);
let body_label = self.label(body);
let continue_label = self.label(continue_target);
let merge_label = self.label(merge);
self.loop_merge(merge_label, continue_label, LoopControl::NONE, [])
.unwrap();
self.branch_conditional(cond_id, body_label, merge_label, [])
.unwrap();
self.compile_block(body);
self.compile_block(continue_target);
self.compile_block(merge);
}
}