use super::{
ir::{Arg, BinaryElemwiseArgs, ElemwiseOp, ElemwisePrecision, UnaryElemwiseArgs},
trace::FuseOnWriteTrace,
trace_builder::FuseOnWriteTraceBuilder,
};
use burn_fusion::{OptimizationBuilder, OptimizationProperties, OptimizationStatus};
use burn_tensor::{
repr::{
BaseOperationDescription, BinaryOperationDescription, FloatOperationDescription,
NumericOperationDescription, OperationDescription, ScalarOperationDescription,
TensorDescription, UnaryOperationDescription,
},
Element,
};
use cubecl::ir::Elem;
pub(crate) struct FuseOnWriteBuilder {
builder: TryFuseBuilder,
current_output_shape: Vec<usize>,
status: OptimizationStatus,
num_ops: usize,
max_bindings: u32,
}
struct TryFuseBuilder {
builder: FuseOnWriteTraceBuilder,
max_bindings: u32,
added_ops: bool,
}
impl TryFuseBuilder {
fn new(max_bindings: u32, bool_precision: ElemwisePrecision) -> Self {
Self {
builder: FuseOnWriteTraceBuilder::new(bool_precision),
max_bindings,
added_ops: false,
}
}
fn register(&mut self, add_ops: impl FnOnce(&mut FuseOnWriteTraceBuilder)) -> bool {
if !self.added_ops {
self.added_ops = true;
add_ops(&mut self.builder);
return true;
}
let mut cloned = self.builder.clone();
add_ops(&mut cloned);
if cloned.estimate_bindings() > self.max_bindings {
return false;
}
self.builder = cloned;
true
}
fn build(&self) -> FuseOnWriteTrace {
self.builder.build()
}
}
impl OptimizationBuilder<FuseOnWriteTrace> for FuseOnWriteBuilder {
fn register(&mut self, op: &OperationDescription) {
if let OptimizationStatus::Closed = self.status {
return;
}
match op {
OperationDescription::BaseFloat(ops) => {
if !self.register_base(ops) {
self.status = OptimizationStatus::Closed;
return;
}
}
OperationDescription::BaseInt(ops) => {
if !self.register_base(ops) {
self.status = OptimizationStatus::Closed;
return;
}
}
OperationDescription::Float(_dtype, ops) => {
if !self.register_float(ops) {
self.status = OptimizationStatus::Closed;
return;
}
}
OperationDescription::NumericFloat(_dtype, ops) => {
if !self.register_numeric::<f32>(ops) {
self.status = OptimizationStatus::Closed;
return;
}
}
OperationDescription::NumericInt(_dtype, ops) => {
if !self.register_numeric::<i32>(ops) {
self.status = OptimizationStatus::Closed;
return;
}
}
_ => {
self.status = OptimizationStatus::Closed;
return;
}
};
self.status = OptimizationStatus::Open;
self.num_ops += 1;
}
fn build(&self) -> FuseOnWriteTrace {
self.builder.build()
}
fn len(&self) -> usize {
self.num_ops
}
fn reset(&mut self) {
self.num_ops = 0;
self.status = OptimizationStatus::Open;
self.builder = TryFuseBuilder::new(self.max_bindings, self.builder.builder.bool_precision);
self.current_output_shape.clear();
}
fn status(&self) -> OptimizationStatus {
self.status
}
fn properties(&self) -> OptimizationProperties {
let ready = self.num_ops > 0;
OptimizationProperties {
ready,
score: self.num_ops as u64,
}
}
}
impl FuseOnWriteBuilder {
pub fn new(max_bindings: u32, bool_precision: ElemwisePrecision) -> Self {
Self {
builder: TryFuseBuilder::new(max_bindings, bool_precision),
num_ops: 0,
max_bindings,
current_output_shape: Vec::new(),
status: OptimizationStatus::Open,
}
}
pub fn close(&mut self) {
self.status = OptimizationStatus::Closed;
}
pub fn input_unhandled(&mut self, tensor: &TensorDescription) -> Arg {
self.builder.builder.input_unhandled(tensor)
}
pub fn output_unhandled(&mut self, tensor: &TensorDescription) -> Arg {
if self.current_output_shape.is_empty() {
self.current_output_shape = tensor.shape.clone();
}
self.builder.builder.output_unhandled(tensor)
}
fn register_base(&mut self, ops: &BaseOperationDescription) -> bool {
match ops {
BaseOperationDescription::Equal(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Equal(BinaryElemwiseArgs { lhs, rhs, out })
}),
BaseOperationDescription::Cast(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Assign(UnaryElemwiseArgs { input, out })
}),
_ => false,
}
}
fn register_float(&mut self, ops: &FloatOperationDescription) -> bool {
match ops {
FloatOperationDescription::Exp(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Exp(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Log(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Log(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Log1p(desc) => self
.register_unary_ops(desc, |input, out| {
ElemwiseOp::Log1p(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Cos(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Cos(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Sin(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Sin(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::PowfScalar(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Powf(BinaryElemwiseArgs { lhs, rhs, out })
}),
FloatOperationDescription::Tanh(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Tanh(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Erf(desc) => self.register_unary_ops(desc, |input, out| {
ElemwiseOp::Erf(UnaryElemwiseArgs { input, out })
}),
FloatOperationDescription::Recip(desc) => self
.register_unary_ops(desc, |input, out| {
ElemwiseOp::Recip(UnaryElemwiseArgs { input, out })
}),
_ => false,
}
}
fn register_numeric<E: Element>(&mut self, op: &NumericOperationDescription<E>) -> bool {
match op {
NumericOperationDescription::Add(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Add(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::AddScalar(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Add(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::Sub(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Sub(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::SubScalar(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Sub(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::Mul(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Mul(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::MulScalar(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Mul(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::Div(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Div(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::DivScalar(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Div(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::Abs(desc) => self
.register_unary_ops(desc, |input, out| {
ElemwiseOp::Abs(UnaryElemwiseArgs { input, out })
}),
NumericOperationDescription::Lower(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Lower(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::LowerElem(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Lower(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::Greater(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Greater(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::GreaterElem(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Greater(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::LowerEqual(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::LowerEqual(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::LowerEqualElem(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::LowerEqual(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::GreaterEqual(desc) => self
.register_binary_ops(desc, |lhs, rhs, out| {
ElemwiseOp::GreaterEqual(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::GreaterEqualElem(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::GreaterEqual(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::EqualElem(desc) => self
.register_scalar_ops(desc, |lhs, rhs, out| {
ElemwiseOp::Equal(BinaryElemwiseArgs { lhs, rhs, out })
}),
NumericOperationDescription::MaskWhere(desc) => {
if !self.output_is_compatible(&desc.out) {
return false;
}
self.builder.register(|build| {
let cond = build.input(&desc.mask);
let lhs = build.input(&desc.value);
let rhs = build.input(&desc.tensor);
let out = build.output(&desc.out);
build.register_operation(ElemwiseOp::ConditionalAssign {
cond,
lhs,
rhs,
out,
})
})
}
NumericOperationDescription::MaskFill(desc) => {
if !self.output_is_compatible(&desc.out) {
return false;
}
self.builder.register(|build| {
let cond = build.input(&desc.mask);
let lhs = build.scalar(&desc.value, desc.out.dtype);
let rhs = build.input(&desc.tensor);
let out = build.output(&desc.out);
build.register_operation(ElemwiseOp::ConditionalAssign {
cond,
lhs,
rhs,
out,
})
})
}
NumericOperationDescription::Ones(desc) => {
if !self.output_is_compatible(desc) {
return false;
}
let elem: Elem = desc.dtype.into();
let precision = elem.into();
let input = Arg::Literal(1, precision);
self.builder.register(|build| {
let out = build.output(desc);
build.register_operation(ElemwiseOp::Assign(UnaryElemwiseArgs { input, out }))
})
}
NumericOperationDescription::Zeros(desc) => {
if !self.output_is_compatible(desc) {
return false;
}
let elem: Elem = desc.dtype.into();
let precision = elem.into();
let input = Arg::Literal(0, precision);
self.builder.register(|build| {
let out = build.output(desc);
build.register_operation(ElemwiseOp::Assign(UnaryElemwiseArgs { input, out }))
})
}
NumericOperationDescription::Full((desc, elem)) => {
if !self.output_is_compatible(desc) {
return false;
}
self.builder.register(|build| {
let input = build.scalar(elem, desc.dtype);
let out = build.output(desc);
build.register_operation(ElemwiseOp::Assign(UnaryElemwiseArgs { input, out }))
})
}
_ => false,
}
}
fn register_binary_ops<Func>(&mut self, desc: &BinaryOperationDescription, func: Func) -> bool
where
Func: Fn(Arg, Arg, Arg) -> ElemwiseOp,
{
if !self.output_is_compatible(&desc.out) {
return false;
}
self.builder.register(|build| {
let lhs = build.input(&desc.lhs);
let rhs = build.input(&desc.rhs);
let out = build.output(&desc.out);
build.register_operation(func(lhs, rhs, out))
})
}
fn register_unary_ops<Func>(&mut self, desc: &UnaryOperationDescription, func: Func) -> bool
where
Func: Fn(Arg, Arg) -> ElemwiseOp,
{
if !self.output_is_compatible(&desc.out) {
return false;
}
self.builder.register(|build| {
let input = build.input(&desc.input);
let out = build.output(&desc.out);
build.register_operation(func(input, out))
})
}
fn register_scalar_ops<Func, E: Element>(
&mut self,
desc: &ScalarOperationDescription<E>,
func: Func,
) -> bool
where
Func: Fn(Arg, Arg, Arg) -> ElemwiseOp,
{
if !self.output_is_compatible(&desc.out) {
return false;
}
self.builder.register(|build| {
let elem = desc.lhs.dtype;
let lhs = build.input(&desc.lhs);
let rhs = build.scalar(&desc.rhs, elem);
let out = build.output(&desc.out);
build.register_operation(func(lhs, rhs, out))
})
}
fn output_is_compatible(&mut self, out: &TensorDescription) -> bool {
if self.current_output_shape.is_empty() {
self.current_output_shape.clone_from(&out.shape);
} else if self.current_output_shape != out.shape {
return false;
}
true
}
}