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

tenferro_extension_macros/
lib.rs
