rustc_codegen_spirv/
attr.rs

1//! `#[spirv(...)]` attribute support.
2//!
3//! The attribute-checking parts of this try to follow `rustc_passes::check_attr`.
4
5use crate::codegen_cx::CodegenCx;
6use crate::symbols::Symbols;
7use rspirv::spirv::{BuiltIn, ExecutionMode, ExecutionModel, StorageClass};
8use rustc_ast::{LitKind, MetaItemInner, MetaItemLit};
9use rustc_hir as hir;
10use rustc_hir::def_id::LocalModDefId;
11use rustc_hir::intravisit::{self, Visitor};
12use rustc_hir::{Attribute, CRATE_HIR_ID, HirId, MethodKind, Target};
13use rustc_middle::hir::nested_filter;
14use rustc_middle::query::Providers;
15use rustc_middle::ty::TyCtxt;
16use rustc_span::{Ident, Span, Symbol};
17use smallvec::SmallVec;
18use std::rc::Rc;
19
20// FIXME(eddyb) replace with `ArrayVec<[Word; 3]>`.
21#[derive(Copy, Clone, Debug)]
22pub struct ExecutionModeExtra {
23    args: [u32; 3],
24    len: u8,
25}
26
27impl ExecutionModeExtra {
28    pub(crate) fn new(args: impl AsRef<[u32]>) -> Self {
29        let _args = args.as_ref();
30        let mut args = [0; 3];
31        args[.._args.len()].copy_from_slice(_args);
32        let len = _args.len() as u8;
33        Self { args, len }
34    }
35}
36
37impl AsRef<[u32]> for ExecutionModeExtra {
38    fn as_ref(&self) -> &[u32] {
39        &self.args[..self.len as _]
40    }
41}
42
43#[derive(Clone, Debug)]
44pub struct Entry {
45    pub execution_model: ExecutionModel,
46    pub execution_modes: Vec<(ExecutionMode, ExecutionModeExtra)>,
47    pub name: Option<Symbol>,
48}
49
50impl From<ExecutionModel> for Entry {
51    fn from(execution_model: ExecutionModel) -> Self {
52        Self {
53            execution_model,
54            execution_modes: Vec::new(),
55            name: None,
56        }
57    }
58}
59
60/// `struct` types that are used to represent special SPIR-V types.
61#[derive(Debug, Clone)]
62pub enum IntrinsicType {
63    GenericImageType,
64    Sampler,
65    AccelerationStructureKhr,
66    SampledImage,
67    RayQueryKhr,
68    RuntimeArray,
69    TypedBuffer,
70    Matrix,
71    Vector,
72    CooperativeMatrixKhr,
73}
74
75#[derive(Copy, Clone, Debug, PartialEq, Eq)]
76pub struct SpecConstant {
77    pub id: u32,
78    pub default: Option<u32>,
79    pub array_count: Option<u32>,
80}
81
82// NOTE(eddyb) when adding new `#[spirv(...)]` attributes, the tests found inside
83// `tests/ui/spirv-attr` should be updated (and new ones added if necessary).
84#[derive(Debug, Clone)]
85pub enum SpirvAttribute {
86    // `struct` attributes:
87    IntrinsicType(IntrinsicType),
88    Block,
89
90    // `fn` attributes:
91    Entry(Entry),
92
93    // (entry) `fn` parameter attributes:
94    StorageClass(StorageClass),
95    Builtin(BuiltIn),
96    DescriptorSet(u32),
97    Binding(u32),
98    Location(u32),
99    Flat,
100    PerPrimitiveExt,
101    Invariant,
102    InputAttachmentIndex(u32),
103    SpecConstant(SpecConstant),
104
105    // `fn`/closure attributes:
106    BufferLoadIntrinsic,
107    BufferStoreIntrinsic,
108}
109
110// HACK(eddyb) this is similar to `rustc_span::Spanned` but with `value` as the
111// field name instead of `node` (which feels inadequate in this context).
112#[derive(Copy, Clone)]
113pub struct Spanned<T> {
114    pub value: T,
115    pub span: Span,
116}
117
118/// Condensed version of a `SpirvAttribute` list, but only keeping one value per
119/// variant of `SpirvAttribute`, and treating multiple such attributes an error.
120// FIXME(eddyb) should this and `fn try_insert_attr` below be generated by a macro?
121#[derive(Default)]
122pub struct AggregatedSpirvAttributes {
123    // `struct` attributes:
124    pub intrinsic_type: Option<Spanned<IntrinsicType>>,
125    pub block: Option<Spanned<()>>,
126
127    // `fn` attributes:
128    pub entry: Option<Spanned<Entry>>,
129
130    // (entry) `fn` parameter attributes:
131    pub storage_class: Option<Spanned<StorageClass>>,
132    pub builtin: Option<Spanned<BuiltIn>>,
133    pub descriptor_set: Option<Spanned<u32>>,
134    pub binding: Option<Spanned<u32>>,
135    pub location: Option<Spanned<u32>>,
136    pub flat: Option<Spanned<()>>,
137    pub invariant: Option<Spanned<()>>,
138    pub per_primitive_ext: Option<Spanned<()>>,
139    pub input_attachment_index: Option<Spanned<u32>>,
140    pub spec_constant: Option<Spanned<SpecConstant>>,
141
142    // `fn`/closure attributes:
143    pub buffer_load_intrinsic: Option<Spanned<()>>,
144    pub buffer_store_intrinsic: Option<Spanned<()>>,
145}
146
147struct MultipleAttrs {
148    prev_span: Span,
149    category: &'static str,
150}
151
152impl AggregatedSpirvAttributes {
153    /// Compute `AggregatedSpirvAttributes` for use during codegen.
154    ///
155    /// Any errors for malformed/duplicate attributes will have been reported
156    /// prior to codegen, by the `attr` check pass.
157    pub fn parse<'tcx>(
158        cx: &CodegenCx<'tcx>,
159        attrs: impl IntoIterator<Item = &'tcx Attribute>,
160    ) -> Self {
161        let mut aggregated_attrs = Self::default();
162
163        // NOTE(eddyb) `span_delayed_bug` ensures that if attribute checking fails
164        // to see an attribute error, it will cause an ICE instead.
165        for parse_attr_result in parse_attrs_for_checking(&cx.sym, attrs) {
166            let (span, parsed_attr) = match parse_attr_result {
167                Ok(span_and_parsed_attr) => span_and_parsed_attr,
168                Err((span, msg)) => {
169                    cx.tcx.dcx().span_delayed_bug(span, msg);
170                    continue;
171                }
172            };
173            match aggregated_attrs.try_insert_attr(parsed_attr, span) {
174                Ok(()) => {}
175                Err(MultipleAttrs {
176                    prev_span: _,
177                    category,
178                }) => {
179                    cx.tcx
180                        .dcx()
181                        .span_delayed_bug(span, format!("multiple {category} attributes"));
182                }
183            }
184        }
185
186        aggregated_attrs
187    }
188
189    fn try_insert_attr(&mut self, attr: SpirvAttribute, span: Span) -> Result<(), MultipleAttrs> {
190        fn try_insert<T>(
191            slot: &mut Option<Spanned<T>>,
192            value: T,
193            span: Span,
194            category: &'static str,
195        ) -> Result<(), MultipleAttrs> {
196            if let Some(prev) = slot {
197                Err(MultipleAttrs {
198                    prev_span: prev.span,
199                    category,
200                })
201            } else {
202                *slot = Some(Spanned { value, span });
203                Ok(())
204            }
205        }
206
207        use SpirvAttribute::*;
208        match attr {
209            IntrinsicType(value) => {
210                try_insert(&mut self.intrinsic_type, value, span, "intrinsic type")
211            }
212            Block => try_insert(&mut self.block, (), span, "#[spirv(block)]"),
213            Entry(value) => try_insert(&mut self.entry, value, span, "entry-point"),
214            StorageClass(value) => {
215                try_insert(&mut self.storage_class, value, span, "storage class")
216            }
217            Builtin(value) => try_insert(&mut self.builtin, value, span, "builtin"),
218            DescriptorSet(value) => try_insert(
219                &mut self.descriptor_set,
220                value,
221                span,
222                "#[spirv(descriptor_set)]",
223            ),
224            Binding(value) => try_insert(&mut self.binding, value, span, "#[spirv(binding)]"),
225            Location(value) => try_insert(&mut self.location, value, span, "#[spirv(location)]"),
226            Flat => try_insert(&mut self.flat, (), span, "#[spirv(flat)]"),
227            Invariant => try_insert(&mut self.invariant, (), span, "#[spirv(invariant)]"),
228            PerPrimitiveExt => try_insert(
229                &mut self.per_primitive_ext,
230                (),
231                span,
232                "#[spirv(per_primitive_ext)]",
233            ),
234            InputAttachmentIndex(value) => try_insert(
235                &mut self.input_attachment_index,
236                value,
237                span,
238                "#[spirv(attachment_index)]",
239            ),
240            SpecConstant(value) => try_insert(
241                &mut self.spec_constant,
242                value,
243                span,
244                "#[spirv(spec_constant)]",
245            ),
246            BufferLoadIntrinsic => try_insert(
247                &mut self.buffer_load_intrinsic,
248                (),
249                span,
250                "#[spirv(buffer_load_intrinsic)]",
251            ),
252            BufferStoreIntrinsic => try_insert(
253                &mut self.buffer_store_intrinsic,
254                (),
255                span,
256                "#[spirv(buffer_store_intrinsic)]",
257            ),
258        }
259    }
260}
261
262// FIXME(eddyb) make this reusable from somewhere in `rustc`.
263fn target_from_impl_item(tcx: TyCtxt<'_>, impl_item: &hir::ImplItem<'_>) -> Target {
264    match impl_item.kind {
265        hir::ImplItemKind::Const(..) => Target::AssocConst,
266        hir::ImplItemKind::Fn(..) => {
267            let parent_owner_id = tcx.hir_get_parent_item(impl_item.hir_id());
268            let containing_item = tcx.hir_expect_item(parent_owner_id.def_id);
269            let containing_impl_is_for_trait = match &containing_item.kind {
270                hir::ItemKind::Impl(hir::Impl { of_trait, .. }) => of_trait.is_some(),
271                _ => unreachable!("parent of an ImplItem must be an Impl"),
272            };
273            if containing_impl_is_for_trait {
274                Target::Method(MethodKind::Trait { body: true })
275            } else {
276                Target::Method(MethodKind::Inherent)
277            }
278        }
279        hir::ImplItemKind::Type(..) => Target::AssocTy,
280    }
281}
282
283struct CheckSpirvAttrVisitor<'tcx> {
284    tcx: TyCtxt<'tcx>,
285    sym: Rc<Symbols>,
286}
287
288impl CheckSpirvAttrVisitor<'_> {
289    fn check_spirv_attributes(&self, hir_id: HirId, target: Target) {
290        let mut aggregated_attrs = AggregatedSpirvAttributes::default();
291
292        let parse_attrs = |attrs| parse_attrs_for_checking(&self.sym, attrs);
293
294        let attrs = self.tcx.hir_attrs(hir_id);
295        for parse_attr_result in parse_attrs(attrs) {
296            let (span, parsed_attr) = match parse_attr_result {
297                Ok(span_and_parsed_attr) => span_and_parsed_attr,
298                Err((span, msg)) => {
299                    self.tcx.dcx().span_err(span, msg);
300                    continue;
301                }
302            };
303
304            /// Error newtype marker used below for readability.
305            struct Expected<T>(T);
306
307            let valid_target = match parsed_attr {
308                SpirvAttribute::IntrinsicType(_) | SpirvAttribute::Block => match target {
309                    Target::Struct => {
310                        // FIXME(eddyb) further check type attribute validity,
311                        // e.g. layout, generics, other attributes, etc.
312                        Ok(())
313                    }
314
315                    _ => Err(Expected("struct")),
316                },
317
318                SpirvAttribute::Entry(_) => match target {
319                    Target::Fn
320                    | Target::Method(MethodKind::Trait { body: true } | MethodKind::Inherent) => {
321                        // FIXME(eddyb) further check entry-point attribute validity,
322                        // e.g. signature, shouldn't have `#[inline]` or generics, etc.
323                        Ok(())
324                    }
325
326                    _ => Err(Expected("function")),
327                },
328
329                SpirvAttribute::StorageClass(_)
330                | SpirvAttribute::Builtin(_)
331                | SpirvAttribute::DescriptorSet(_)
332                | SpirvAttribute::Binding(_)
333                | SpirvAttribute::Location(_)
334                | SpirvAttribute::Flat
335                | SpirvAttribute::Invariant
336                | SpirvAttribute::PerPrimitiveExt
337                | SpirvAttribute::InputAttachmentIndex(_)
338                | SpirvAttribute::SpecConstant(_) => match target {
339                    Target::Param => {
340                        let parent_hir_id = self.tcx.parent_hir_id(hir_id);
341                        let parent_is_entry_point = parse_attrs(self.tcx.hir_attrs(parent_hir_id))
342                            .filter_map(|r| r.ok())
343                            .any(|(_, attr)| matches!(attr, SpirvAttribute::Entry(_)));
344                        if !parent_is_entry_point {
345                            self.tcx.dcx().span_err(
346                                span,
347                                "attribute is only valid on a parameter of an entry-point function",
348                            );
349                        } else {
350                            // FIXME(eddyb) should we just remove all 5 of these storage class
351                            // attributes, instead of disallowing them here?
352                            if let SpirvAttribute::StorageClass(storage_class) = parsed_attr {
353                                let valid = match storage_class {
354                                    StorageClass::Input | StorageClass::Output => {
355                                        Err("is the default and should not be explicitly specified")
356                                    }
357
358                                    StorageClass::Private
359                                    | StorageClass::Function
360                                    | StorageClass::Generic => {
361                                        Err("can not be used as part of an entry's interface")
362                                    }
363
364                                    _ => Ok(()),
365                                };
366
367                                if let Err(msg) = valid {
368                                    self.tcx.dcx().span_err(
369                                        span,
370                                        format!("`{storage_class:?}` storage class {msg}"),
371                                    );
372                                }
373                            }
374                        }
375                        Ok(())
376                    }
377
378                    _ => Err(Expected("function parameter")),
379                },
380                SpirvAttribute::BufferLoadIntrinsic | SpirvAttribute::BufferStoreIntrinsic => {
381                    match target {
382                        Target::Fn => Ok(()),
383                        _ => Err(Expected("function")),
384                    }
385                }
386            };
387            match valid_target {
388                Err(Expected(expected_target)) => {
389                    self.tcx.dcx().span_err(
390                        span,
391                        format!(
392                            "attribute is only valid on a {expected_target}, not on a {target}"
393                        ),
394                    );
395                }
396                Ok(()) => match aggregated_attrs.try_insert_attr(parsed_attr, span) {
397                    Ok(()) => {}
398                    Err(MultipleAttrs {
399                        prev_span,
400                        category,
401                    }) => {
402                        self.tcx
403                            .dcx()
404                            .struct_span_err(
405                                span,
406                                format!("only one {category} attribute is allowed on a {target}"),
407                            )
408                            .with_span_note(prev_span, format!("previous {category} attribute"))
409                            .emit();
410                    }
411                },
412            }
413        }
414
415        // At this point we have all of the attributes (valid for this target),
416        // so we can perform further checks, emit warnings, etc.
417
418        if let Some(block_attr) = aggregated_attrs.block {
419            self.tcx.dcx().span_warn(
420                block_attr.span,
421                "#[spirv(block)] is no longer needed and should be removed",
422            );
423        }
424    }
425}
426
427// FIXME(eddyb) DRY this somehow and make it reusable from somewhere in `rustc`.
428impl<'tcx> Visitor<'tcx> for CheckSpirvAttrVisitor<'tcx> {
429    type NestedFilter = nested_filter::OnlyBodies;
430
431    fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
432        self.tcx
433    }
434
435    fn visit_item(&mut self, item: &'tcx hir::Item<'tcx>) {
436        let target = Target::from_item(item);
437        self.check_spirv_attributes(item.hir_id(), target);
438        intravisit::walk_item(self, item);
439    }
440
441    fn visit_generic_param(&mut self, generic_param: &'tcx hir::GenericParam<'tcx>) {
442        let target = Target::from_generic_param(generic_param);
443        self.check_spirv_attributes(generic_param.hir_id, target);
444        intravisit::walk_generic_param(self, generic_param);
445    }
446
447    fn visit_trait_item(&mut self, trait_item: &'tcx hir::TraitItem<'tcx>) {
448        let target = Target::from_trait_item(trait_item);
449        self.check_spirv_attributes(trait_item.hir_id(), target);
450        intravisit::walk_trait_item(self, trait_item);
451    }
452
453    fn visit_field_def(&mut self, field: &'tcx hir::FieldDef<'tcx>) {
454        self.check_spirv_attributes(field.hir_id, Target::Field);
455        intravisit::walk_field_def(self, field);
456    }
457
458    fn visit_arm(&mut self, arm: &'tcx hir::Arm<'tcx>) {
459        self.check_spirv_attributes(arm.hir_id, Target::Arm);
460        intravisit::walk_arm(self, arm);
461    }
462
463    fn visit_foreign_item(&mut self, f_item: &'tcx hir::ForeignItem<'tcx>) {
464        let target = Target::from_foreign_item(f_item);
465        self.check_spirv_attributes(f_item.hir_id(), target);
466        intravisit::walk_foreign_item(self, f_item);
467    }
468
469    fn visit_impl_item(&mut self, impl_item: &'tcx hir::ImplItem<'tcx>) {
470        let target = target_from_impl_item(self.tcx, impl_item);
471        self.check_spirv_attributes(impl_item.hir_id(), target);
472        intravisit::walk_impl_item(self, impl_item);
473    }
474
475    fn visit_stmt(&mut self, stmt: &'tcx hir::Stmt<'tcx>) {
476        // When checking statements ignore expressions, they will be checked later.
477        if let hir::StmtKind::Let(l) = stmt.kind {
478            self.check_spirv_attributes(l.hir_id, Target::Statement);
479        }
480        intravisit::walk_stmt(self, stmt);
481    }
482
483    fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
484        let target = match expr.kind {
485            hir::ExprKind::Closure { .. } => Target::Closure,
486            _ => Target::Expression,
487        };
488
489        self.check_spirv_attributes(expr.hir_id, target);
490        intravisit::walk_expr(self, expr);
491    }
492
493    fn visit_variant(&mut self, variant: &'tcx hir::Variant<'tcx>) {
494        self.check_spirv_attributes(variant.hir_id, Target::Variant);
495        intravisit::walk_variant(self, variant);
496    }
497
498    fn visit_param(&mut self, param: &'tcx hir::Param<'tcx>) {
499        self.check_spirv_attributes(param.hir_id, Target::Param);
500
501        intravisit::walk_param(self, param);
502    }
503}
504
505// FIXME(eddyb) DRY this somehow and make it reusable from somewhere in `rustc`.
506fn check_mod_attrs(tcx: TyCtxt<'_>, module_def_id: LocalModDefId) {
507    let check_spirv_attr_visitor = &mut CheckSpirvAttrVisitor {
508        tcx,
509        sym: Symbols::get(),
510    };
511    tcx.hir_visit_item_likes_in_module(module_def_id, check_spirv_attr_visitor);
512    if module_def_id.is_top_level_module() {
513        check_spirv_attr_visitor.check_spirv_attributes(CRATE_HIR_ID, Target::Mod);
514    }
515}
516
517pub(crate) fn provide(providers: &mut Providers) {
518    *providers = Providers {
519        check_mod_attrs: |tcx, module_def_id| {
520            // Run both the default checks, and our `#[spirv(...)]` ones.
521            (rustc_interface::DEFAULT_QUERY_PROVIDERS
522                .queries
523                .check_mod_attrs)(tcx, module_def_id);
524            check_mod_attrs(tcx, module_def_id);
525        },
526        // HACK(LegNeato) keep `#[spirv(...)]` entry-points from getting
527        // internalized and DCE'd before codegen. See issue #590.
528        cross_crate_inlinable: |tcx, def_id| {
529            let sym = Symbols::get();
530            let path = [sym.rust_gpu, sym.spirv_attr_with_version];
531            let attrs: Vec<_> = tcx.get_attrs_by_path(def_id.to_def_id(), &path).collect();
532            let is_entry_point = parse_attrs_for_checking(&sym, attrs)
533                .any(|result| matches!(result, Ok((_, SpirvAttribute::Entry(_)))));
534            if is_entry_point {
535                false
536            } else {
537                (rustc_interface::DEFAULT_QUERY_PROVIDERS
538                    .queries
539                    .cross_crate_inlinable)(tcx, def_id)
540            }
541        },
542        ..*providers
543    };
544}
545
546// FIXME(eddyb) find something nicer for the error type.
547type ParseAttrError = (Span, String);
548
549#[allow(clippy::get_first)]
550fn parse_attrs_for_checking<'sym, 'attr, I>(
551    sym: &'sym Symbols,
552    attrs: I,
553) -> impl Iterator<Item = Result<(Span, SpirvAttribute), ParseAttrError>> + 'sym
554where
555    I: IntoIterator<Item = &'attr Attribute> + 'sym,
556    I::IntoIter: 'sym,
557    'attr: 'sym,
558{
559    attrs
560        .into_iter()
561        .map(move |attr| {
562            // parse the #[rust_gpu::spirv(...)] attr and return the inner list
563            match attr {
564                Attribute::Unparsed(item) => {
565                    // #[...]
566                    let s = &item.path.segments;
567                    if let Some(rust_gpu) = s.get(0) && *rust_gpu == sym.rust_gpu {
568                        // #[rust_gpu ...]
569                        match s.get(1) {
570                            Some(command) if *command == sym.spirv_attr_with_version => {
571                                // #[rust_gpu::spirv ...]
572                                if let Some(args) = attr.meta_item_list() {
573                                    // #[rust_gpu::spirv(...)]
574                                    Ok(parse_spirv_attr(sym, args.iter()))
575                                } else {
576                                    // #[rust_gpu::spirv]
577                                    Err((
578                                        attr.span(),
579                                        "#[spirv(..)] attribute must have at least one argument"
580                                            .to_string(),
581                                    ))
582                                }
583                            }
584                            Some(command) if *command == sym.vector => {
585                                // #[rust_gpu::vector ...]
586                                match s.get(2) {
587                                    // #[rust_gpu::vector::v1]
588                                    Some(version) if *version == sym.v1 => {
589                                        Ok(SmallVec::from_iter([
590                                            Ok((attr.span(), SpirvAttribute::IntrinsicType(IntrinsicType::Vector)))
591                                        ]))
592                                    },
593                                    _ => Err((
594                                        attr.span(),
595                                        "unknown `rust_gpu::vector` version, expected `rust_gpu::vector::v1`"
596                                            .to_string(),
597                                    )),
598                                }
599                            }
600                            _ => {
601                                // #[rust_gpu::...] but not a know version
602                                let spirv = sym.spirv_attr_with_version.as_str();
603                                Err((
604                                    attr.span(),
605                                    format!("unknown `rust_gpu` attribute, expected `rust_gpu::{spirv}`. \
606                                Do the versions of `spirv-std` and `rustc_codegen_spirv` match?"),
607                                ))
608                            }
609                        }
610                    } else {
611                        // #[...] but not #[rust_gpu ...]
612                        Ok(Default::default())
613                    }
614                }
615                Attribute::Parsed(_) => Ok(Default::default()),
616            }
617        })
618        .flat_map(|result| {
619            result
620                .unwrap_or_else(|err| SmallVec::from_iter([Err(err)]))
621                .into_iter()
622        })
623}
624
625fn parse_spirv_attr<'a>(
626    sym: &Symbols,
627    iter: impl Iterator<Item = &'a MetaItemInner>,
628) -> SmallVec<[Result<(Span, SpirvAttribute), ParseAttrError>; 4]> {
629    iter.map(|arg| {
630        let span = arg.span();
631        let parsed_attr =
632            if arg.has_name(sym.descriptor_set) {
633                SpirvAttribute::DescriptorSet(parse_attr_int_value(arg)?)
634            } else if arg.has_name(sym.binding) {
635                SpirvAttribute::Binding(parse_attr_int_value(arg)?)
636            } else if arg.has_name(sym.location) {
637                SpirvAttribute::Location(parse_attr_int_value(arg)?)
638            } else if arg.has_name(sym.input_attachment_index) {
639                SpirvAttribute::InputAttachmentIndex(parse_attr_int_value(arg)?)
640            } else if arg.has_name(sym.spec_constant) {
641                SpirvAttribute::SpecConstant(parse_spec_constant_attr(sym, arg)?)
642            } else {
643                let name = match arg.ident() {
644                    Some(i) => i,
645                    None => {
646                        return Err((
647                            span,
648                            "#[spirv(..)] attribute argument must be single identifier".to_string(),
649                        ));
650                    }
651                };
652                sym.attributes.get(&name.name).map_or_else(
653                    || Err((name.span, "unknown argument to spirv attribute".to_string())),
654                    |a| {
655                        Ok(match a {
656                            SpirvAttribute::Entry(entry) => SpirvAttribute::Entry(
657                                parse_entry_attrs(sym, arg, &name, entry.execution_model)?,
658                            ),
659                            _ => a.clone(),
660                        })
661                    },
662                )?
663            };
664        Ok((span, parsed_attr))
665    })
666    .collect()
667}
668
669fn parse_spec_constant_attr(
670    sym: &Symbols,
671    arg: &MetaItemInner,
672) -> Result<SpecConstant, ParseAttrError> {
673    let mut id = None;
674    let mut default = None;
675
676    if let Some(attrs) = arg.meta_item_list() {
677        for attr in attrs {
678            if attr.has_name(sym.id) {
679                if id.is_none() {
680                    id = Some(parse_attr_int_value(attr)?);
681                } else {
682                    return Err((attr.span(), "`id` may only be specified once".into()));
683                }
684            } else if attr.has_name(sym.default) {
685                if default.is_none() {
686                    default = Some(parse_attr_int_value(attr)?);
687                } else {
688                    return Err((attr.span(), "`default` may only be specified once".into()));
689                }
690            } else {
691                return Err((attr.span(), "expected `id = ...` or `default = ...`".into()));
692            }
693        }
694    }
695    Ok(SpecConstant {
696        id: id.ok_or_else(|| (arg.span(), "expected `spec_constant(id = ...)`".into()))?,
697        default,
698        // to be set later
699        array_count: None,
700    })
701}
702
703fn parse_attr_int_value(arg: &MetaItemInner) -> Result<u32, ParseAttrError> {
704    let arg = match arg.meta_item() {
705        Some(arg) => arg,
706        None => return Err((arg.span(), "attribute must have value".to_string())),
707    };
708    match arg.name_value_literal() {
709        Some(&MetaItemLit {
710            kind: LitKind::Int(x, ..),
711            ..
712        }) if x <= u32::MAX as u128 => Ok(x.get() as u32),
713        _ => Err((arg.span, "attribute value must be integer".to_string())),
714    }
715}
716
717fn parse_local_size_attr(arg: &MetaItemInner) -> Result<[u32; 3], ParseAttrError> {
718    let arg = match arg.meta_item() {
719        Some(arg) => arg,
720        None => return Err((arg.span(), "attribute must have value".to_string())),
721    };
722    match arg.meta_item_list() {
723        Some(tuple) if !tuple.is_empty() && tuple.len() < 4 => {
724            let mut local_size = [1; 3];
725            for (idx, lit) in tuple.iter().enumerate() {
726                match lit {
727                    MetaItemInner::Lit(MetaItemLit {
728                                           kind: LitKind::Int(x, ..),
729                                           ..
730                                       }) if *x <= u32::MAX as u128 => local_size[idx] = x.get() as u32,
731                    _ => return Err((lit.span(), "must be a u32 literal".to_string())),
732                }
733            }
734            Ok(local_size)
735        }
736        Some([]) => Err((
737            arg.span,
738            "#[spirv(compute(threads(x, y, z)))] must have the x dimension specified, trailing ones may be elided".to_string(),
739        )),
740        Some(tuple) if tuple.len() > 3 => Err((
741            arg.span,
742            "#[spirv(compute(threads(x, y, z)))] is three dimensional".to_string(),
743        )),
744        _ => Err((
745            arg.span,
746            "#[spirv(compute(threads(x, y, z)))] must have 1 to 3 parameters, trailing ones may be elided".to_string(),
747        )),
748    }
749}
750
751// for a given entry, gather up the additional attributes
752// in this case ExecutionMode's, some have extra arguments
753// others are specified with x, y, or z components
754// ie #[spirv(fragment(origin_lower_left))] or #[spirv(gl_compute(local_size_x=64, local_size_y=8))]
755fn parse_entry_attrs(
756    sym: &Symbols,
757    arg: &MetaItemInner,
758    name: &Ident,
759    execution_model: ExecutionModel,
760) -> Result<Entry, ParseAttrError> {
761    use ExecutionMode::*;
762    use ExecutionModel::*;
763    let mut entry = Entry::from(execution_model);
764    let mut origin_mode: Option<ExecutionMode> = None;
765    let mut local_size: Option<[u32; 3]> = None;
766    let mut local_size_hint: Option<[u32; 3]> = None;
767    // Reserved
768    //let mut max_workgroup_size_intel: Option<[u32; 3]> = None;
769    if let Some(attrs) = arg.meta_item_list() {
770        for attr in attrs {
771            if let Some(attr_name) = attr.ident() {
772                if let Some((execution_mode, extra_dim)) = sym.execution_modes.get(&attr_name.name)
773                {
774                    use crate::symbols::ExecutionModeExtraDim::*;
775                    let val = match extra_dim {
776                        None | Tuple => Option::None,
777                        _ => Some(parse_attr_int_value(attr)?),
778                    };
779                    match execution_mode {
780                        OriginUpperLeft | OriginLowerLeft => {
781                            origin_mode.replace(*execution_mode);
782                        }
783                        LocalSize => {
784                            if local_size.is_none() {
785                                local_size.replace(parse_local_size_attr(attr)?);
786                            } else {
787                                return Err((
788                                    attr_name.span,
789                                    String::from(
790                                        "`#[spirv(compute(threads))]` may only be specified once",
791                                    ),
792                                ));
793                            }
794                        }
795                        LocalSizeHint => {
796                            let val = val.unwrap();
797                            if local_size_hint.is_none() {
798                                local_size_hint.replace([1, 1, 1]);
799                            }
800                            let local_size_hint = local_size_hint.as_mut().unwrap();
801                            match extra_dim {
802                                X => {
803                                    local_size_hint[0] = val;
804                                }
805                                Y => {
806                                    local_size_hint[1] = val;
807                                }
808                                Z => {
809                                    local_size_hint[2] = val;
810                                }
811                                _ => unreachable!(),
812                            }
813                        }
814                        // Reserved
815                        /*MaxWorkgroupSizeINTEL => {
816                            let val = val.unwrap();
817                            if max_workgroup_size_intel.is_none() {
818                                max_workgroup_size_intel.replace([1, 1, 1]);
819                            }
820                            let max_workgroup_size_intel = max_workgroup_size_intel.as_mut()
821                                .unwrap();
822                            match extra_dim {
823                                X => {
824                                    max_workgroup_size_intel[0] = val;
825                                },
826                                Y => {
827                                    max_workgroup_size_intel[1] = val;
828                                },
829                                Z => {
830                                    max_workgroup_size_intel[2] = val;
831                                },
832                                _ => unreachable!(),
833                            }
834                        },*/
835                        _ => {
836                            if let Some(val) = val {
837                                entry
838                                    .execution_modes
839                                    .push((*execution_mode, ExecutionModeExtra::new([val])));
840                            } else {
841                                entry
842                                    .execution_modes
843                                    .push((*execution_mode, ExecutionModeExtra::new([])));
844                            }
845                        }
846                    }
847                } else if attr_name.name == sym.entry_point_name {
848                    match attr.value_str() {
849                        Some(sym) => {
850                            entry.name = Some(sym);
851                        }
852                        None => {
853                            return Err((
854                                attr_name.span,
855                                format!(
856                                    "#[spirv({name}(..))] unknown attribute argument {attr_name}"
857                                ),
858                            ));
859                        }
860                    }
861                } else {
862                    return Err((
863                        attr_name.span,
864                        format!("#[spirv({name}(..))] unknown attribute argument {attr_name}",),
865                    ));
866                }
867            } else {
868                return Err((
869                    arg.span(),
870                    format!("#[spirv({name}(..))] attribute argument must be single identifier"),
871                ));
872            }
873        }
874    }
875    match entry.execution_model {
876        Fragment => {
877            let origin_mode = origin_mode.unwrap_or(OriginUpperLeft);
878            entry
879                .execution_modes
880                .push((origin_mode, ExecutionModeExtra::new([])));
881        }
882        GLCompute | MeshNV | TaskNV | TaskEXT | MeshEXT => {
883            if let Some(local_size) = local_size {
884                entry
885                    .execution_modes
886                    .push((LocalSize, ExecutionModeExtra::new(local_size)));
887            } else {
888                return Err((
889                    arg.span(),
890                    String::from(
891                        "The `threads` argument must be specified when using `#[spirv(compute)]`, `#[spirv(mesh_nv)]`, `#[spirv(task_nv)]`, `#[spirv(task_ext)]` or `#[spirv(mesh_ext)]`",
892                    ),
893                ));
894            }
895        }
896        //TODO: Cover more defaults
897        _ => {}
898    }
899    Ok(entry)
900}
rustc_codegen_spirv/attr.rs

rustc_codegen_spirv/
attr.rs
