Skip to main content

sim_lib_numbers_tensor/implementation/
tensor_site.rs

1//! Tensor execution site over the kernel EvalFabric contract.
2
3use std::sync::Arc;
4
5use sim_kernel::{
6    CapabilityName, ClassRef, Cx, DefaultFactory, Env, Error, EvalFabric, EvalReply, EvalRequest,
7    Factory, Object, Result, ShapeId, Symbol, Value,
8};
9
10use super::execution::{
11    CpuTensorExecutor, TensorExecutor, TensorExecutorCard, tensor_executor_symbol,
12    tensor_executor_value, tensor_site_symbol,
13};
14
15/// Eval-fabric wrapper that binds one tensor executor for a realization.
16#[derive(Clone)]
17pub struct TensorSite {
18    symbol: Symbol,
19    executor: Arc<dyn TensorExecutor>,
20    capabilities: Arc<[CapabilityName]>,
21}
22
23impl TensorSite {
24    /// Builds a tensor site around an executor and required capability set.
25    pub fn new(
26        symbol: Symbol,
27        executor: Arc<dyn TensorExecutor>,
28        capabilities: Vec<CapabilityName>,
29    ) -> Self {
30        Self {
31            symbol,
32            executor,
33            capabilities: capabilities.into(),
34        }
35    }
36
37    /// Builds the default local tensor site using the CPU executor.
38    pub fn local_cpu() -> Self {
39        Self::new(
40            tensor_site_symbol(),
41            Arc::new(CpuTensorExecutor::new()),
42            Vec::new(),
43        )
44    }
45
46    /// Returns the site symbol.
47    pub fn symbol(&self) -> &Symbol {
48        &self.symbol
49    }
50
51    /// Returns the executor card exposed by this site.
52    pub fn card(&self) -> TensorExecutorCard {
53        self.executor.card()
54    }
55
56    /// Returns the capabilities this site requires for realization.
57    pub fn capabilities(&self) -> &[CapabilityName] {
58        &self.capabilities
59    }
60}
61
62impl EvalFabric for TensorSite {
63    fn realize(&self, cx: &mut Cx, request: EvalRequest) -> Result<EvalReply> {
64        cx.require_all(&self.capabilities)?;
65        cx.require_all(&request.required_capabilities)?;
66
67        let executor = tensor_executor_value(self.executor.clone())?;
68        let mut child = Env::child(Arc::new(cx.env().clone()));
69        child.define(tensor_executor_symbol(), executor);
70        let value = cx.with_env(child, |cx| cx.eval_expr(request.expr))?;
71        if let Some(shape_value) = request.result_shape.clone() {
72            check_result_shape(cx, &shape_value, value.clone())?;
73        }
74        Ok(EvalReply {
75            value,
76            diagnostics: Vec::new(),
77            trace: request.trace.then(|| {
78                DefaultFactory
79                    .symbol(Symbol::qualified("tensor", "trace/local"))
80                    .expect("trace symbol should be boxable")
81            }),
82        })
83    }
84}
85
86impl Object for TensorSite {
87    fn display(&self, _cx: &mut Cx) -> Result<String> {
88        Ok(format!("#<tensor-site {}>", self.symbol))
89    }
90
91    fn as_any(&self) -> &dyn std::any::Any {
92        self
93    }
94}
95
96impl sim_kernel::ObjectCompat for TensorSite {
97    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
98        if let Some(value) = cx
99            .registry()
100            .class_by_symbol(&Symbol::qualified("core", "EvalFabric"))
101        {
102            return Ok(value.clone());
103        }
104        DefaultFactory.class_stub(
105            sim_kernel::CORE_EVAL_REQUEST_CLASS_ID,
106            Symbol::qualified("core", "EvalFabric"),
107        )
108    }
109
110    fn as_eval_fabric(&self) -> Option<&dyn EvalFabric> {
111        Some(self)
112    }
113
114    fn as_table(&self, cx: &mut Cx) -> Result<Value> {
115        let card = self.card();
116        let operations = card
117            .operations
118            .iter()
119            .map(|symbol| cx.factory().symbol(symbol.clone()))
120            .collect::<Result<Vec<_>>>()?;
121        let device_capability = match card.device_capability {
122            Some(capability) => cx.factory().string(capability.as_str().to_owned())?,
123            None => cx.factory().nil()?,
124        };
125        cx.factory().table(vec![
126            (
127                Symbol::new("site"),
128                cx.factory().symbol(self.symbol.clone())?,
129            ),
130            (Symbol::new("executor"), cx.factory().symbol(card.symbol)?),
131            (Symbol::new("provider"), cx.factory().string(card.provider)?),
132            (Symbol::new("locality"), cx.factory().symbol(card.locality)?),
133            (Symbol::new("operations"), cx.factory().list(operations)?),
134            (Symbol::new("device-capability"), device_capability),
135        ])
136    }
137}
138
139fn check_result_shape(cx: &mut Cx, shape_value: &Value, value: Value) -> Result<()> {
140    let shape = shape_value.object().as_shape().ok_or(Error::TypeMismatch {
141        expected: "shape",
142        found: "non-shape",
143    })?;
144    let matched = shape.check_value(cx, value)?;
145    if matched.accepted {
146        Ok(())
147    } else {
148        Err(Error::WrongShape {
149            expected: shape.id().unwrap_or(ShapeId(0)),
150            diagnostics: matched.diagnostics,
151        })
152    }
153}