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

tenferro_extension_macros/
lib.rs
