tenferro_extension_macros/
lib.rs

1//! Procedural macros for tenferro extension crates.
2//!
3//! # Examples
4//!
5//! ```
6//! use tenferro_extension_macros::ExtensionFamilyId;
7//!
8//! #[derive(ExtensionFamilyId)]
9//! #[tenferro_extension(namespace = "my-crate", name = "fft", version = 1)]
10//! struct FftOp;
11//!
12//! assert_eq!(FftOp::FAMILY_ID, "my-crate.fft.v1");
13//! ```
14
15use proc_macro::TokenStream;
16use quote::{format_ident, quote};
17use syn::parse::{Parse, ParseStream};
18use syn::{parse_macro_input, DeriveInput, Expr, ExprLit, Ident, Lit, Path, Token};
19
20#[derive(Debug, Default)]
21struct ExtensionArgs {
22    namespace: Option<String>,
23    name: Option<String>,
24    version: Option<u64>,
25}
26
27struct RuntimeArgs {
28    runtime: Ident,
29    family_id: Path,
30    op_type: Path,
31    execute_reads: Option<Path>,
32    execute_in_session: Option<Path>,
33    session_supported: Option<Path>,
34    backend_bound: Path,
35}
36
37impl Parse for ExtensionArgs {
38    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
39        let mut args = Self::default();
40        while !input.is_empty() {
41            let key: syn::Ident = input.parse()?;
42            input.parse::<Token![=]>()?;
43            let value: Expr = input.parse()?;
44            match key.to_string().as_str() {
45                "namespace" => args.namespace = Some(expect_string(value, "namespace")?),
46                "name" => args.name = Some(expect_string(value, "name")?),
47                "version" => args.version = Some(expect_u64(value, "version")?),
48                other => {
49                    return Err(syn::Error::new(
50                        key.span(),
51                        format!("unsupported tenferro_extension argument {other:?}"),
52                    ));
53                }
54            }
55            if input.is_empty() {
56                break;
57            }
58            input.parse::<Token![,]>()?;
59        }
60        Ok(args)
61    }
62}
63
64impl Parse for RuntimeArgs {
65    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
66        let mut runtime = None;
67        let mut family_id = None;
68        let mut op_type = None;
69        let mut execute_reads = None;
70        let mut execute_in_session = None;
71        let mut session_supported = None;
72        let mut backend_bound = None;
73
74        while !input.is_empty() {
75            let key: Ident = input.parse()?;
76            input.parse::<Token![=]>()?;
77            match key.to_string().as_str() {
78                "runtime" => runtime = Some(input.parse()?),
79                "family_id" => family_id = Some(input.parse()?),
80                "op_type" => op_type = Some(input.parse()?),
81                "execute_reads" => execute_reads = Some(input.parse()?),
82                "execute_in_session" => execute_in_session = Some(input.parse()?),
83                "session_supported" => session_supported = Some(input.parse()?),
84                "backend_bound" => backend_bound = Some(input.parse()?),
85                other => {
86                    return Err(syn::Error::new(
87                        key.span(),
88                        format!("unsupported define_extension_runtime argument {other:?}"),
89                    ));
90                }
91            }
92            if input.is_empty() {
93                break;
94            }
95            input.parse::<Token![,]>()?;
96        }
97
98        Ok(Self {
99            runtime: required(runtime, "runtime")?,
100            family_id: required(family_id, "family_id")?,
101            op_type: required(op_type, "op_type")?,
102            execute_reads,
103            execute_in_session,
104            session_supported,
105            backend_bound: backend_bound
106                .unwrap_or_else(|| syn::parse_quote!(tenferro_tensor::Engine)),
107        })
108    }
109}
110
111/// Derive an inherent `FAMILY_ID` constant for an extension payload type.
112///
113/// The required attribute is:
114/// `#[tenferro_extension(namespace = "...", version = N)]`.
115/// `name = "..."` is optional; when omitted, the Rust type name is converted
116/// to snake_case.
117#[proc_macro_derive(ExtensionFamilyId, attributes(tenferro_extension))]
118pub fn derive_extension_family_id(input: TokenStream) -> TokenStream {
119    let input = parse_macro_input!(input as DeriveInput);
120    match expand_extension_family_id(input) {
121        Ok(tokens) => tokens.into(),
122        Err(err) => err.to_compile_error().into(),
123    }
124}
125
126/// Generate a standard extension module, preparation engine, and prepared
127/// operation.
128///
129/// Exactly one execution route must be supplied.
130///
131/// * `execute_reads` runs the operation inside the runtime-formed context and
132///   has this signature:
133///   `fn<B: Engine + ?Sized>(&OpType, &[TensorRead<'_>], &mut ExtensionExecutionContext<'_, B>)`.
134///   The generated owner entry opens the backend session and the context
135///   borrows it (`B` is `dyn Engine`), so the callback never runs on
136///   the owner. Prefer the session route below for new extensions.
137/// * the session route (`execute_in_session` + `session_supported`) is the
138///   preferred one: `session_supported` has signature
139///   `fn<B: BackendBound + 'static>(&OpType) -> bool` and `execute_in_session`
140///   has signature
141///   `fn(&OpType, &mut dyn Engine, &mut ExtensionCacheStore, &[TensorRead<'_>])`.
142///   With it, the generated owner entry forms the session itself and hands the
143///   extension only a session, so no operation ever runs on the owner.
144///
145/// The removed legacy `execute` argument is rejected, not silently ignored.
146#[proc_macro]
147pub fn define_extension_runtime(input: TokenStream) -> TokenStream {
148    let args = parse_macro_input!(input as RuntimeArgs);
149    match expand_extension_runtime(args) {
150        Ok(tokens) => tokens.into(),
151        Err(err) => err.to_compile_error().into(),
152    }
153}
154
155fn expand_extension_family_id(input: DeriveInput) -> syn::Result<proc_macro2::TokenStream> {
156    let mut parsed = None;
157    for attr in &input.attrs {
158        if attr.path().is_ident("tenferro_extension") {
159            let args = attr.parse_args::<ExtensionArgs>()?;
160            parsed = Some(args);
161        }
162    }
163    let args = parsed.ok_or_else(|| {
164        syn::Error::new_spanned(
165            &input.ident,
166            "missing #[tenferro_extension(namespace = \"...\", version = N)]",
167        )
168    })?;
169    let namespace = args.namespace.ok_or_else(|| {
170        syn::Error::new_spanned(&input.ident, "missing tenferro_extension namespace")
171    })?;
172    let version = args.version.ok_or_else(|| {
173        syn::Error::new_spanned(&input.ident, "missing tenferro_extension version")
174    })?;
175    let name = args
176        .name
177        .unwrap_or_else(|| to_snake_case(&input.ident.to_string()));
178    let family_id = format!("{namespace}.{name}.v{version}");
179    let ident = input.ident;
180
181    Ok(quote! {
182        impl #ident {
183            /// Stable extension family identifier generated by `ExtensionFamilyId`.
184            pub const FAMILY_ID: &'static str = #family_id;
185        }
186    })
187}
188
189fn expand_extension_runtime(args: RuntimeArgs) -> syn::Result<proc_macro2::TokenStream> {
190    let RuntimeArgs {
191        runtime,
192        family_id,
193        op_type,
194        execute_reads,
195        execute_in_session,
196        session_supported,
197        backend_bound,
198    } = args;
199    let module = format_ident!("{}Module", runtime);
200    let planning_config = format_ident!("{}PlanningConfig", runtime);
201    let prepared_operation = format_ident!("{}PreparedOperation", runtime);
202    let owner_execute = match (&execute_reads, &execute_in_session) {
203        (None, Some(execute_in_session)) => quote! {
204            fn execute(
205                &self,
206                context: &mut tenferro_runtime::ErasedExecutionContext<'_>,
207                extension_caches: &mut tenferro_runtime::ExtensionCacheStore,
208                inputs: &[tenferro_tensor::TensorRead<'_>],
209            ) -> tenferro_runtime::Result<Vec<tenferro_tensor::Tensor>> {
210                // The context is the engine itself; the extension runs on it.
211                let engine = context
212                    .downcast_mut::<B>(self.binding.context_identity())
213                    .map_err(|source| tenferro_runtime::Error::runtime_state_source(
214                        "extension",
215                        tenferro_runtime::ErrorPhase::Execution,
216                        source,
217                    ))?;
218                let engine: &mut dyn tenferro_tensor::Engine = engine;
219                Ok(#execute_in_session(&self.op, engine, extension_caches, inputs)?)
220            }
221        },
222        (Some(execute_reads), None) => quote! {
223            fn execute(
224                &self,
225                context: &mut tenferro_runtime::ErasedExecutionContext<'_>,
226                extension_caches: &mut tenferro_runtime::ExtensionCacheStore,
227                inputs: &[tenferro_tensor::TensorRead<'_>],
228            ) -> tenferro_runtime::Result<Vec<tenferro_tensor::Tensor>> {
229                let engine = context
230                    .downcast_mut::<B>(self.binding.context_identity())
231                    .map_err(|source| tenferro_runtime::Error::runtime_state_source(
232                        "extension",
233                        tenferro_runtime::ErrorPhase::Execution,
234                        source,
235                    ))?;
236                let engine: &mut dyn tenferro_tensor::Engine = engine;
237                let mut ctx = tenferro_runtime::ExtensionExecutionContext::new(
238                    engine,
239                    extension_caches,
240                );
241                Ok(#execute_reads(&self.op, inputs, &mut ctx)?)
242            }
243        },
244        (Some(execute_reads), Some(_)) => {
245            return Err(syn::Error::new_spanned(
246                execute_reads,
247                "supply either execute_reads or execute_in_session, not both",
248            ))
249        }
250        (None, None) => {
251            return Err(syn::Error::new_spanned(
252                &op_type,
253                "either execute_reads or execute_in_session must be supplied",
254            ))
255        }
256    };
257    let session_methods = match (execute_in_session, session_supported) {
258        (Some(execute_in_session), Some(session_supported)) => quote! {
259            fn supports_session(&self) -> bool {
260                #session_supported::<B>(&self.op)
261            }
262
263            fn execute_in_session(
264                &self,
265                session: &mut dyn tenferro_tensor::Engine,
266                extension_caches: &mut tenferro_runtime::ExtensionCacheStore,
267                inputs: &[tenferro_tensor::TensorRead<'_>],
268            ) -> tenferro_runtime::Result<Vec<tenferro_tensor::Tensor>> {
269                Ok(#execute_in_session(&self.op, session, extension_caches, inputs)?)
270            }
271        },
272        (None, None) => quote! {},
273        (Some(path), None) | (None, Some(path)) => {
274            return Err(syn::Error::new_spanned(
275                path,
276                "execute_in_session and session_supported must be supplied together",
277            ))
278        }
279    };
280    Ok(quote! {
281        pub(crate) struct #runtime<B: #backend_bound + 'static> {
282            engine_id: tenferro_runtime::EngineId,
283            _backend: std::marker::PhantomData<fn() -> B>,
284        }
285
286        pub(crate) struct #module<B: #backend_bound + 'static> {
287            module_id: tenferro_runtime::ExtensionModuleId,
288            engine_id: tenferro_runtime::EngineId,
289            _backend: std::marker::PhantomData<fn() -> B>,
290        }
291
292        #[derive(Debug, Default)]
293        pub(crate) struct #planning_config;
294
295        pub(crate) struct #prepared_operation<B: #backend_bound + 'static> {
296            binding: tenferro_runtime::PreparedOperationBinding,
297            specialization: tenferro_runtime::SpecializationProjection,
298            op: #op_type,
299            _backend: std::marker::PhantomData<fn() -> B>,
300        }
301
302        impl<B: #backend_bound + 'static> std::fmt::Debug for #runtime<B> {
303            fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
304                formatter
305                    .debug_struct(stringify!(#runtime))
306                    .field("family_id", &#family_id)
307                    .field("engine_id", &self.engine_id)
308                    .field("backend_type", &std::any::type_name::<B>())
309                    .finish()
310            }
311        }
312
313        impl<B: #backend_bound + 'static> std::fmt::Debug for #module<B> {
314            fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
315                formatter
316                    .debug_struct(stringify!(#module))
317                    .field("module_id", &self.module_id)
318                    .field("engine_id", &self.engine_id)
319                    .field("backend_type", &std::any::type_name::<B>())
320                    .finish()
321            }
322        }
323
324        impl<B: #backend_bound + 'static> std::fmt::Debug for #prepared_operation<B> {
325            fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
326                formatter
327                    .debug_struct(stringify!(#prepared_operation))
328                    .field("family_id", &#family_id)
329                    .field("binding", &self.binding)
330                    .field("specialization", &self.specialization)
331                    .field("backend_type", &std::any::type_name::<B>())
332                    .finish_non_exhaustive()
333            }
334        }
335
336        impl<B: #backend_bound + 'static> tenferro_runtime::ExtensionEngine for #runtime<B>
337        where
338            #op_type: Clone + Send + Sync + 'static,
339        {
340            fn family_id(&self) -> &'static str {
341                #family_id
342            }
343
344            fn engine_id(&self) -> &tenferro_runtime::EngineId {
345                &self.engine_id
346            }
347
348            fn context_identity(&self) -> tenferro_runtime::ExecutionContextIdentity {
349                tenferro_runtime::ExecutionContextIdentity::of::<B>()
350            }
351
352            fn prepare(
353                &self,
354                request: tenferro_runtime::ExtensionPrepareRequest<'_>,
355            ) -> std::result::Result<tenferro_runtime::PrepareCapability, tenferro_runtime::PrepareError> {
356                let op = request
357                    .operation()
358                    .as_any()
359                    .downcast_ref::<#op_type>()
360                    .cloned()
361                    .ok_or_else(|| tenferro_runtime::PrepareError::ProviderContract {
362                        source: tenferro_runtime::ProviderContractError::WrongOperationFamily {
363                            expected: tenferro_runtime::CoreCapabilityKind::Elementwise,
364                            operation: #family_id,
365                        },
366                    })?;
367                let prepared = std::sync::Arc::new(#prepared_operation::<B> {
368                    binding: request.binding().clone(),
369                    specialization: request.specialization().clone(),
370                    op,
371                    _backend: std::marker::PhantomData,
372                });
373                Ok(tenferro_runtime::PrepareCapability::Prepared(
374                    tenferro_runtime::PreparedOperationPlan::executable(prepared.clone(), prepared)
375                ))
376            }
377        }
378
379        impl tenferro_runtime::ExtensionPlanningConfig for #planning_config {
380            fn family_id(&self) -> &'static str {
381                #family_id
382            }
383
384            fn as_any(&self) -> &dyn std::any::Any {
385                self
386            }
387
388            fn payload_hash(&self, state: &mut dyn std::hash::Hasher) {
389                state.write_u8(0);
390            }
391
392            fn payload_eq(&self, other: &dyn tenferro_runtime::ExtensionPlanningConfig) -> bool {
393                other.as_any().downcast_ref::<Self>().is_some()
394            }
395
396            fn retained_bytes(&self) -> usize {
397                0
398            }
399        }
400
401        impl<B: #backend_bound + 'static> tenferro_runtime::PreparedOperation for #prepared_operation<B>
402        where
403            #op_type: Clone + Send + Sync + 'static,
404        {
405            fn binding(&self) -> &tenferro_runtime::PreparedOperationBinding {
406                &self.binding
407            }
408
409            fn specialization(&self) -> &tenferro_runtime::SpecializationProjection {
410                &self.specialization
411            }
412
413            fn retained_bytes(&self) -> usize {
414                0
415            }
416
417        }
418
419        impl<B: #backend_bound + 'static> tenferro_runtime::PreparedOperationExecutor for #prepared_operation<B>
420        where
421            #op_type: Clone + Send + Sync + 'static,
422        {
423            #owner_execute
424
425            #session_methods
426        }
427
428        impl<B: #backend_bound + 'static> tenferro_runtime::ExtensionModule for #module<B>
429        where
430            #op_type: Clone + Send + Sync + 'static,
431        {
432            fn module_id(&self) -> &tenferro_runtime::ExtensionModuleId {
433                &self.module_id
434            }
435
436            fn configure(
437                &self,
438                registrar: &mut tenferro_runtime::ExtensionModuleRegistrar<'_>,
439            ) -> std::result::Result<(), tenferro_runtime::ExtensionModuleError> {
440                registrar.register_engine(std::sync::Arc::new(#runtime::<B> {
441                    engine_id: self.engine_id.clone(),
442                    _backend: std::marker::PhantomData,
443                }))?;
444                registrar.register_planning_config(
445                    self.engine_id.clone(),
446                    std::sync::Arc::new(#planning_config),
447                )?;
448                Ok(())
449            }
450        }
451
452        #[doc = "Build this extension module for one runtime engine."]
453        #[doc = "\n# Errors\n\nReturns `RuntimeConfigError::MalformedIdentity` when the generated module identifier is invalid."]
454        pub fn extension_module<B: #backend_bound + 'static>(
455            engine_id: tenferro_runtime::EngineId,
456        ) -> std::result::Result<
457            std::sync::Arc<dyn tenferro_runtime::ExtensionModule>,
458            tenferro_runtime::RuntimeConfigError,
459        >
460        where
461            #op_type: Clone + Send + Sync + 'static,
462        {
463            Ok(std::sync::Arc::new(#module::<B> {
464                module_id: tenferro_runtime::ExtensionModuleId::new(format!("{}.module", #family_id))?,
465                engine_id,
466                _backend: std::marker::PhantomData,
467            }))
468        }
469
470    })
471}
472
473fn expect_string(value: Expr, field: &str) -> syn::Result<String> {
474    match value {
475        Expr::Lit(ExprLit {
476            lit: Lit::Str(value),
477            ..
478        }) => Ok(value.value()),
479        other => Err(syn::Error::new_spanned(
480            other,
481            format!("{field} must be a string literal"),
482        )),
483    }
484}
485
486fn expect_u64(value: Expr, field: &str) -> syn::Result<u64> {
487    match value {
488        Expr::Lit(ExprLit {
489            lit: Lit::Int(value),
490            ..
491        }) => value.base10_parse(),
492        other => Err(syn::Error::new_spanned(
493            other,
494            format!("{field} must be an integer literal"),
495        )),
496    }
497}
498
499fn required<T>(value: Option<T>, field: &str) -> syn::Result<T> {
500    value.ok_or_else(|| syn::Error::new(proc_macro2::Span::call_site(), format!("missing {field}")))
501}
502
503fn to_snake_case(input: &str) -> String {
504    let mut out = String::new();
505    let mut prev_lower_or_digit = false;
506    for ch in input.chars() {
507        if ch.is_ascii_uppercase() {
508            if prev_lower_or_digit {
509                out.push('_');
510            }
511            out.push(ch.to_ascii_lowercase());
512            prev_lower_or_digit = false;
513        } else {
514            prev_lower_or_digit = ch.is_ascii_lowercase() || ch.is_ascii_digit();
515            out.push(ch);
516        }
517    }
518    out
519}
520
521#[cfg(test)]
522mod tests;
tenferro_extension_macros/lib.rs

tenferro_extension_macros/
lib.rs
