use std::borrow::Borrow;
use std::cell::RefCell;
use std::fmt::{Debug, Display};
use multithread::Executor;
use crate::internal::*;
use crate::model::{Fact, Graph, OutletId};
use crate::ops::FrozenOpState;
use crate::ops::konst::Const;
use crate::runtime::RunOptions;
use self::order::{build_flush_list, eval_order_for_nodes, eval_order_opt_ram_for_nodes};
pub struct TurnState {
pub resolved_symbols: SymbolValues,
pub scenario: Option<usize>,
pub cached_mmm_scratch_space: RefCell<Option<Box<dyn tract_linalg::mmm::ScratchSpace>>>,
pub scratch_extensions: anymap3::Map,
pub values: Vec<Option<TVec<TValue>>>,
}
impl Default for TurnState {
fn default() -> Self {
TurnState {
resolved_symbols: SymbolValues::default(),
scenario: None,
cached_mmm_scratch_space: None.into(),
scratch_extensions: anymap3::Map::new(),
values: vec![],
}
}
}
impl Clone for TurnState {
fn clone(&self) -> Self {
TurnState {
resolved_symbols: self.resolved_symbols.clone(),
scenario: self.scenario,
cached_mmm_scratch_space: None.into(),
scratch_extensions: anymap3::Map::new(),
values: vec![],
}
}
}
pub trait SessionStateHandler: Send + Sync + Debug {
fn before_plan_eval(&self, session_state: &mut TurnState) -> TractResult<()>;
fn after_plan_eval(&self, session_state: &mut TurnState) -> TractResult<()>;
}
impl Debug for TurnState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "SessionState({:?})", self.resolved_symbols)
}
}
#[derive(Debug, Clone)]
pub struct SimplePlan<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub(crate) model: Arc<Graph<F, O>>,
outputs: Vec<OutletId>,
order: Vec<usize>,
flush_lists: Vec<TVec<usize>>,
has_unresolved_symbols: bool,
executor: Option<Executor>,
session_handler: Option<Arc<dyn SessionStateHandler + 'static>>,
}
impl<F, O> SimplePlan<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub fn new(model: impl Into<Arc<Graph<F, O>>>) -> TractResult<Arc<SimplePlan<F, O>>> {
let model = model.into();
Self::build(model, &RunOptions::default()).map(Arc::new)
}
pub fn new_with_options(
model: impl Into<Arc<Graph<F, O>>>,
options: &RunOptions,
) -> TractResult<Arc<SimplePlan<F, O>>> {
let model = model.into();
Self::build(model, options).map(Arc::new)
}
#[deprecated]
pub fn new_for_output(
model: Graph<F, O>,
output: OutletId,
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, &[output], &[], &RunOptions::default())
.map(Arc::new)
}
#[deprecated]
pub fn new_for_outputs(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, outputs, &[], &RunOptions::default()).map(Arc::new)
}
pub fn with_session_handler<H: SessionStateHandler + 'static>(
mut self,
session_handler: H,
) -> Self {
self.session_handler = Some(Arc::new(session_handler));
self
}
#[deprecated]
pub fn new_for_outputs_and_deps(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
deps: &[(usize, usize)],
) -> TractResult<Arc<SimplePlan<F, O>>> {
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, outputs, deps, &RunOptions::default())
.map(Arc::new)
}
pub fn build(
model: impl Into<Arc<Graph<F, O>>>,
options: &RunOptions,
) -> TractResult<SimplePlan<F, O>> {
let model = model.into();
let outputs = model.outputs.clone();
#[allow(deprecated)]
Self::build_with_outputs_and_deps(model, &outputs, &[], options)
}
#[deprecated]
pub fn build_with_outputs_and_deps(
model: impl Into<Arc<Graph<F, O>>>,
outputs: &[OutletId],
deps: &[(usize, usize)],
options: &RunOptions,
) -> TractResult<SimplePlan<F, O>> {
let model = model.into();
let inputs = model.input_outlets()?.iter().map(|n| n.node).collect::<Vec<usize>>();
let outputs_nodes = outputs.iter().map(|n| n.node).collect::<Vec<usize>>();
let mut order = if options.skip_order_opt_ram {
eval_order_for_nodes(model.nodes(), &inputs, &outputs_nodes, deps)?
} else {
eval_order_opt_ram_for_nodes(model.nodes(), &inputs, &outputs_nodes, deps)?
};
order.retain(|node| !model.node(*node).op_is::<Const>());
let flush_lists = build_flush_list(&*model, &order, outputs, |n| !n.op_is::<Const>());
#[allow(clippy::mutable_key_type)]
let mut symbols: std::collections::HashSet<Symbol> = Default::default();
for node in &model.nodes {
for output in &node.outputs {
if let Ok(fact) = output.fact.to_typed_fact() {
symbols.extend(fact.shape.iter().flat_map(|d| d.symbols()))
}
}
}
Ok(SimplePlan {
model,
order,
flush_lists,
outputs: outputs.to_vec(),
has_unresolved_symbols: !symbols.is_empty(),
executor: options.executor.clone(),
session_handler: None,
})
}
pub fn order_without_consts(&self) -> &[usize] {
&self.order
}
pub fn run(self: &Arc<Self>, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
let mut state = self.spawn()?;
state.run(inputs)
}
pub fn model(&self) -> &Graph<F, O> {
self.model.borrow()
}
pub fn spawn(self: &Arc<Self>) -> TractResult<SimpleState<F, O>> {
SimpleState::new(self)
}
}
#[derive(Clone, Debug)]
pub struct SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub(crate) plan: Arc<SimplePlan<F, O>>,
pub op_states: Vec<Option<Box<dyn OpState>>>,
pub turn_state: TurnState,
}
impl<F, O> SimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub fn new(plan: &Arc<SimplePlan<F, O>>) -> TractResult<SimpleState<F, O>> {
let plan = Arc::clone(plan);
let turn = TurnState::default();
let model = plan.model();
let states: Vec<Option<Box<dyn OpState>>> = vec![None; model.nodes.len()];
let mut state = SimpleState { plan, op_states: states, turn_state: turn };
state.reset_op_states()?;
Ok(state)
}
pub fn new_from_inputs(
plan: &Arc<SimplePlan<F, O>>,
inputs: TVec<TValue>,
) -> TractResult<SimpleState<F, O>> {
let mut state = SimpleState::new(plan)?;
state.set_inputs(inputs)?;
state.resolve_symbols_with_states()?;
Ok(state)
}
fn ready_turn(&mut self) {
if self.turn_state.values.len() == 0 {
self.turn_state.values = vec![None; self.plan.model.nodes().len()];
for node in &self.plan.model.nodes {
if let Some(k) = node.op_as::<Const>() {
self.turn_state.values[node.id] = Some(tvec!(k.val().clone().into_tvalue()));
}
}
}
}
pub fn reset_turn(&mut self) -> TractResult<()> {
for node in &self.plan.order {
self.turn_state.values[*node] = None;
}
self.turn_state.resolved_symbols = SymbolValues::default();
Ok(())
}
fn reset_op_states(&mut self) -> TractResult<()> {
let &mut SimpleState { ref plan, ref mut turn_state, op_states: ref mut states, .. } = self;
for (ix, n) in plan.model.nodes.iter().enumerate() {
states[ix] = if n.op().is_stateless() { None } else { n.op().state(turn_state, ix)? };
}
Ok(())
}
fn resolve_symbols_with_states(&mut self) -> TractResult<()> {
for state in self
.op_states
.iter_mut()
.filter_map(Option::as_mut)
.filter(|s| s.init_tensor_fact().is_some())
{
state.resolve_symbols(&mut self.turn_state)?;
}
Ok(())
}
pub fn run(&mut self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
self.run_plan_with_eval(inputs, self::eval)
}
pub fn exec(&mut self) -> TractResult<()> {
self.exec_plan_with_eval(self::eval)
}
pub fn run_plan_with_eval<Eval, E>(
&mut self,
inputs: TVec<TValue>,
eval: Eval,
) -> TractResult<TVec<TValue>>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a mut TurnState,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
self.set_inputs(inputs)?;
self.resolve_symbols_with_states()?;
self.exec_plan_with_eval(eval)?;
let outputs = self.outputs()?;
self.reset_turn()?;
Ok(outputs)
}
pub fn exec_plan_with_eval<Eval, E>(&mut self, eval: Eval) -> TractResult<()>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a mut TurnState,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
if let Some(executor) = self.plan().executor.as_ref() {
tract_linalg::multithread::multithread_tract_scope(executor.clone(), || {
self.do_exec_plan_with_eval(eval)
})
} else {
self.do_exec_plan_with_eval(eval)
}
}
fn do_exec_plan_with_eval<Eval, E>(&mut self, mut eval: Eval) -> TractResult<()>
where
Eval: for<'a, 'b, 'c> FnMut(
&'a mut TurnState,
Option<&'b mut (dyn OpState + 'static)>,
&'c Node<F, O>,
TVec<TValue>,
) -> Result<TVec<TValue>, E>,
E: Into<anyhow::Error> + Send + Sync + 'static,
{
{
self.ready_turn();
self.plan
.session_handler
.as_ref()
.map(|it| it.before_plan_eval(&mut self.turn_state))
.transpose()?;
for (step, n) in self.plan.order.iter().enumerate() {
let node = self.plan.model.node(*n);
trace!("Running step {step}, node {node}");
let mut inputs: TVec<TValue> = tvec![];
for i in &node.inputs {
trace!(" use input {i:?}");
let prec_node = self.plan.model.node(i.node);
let prec = self.turn_state.values[i.node].as_ref().ok_or_else(|| {
format_err!("Computing {}, precursor {} not done:", node, prec_node)
})?;
inputs.push(prec[i.slot].clone())
}
for flush in &self.plan.flush_lists[step] {
trace!(" Ran {} can now flush {}", node, self.plan.model.node(*flush));
self.turn_state.values[*flush] = None;
}
if cfg!(debug_assertions) {
let facts = self.plan.model.node_input_facts(node.id)?;
if facts.len() != inputs.len() {
bail!(
"Evaluating {}: expected {} inputs, got {}",
node,
facts.len(),
inputs.len()
);
}
for (ix, (v, f)) in inputs.iter().zip(facts.iter()).enumerate() {
if !f.matches(v, Some(&self.turn_state.resolved_symbols))? {
bail!(
"Evaluating {}: input {:?}, expected {:?}, got {:?}",
node,
ix,
f,
v
);
}
}
}
let vs = eval(
&mut self.turn_state,
self.op_states[node.id].as_deref_mut(),
node,
inputs,
)
.map_err(|e| e.into())?;
if self.plan.has_unresolved_symbols {
for (o, v) in node.outputs.iter().zip(vs.iter()) {
if let Ok(f) = o.fact.to_typed_fact() {
for (dim_abstract, dim_concrete) in f.shape.iter().zip(v.shape()) {
Self::resolve(
&mut self.turn_state,
dim_abstract,
*dim_concrete as i64,
)?;
}
}
}
}
if cfg!(debug_assertions) {
let facts = self.plan.model.node_output_facts(node.id)?;
if facts.len() != vs.len() {
bail!(
"Evaluating {}: expected {} outputs, got {}",
node,
facts.len(),
vs.len()
);
}
for (ix, (v, f)) in vs.iter().zip(facts.iter()).enumerate() {
if node.outputs[ix].successors.len() == 0 {
continue;
}
if !f.matches(v, Some(&self.turn_state.resolved_symbols))? {
bail!(
"Evaluating {}: output {:?}, expected {:?}, got {:?}",
node,
ix,
f,
v
);
}
}
}
self.turn_state.values[node.id] = Some(vs);
}
self.plan
.session_handler
.as_ref()
.map(|it| it.after_plan_eval(&mut self.turn_state))
.transpose()?;
}
Ok(())
}
pub fn set_inputs(&mut self, inputs: TVec<TValue>) -> TractResult<()> {
ensure!(
inputs.len() == self.model().inputs.len(),
"Wrong number of inputs for model. Expected {} got {}",
self.model().inputs.len(),
inputs.len()
);
for (ix, t) in inputs.into_iter().enumerate() {
self.set_input(ix, t)?
}
Ok(())
}
fn resolve(state: &mut TurnState, expression: &TDim, provided: i64) -> TractResult<()> {
let expected = expression.eval(&state.resolved_symbols);
if let Ok(x) = expected.to_i64()
&& x != provided
{
bail!("Clashing resolution for expression. {expression}={x} != {provided}. ({state:?})")
}
if expected.symbols().len() == 1 {
let sym = expected.symbols().into_iter().next().unwrap();
if let Some(v) = solve_for(&sym, &expected, &provided.to_dim()) {
debug!("Determined symbol {sym}={v}");
state.resolved_symbols.set(&sym, v.to_i64().unwrap());
}
if state.scenario.is_none() {
let scope = sym
.scope()
.with_context(|| format!("Symbol {sym:?} points to an invalid (dead ?) SymbolScope. Make sure to create symbols using the model-managed SymbolScope."))?;
state.scenario = scope.guess_scenario(&state.resolved_symbols)?;
}
}
Ok(())
}
pub fn set_input(&mut self, input: usize, t: TValue) -> TractResult<()> {
let outlet: OutletId = *self
.model()
.input_outlets()?
.get(input)
.with_context(|| format!("Invalid input id for model ({input})."))?;
if let Ok(fact) = self.plan.model.outlet_fact(outlet)?.to_typed_fact() {
for (expected, provided) in fact.shape.iter().zip(t.shape()) {
Self::resolve(&mut self.turn_state, expected, *provided as i64)?;
}
}
let fact = self.plan.model.outlet_fact(outlet)?;
ensure!(
fact.matches(&t, Some(&self.turn_state.resolved_symbols))
.with_context(|| format!("Setting input {input}"))?,
"Input at index {input} has incorrect dtype or shape (got {t:?}, expected to match fact {fact:?})",
);
self.ready_turn();
self.turn_state.values[outlet.node] = Some(tvec!(t));
Ok(())
}
pub fn output(&self, id: usize) -> TractResult<&TValue> {
let outlet = self.model().output_outlets()?.get(id).with_context(|| {
format!(
"Required output {}, only have {}",
id,
self.model().output_outlets().unwrap().len()
)
})?;
let value: &TValue = self
.turn_state
.values
.get(outlet.node)
.context("node id for output beyond node values array")?
.as_ref()
.context("node is not an output")?
.get(outlet.slot)
.context("slot id too high")?;
Ok(value)
}
pub fn outputs(&mut self) -> TractResult<TVec<TValue>> {
let &mut SimpleState { ref plan, ref mut turn_state, .. } = self;
let mut v = tvec![];
for o in plan.outputs.iter() {
let vs = turn_state.values[o.node].as_mut().ok_or_else(|| {
format_err!("Outputs of {:?} are not computed", &plan.model.nodes()[o.node])
})?;
v.push(vs[o.slot].clone())
}
Ok(v)
}
pub fn set_values(&mut self, id: usize, values: TVec<TValue>) -> TractResult<()> {
self.turn_state.values[id] = Some(values);
Ok(())
}
pub fn set_value(&mut self, id: usize, value: TValue) -> TractResult<()> {
self.set_values(id, tvec!(value))
}
pub fn prepare_inputs(&self, node: usize) -> TractResult<TVec<TValue>> {
let SimpleState { plan, turn_state, .. } = self;
let nodes = plan.model.nodes();
let node = &nodes[node];
let mut inputs: TVec<TValue> = tvec![];
for i in &node.inputs {
let prec_node = &nodes[i.node];
let prec = turn_state.values[i.node].as_ref().ok_or_else(|| {
format_err!("Computing {}, precursor {} not done.", node, prec_node)
})?;
inputs.push(prec[i.slot].clone())
}
Ok(inputs)
}
pub fn compute_one(&mut self, node: usize) -> TractResult<()> {
let inputs = self.prepare_inputs(node)?;
self.compute_one_with_inputs(node, inputs)
}
pub fn compute_one_with_inputs(
&mut self,
node: usize,
inputs: TVec<TValue>,
) -> TractResult<()> {
let &mut SimpleState { ref plan, ref mut turn_state, op_states: ref mut states, .. } = self;
let nodes = plan.model.nodes();
let node = &nodes[node];
let vs = eval(turn_state, states[node.id].as_deref_mut(), node, inputs)?;
turn_state.values[node.id] = Some(vs);
Ok(())
}
pub fn compute_recursively(&mut self, node: usize) -> TractResult<&[TValue]> {
let values = {
#[allow(clippy::needless_collect)] let precs: Vec<usize> =
self.model().nodes()[node].inputs.iter().map(|i| i.node).collect();
for i in precs.into_iter() {
if self.turn_state.values[i].is_none() {
let _ = self.compute_recursively(i)?;
}
}
let mut inputs: TVec<TValue> = tvec![];
{
let node = &self.model().nodes()[node];
for i in &node.inputs {
inputs.push(self.turn_state.values[i.node].as_ref().unwrap()[i.slot].clone())
}
}
let &mut Self {
op_states: ref mut states,
turn_state: ref mut session_state,
ref plan,
..
} = self;
eval(session_state, states[node].as_deref_mut(), &plan.model().nodes[node], inputs)?
};
self.turn_state.values[node] = Some(values);
Ok(self.turn_state.values[node].as_ref().unwrap())
}
pub fn take_by_name(&mut self, name: &str) -> TractResult<TVec<Tensor>> {
let id = self.model().node_by_name(name)?.id;
Self::take(self, id)
}
pub fn take(&mut self, id: usize) -> TractResult<TVec<Tensor>> {
Ok(self.turn_state.values[id]
.take()
.ok_or_else(|| format_err!("Node is not computed"))?
.into_iter()
.map(|v| v.into_tensor())
.collect())
}
pub fn plan(&self) -> &Arc<SimplePlan<F, O>> {
&self.plan
}
pub fn model(&self) -> &Graph<F, O> {
&self.plan.model
}
pub fn freeze(&self) -> FrozenSimpleState<F, O> {
FrozenSimpleState {
plan: self.plan.clone(),
resolved_symbols: self.turn_state.resolved_symbols.clone(),
scenario: self.turn_state.scenario,
states: self.op_states.iter().map(|s| s.as_ref().map(|s| s.freeze())).collect(),
values: self
.turn_state
.values
.iter()
.enumerate()
.map(|(ix, t)| {
if self.model().nodes[ix].op_is::<Const>() {
t.as_ref().map(|t| t.iter().map(|t| t.clone().into_tensor()).collect())
} else {
None
}
})
.collect(),
}
}
pub fn freeze_into(self) -> FrozenSimpleState<F, O> {
let plan = self.plan;
let model = &plan.model;
FrozenSimpleState {
resolved_symbols: self.turn_state.resolved_symbols,
scenario: self.turn_state.scenario,
states: self.op_states.into_iter().map(|s| s.map(|s| s.freeze_into())).collect(),
values: self
.turn_state
.values
.into_iter()
.enumerate()
.map(|(ix, t)| {
if model.nodes[ix].op_is::<Const>() {
t.map(|t| t.into_iter().map(|t| t.into_tensor()).collect())
} else {
None
}
})
.collect(),
plan,
}
}
}
pub fn eval<F, O>(
session_state: &mut TurnState,
mut state: Option<&mut (dyn OpState + 'static)>,
node: &Node<F, O>,
input: TVec<TValue>,
) -> TractResult<TVec<TValue>>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
#[allow(clippy::let_and_return)]
let r = match state {
Some(ref mut state) => state.eval(session_state, node.op(), input),
None => node.op().eval_with_session(node.id, session_state, input),
}
.with_context(|| format!("Evaluating {node}"));
r
}
#[derive(Clone, Debug)]
pub struct FrozenSimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
plan: Arc<SimplePlan<F, O>>,
pub resolved_symbols: SymbolValues,
pub scenario: Option<usize>,
pub states: Vec<Option<Box<dyn FrozenOpState>>>,
pub values: Vec<Option<TVec<Tensor>>>,
}
impl<F, O> FrozenSimpleState<F, O>
where
F: Fact + Clone + 'static,
O: Debug + Display + AsRef<dyn Op> + AsMut<dyn Op> + Clone + 'static,
{
pub fn plan(&self) -> &Arc<SimplePlan<F, O>> {
&self.plan
}
pub fn unfreeze(&self) -> SimpleState<F, O> {
SimpleState {
plan: self.plan.clone(),
turn_state: TurnState {
resolved_symbols: self.resolved_symbols.clone(),
scenario: self.scenario,
cached_mmm_scratch_space: None.into(),
scratch_extensions: anymap3::Map::new(),
values: self
.values
.iter()
.map(|t| {
t.as_ref().map(|t| t.iter().map(|t| t.clone().into_tvalue()).collect())
})
.collect(),
},
op_states: self.states.iter().map(|s| s.as_ref().map(|s| s.unfreeze())).collect(),
}
}
}
#[cfg(test)]
mod test {
use super::*;
fn is_send<T: Send>() {}
fn is_sync<T: Sync>() {}
#[test]
fn type_model_is_sync() {
is_sync::<TypedModel>();
}
#[test]
fn type_model_is_send() {
is_send::<TypedModel>();
}
#[test]
fn type_plan_is_send() {
is_send::<TypedSimplePlan>();
}
#[test]
fn type_plan_is_sync() {
is_sync::<TypedSimplePlan>();
}
#[test]
fn frozen_type_state_is_send() {
is_send::<TypedFrozenSimpleState>();
}
}