sim_lib_numbers_tensor/implementation/
tensor_site.rs1use 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#[derive(Clone)]
17pub struct TensorSite {
18 symbol: Symbol,
19 executor: Arc<dyn TensorExecutor>,
20 capabilities: Arc<[CapabilityName]>,
21}
22
23impl TensorSite {
24 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 pub fn local_cpu() -> Self {
39 Self::new(
40 tensor_site_symbol(),
41 Arc::new(CpuTensorExecutor::new()),
42 Vec::new(),
43 )
44 }
45
46 pub fn symbol(&self) -> &Symbol {
48 &self.symbol
49 }
50
51 pub fn card(&self) -> TensorExecutorCard {
53 self.executor.card()
54 }
55
56 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}