use crate::ir::graph::*;
pub struct ScGraphBuilder {
graph: ScGraph,
}
impl ScGraphBuilder {
pub fn new(name: impl Into<String>) -> Self {
Self {
graph: ScGraph::new(name),
}
}
pub fn input(&mut self, name: impl Into<String>, ty: ScType) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Input {
id,
name: name.into(),
ty,
})
}
pub fn output(&mut self, name: impl Into<String>, source: ValueId) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Output {
id,
name: name.into(),
source,
})
}
pub fn constant(&mut self, value: ScConst, ty: ScType) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Constant { id, value, ty })
}
pub fn encode(&mut self, prob: ValueId, length: usize, seed: u16) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Encode {
id,
prob,
length,
seed,
})
}
pub fn bitwise_and(&mut self, lhs: ValueId, rhs: ValueId) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::BitwiseAnd { id, lhs, rhs })
}
pub fn popcount(&mut self, input: ValueId) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Popcount { id, input })
}
pub fn lif_step(
&mut self,
current: ValueId,
leak: ValueId,
gain: ValueId,
noise: ValueId,
params: LifParams,
) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::LifStep {
id,
current,
leak,
gain,
noise,
params,
})
}
pub fn dense_forward(
&mut self,
inputs: ValueId,
weights: ValueId,
leak: ValueId,
gain: ValueId,
params: DenseParams,
) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::DenseForward {
id,
inputs,
weights,
leak,
gain,
params,
})
}
pub fn bitwise_xor(&mut self, lhs: ValueId, rhs: ValueId) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::BitwiseXor { id, lhs, rhs })
}
pub fn reduce(&mut self, input: ValueId, mode: ReduceMode) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Reduce { id, input, mode })
}
pub fn graph_forward(
&mut self,
features: ValueId,
adjacency: ValueId,
n_nodes: usize,
n_features: usize,
) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::GraphForward {
id,
features,
adjacency,
n_nodes,
n_features,
})
}
pub fn softmax_attention(
&mut self,
q: ValueId,
k: ValueId,
v: ValueId,
dim_k: usize,
) -> ValueId {
let id = self.graph.fresh_id();
self.graph
.push(ScOp::SoftmaxAttention { id, q, k, v, dim_k })
}
pub fn kuramoto_step(
&mut self,
phases: ValueId,
omega: ValueId,
coupling: ValueId,
dt: f64,
) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::KuramotoStep {
id,
phases,
omega,
coupling,
dt,
})
}
pub fn scale(&mut self, input: ValueId, factor: f64) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Scale { id, input, factor })
}
pub fn offset(&mut self, input: ValueId, offset_val: f64) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::Offset {
id,
input,
offset: offset_val,
})
}
pub fn div_const(&mut self, input: ValueId, divisor: u64) -> ValueId {
let id = self.graph.fresh_id();
self.graph.push(ScOp::DivConst { id, input, divisor })
}
pub fn build(self) -> ScGraph {
self.graph
}
}