rustc_monomorphize/offload/
manifest.rs

1//! Offload manifest: communicates required generic kernel instantiations
2//! between host-metadata and device compilation passes.
3//!
4//! Uses `TyEncoder`/`TyDecoder` to serialize `ty::Instance`. DefIds are
5//! encoded as (crate name, DefPath) pairs for stability.
6
7use std::fs;
8
9use rustc_data_structures::fx::FxHashMap;
10use rustc_data_structures::sync::Lock;
11use rustc_hir::def_id::{DefId, DefIndex, LOCAL_CRATE, StableCrateId};
12use rustc_middle::bug;
13use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
14use rustc_middle::mono::MonoItem;
15use rustc_middle::ty::codec::{TyDecoder, TyEncoder};
16use rustc_middle::ty::{self, Ty, TyCtxt};
17use rustc_serialize::opaque::{FileEncoder, MemDecoder};
18use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
19use rustc_span::{
20    BlobDecoder, BytePos, ByteSymbol, Pos, Span, SpanDecoder, SpanEncoder, Symbol, SyntaxContext,
21};
22
23pub(crate) struct OffloadManifestEncoder<'a, 'tcx> {
24    encoder: FileEncoder<'a>,
25    type_shorthands: FxHashMap<Ty<'tcx>, usize>,
26    predicate_shorthands: FxHashMap<ty::PredicateKind<'tcx>, usize>,
27    tcx: TyCtxt<'tcx>,
28}
29
30impl<'a, 'tcx> OffloadManifestEncoder<'a, 'tcx> {
31    pub(crate) fn new(path: &'a std::path::Path, tcx: TyCtxt<'tcx>) -> std::io::Result<Self> {
32        let encoder = FileEncoder::new(path)?;
33        Ok(OffloadManifestEncoder {
34            encoder,
35            type_shorthands: FxHashMap::default(),
36            predicate_shorthands: FxHashMap::default(),
37            tcx,
38        })
39    }
40
41    pub(crate) fn finish(mut self) -> std::io::Result<()> {
42        self.encoder.finish().map(|_| ()).map_err(|(_, e)| e)
43    }
44}
45
46impl<'a, 'tcx> Encoder for OffloadManifestEncoder<'a, 'tcx> {
47    fn emit_usize(&mut self, v: usize) {
48        self.encoder.emit_usize(v);
49    }
50    fn emit_u128(&mut self, v: u128) {
51        self.encoder.emit_u128(v);
52    }
53    fn emit_u64(&mut self, v: u64) {
54        self.encoder.emit_u64(v);
55    }
56    fn emit_u32(&mut self, v: u32) {
57        self.encoder.emit_u32(v);
58    }
59    fn emit_u16(&mut self, v: u16) {
60        self.encoder.emit_u16(v);
61    }
62    fn emit_u8(&mut self, v: u8) {
63        self.encoder.emit_u8(v);
64    }
65    fn emit_isize(&mut self, v: isize) {
66        self.encoder.emit_isize(v);
67    }
68    fn emit_i128(&mut self, v: i128) {
69        self.encoder.emit_i128(v);
70    }
71    fn emit_i64(&mut self, v: i64) {
72        self.encoder.emit_i64(v);
73    }
74    fn emit_i32(&mut self, v: i32) {
75        self.encoder.emit_i32(v);
76    }
77    fn emit_i16(&mut self, v: i16) {
78        self.encoder.emit_i16(v);
79    }
80    fn emit_i8(&mut self, v: i8) {
81        self.encoder.emit_i8(v);
82    }
83    fn emit_raw_bytes(&mut self, v: &[u8]) {
84        self.encoder.emit_raw_bytes(v);
85    }
86}
87
88impl<'a, 'tcx> SpanEncoder for OffloadManifestEncoder<'a, 'tcx> {
89    fn encode_span(&mut self, _span: Span) {
90        // Spans are not needed in the manifest, encode a dummy span.
91        self.emit_usize(0);
92        self.emit_usize(0);
93        self.emit_u32(0);
94    }
95
96    fn encode_symbol(&mut self, sym: rustc_span::Symbol) {
97        sym.as_str().encode(self);
98    }
99
100    fn encode_byte_symbol(&mut self, byte_sym: ByteSymbol) {
101        let bytes = byte_sym.as_byte_str();
102        if true {
    if !bytes.is_empty() {
        {
            ::core::panicking::panic_fmt(format_args!("ByteSymbols with content are not expected in offload manifests"));
        }
    };
};debug_assert!(
103            bytes.is_empty(),
104            "ByteSymbols with content are not expected in offload manifests"
105        );
106        self.emit_usize(bytes.len());
107        self.emit_raw_bytes(bytes.as_ref())
108    }
109
110    fn encode_expn_id(&mut self, _expn_id: rustc_span::ExpnId) {
111        self.emit_u32(0);
112    }
113
114    fn encode_syntax_context(&mut self, _syntax_context: SyntaxContext) {
115        self.emit_u32(0);
116    }
117
118    fn encode_crate_num(&mut self, crate_num: rustc_span::def_id::CrateNum) {
119        self.tcx.stable_crate_id(crate_num).encode(self);
120    }
121
122    fn encode_def_index(&mut self, def_index: rustc_span::def_id::DefIndex) {
123        def_index.as_u32().encode(self);
124    }
125
126    fn encode_def_id(&mut self, def_id: rustc_span::def_id::DefId) {
127        let crate_name = self.tcx.crate_name(def_id.krate);
128        let def_path = self.tcx.def_path(def_id);
129        crate_name.encode(self);
130        def_path.to_string_no_crate_verbose().encode(self);
131    }
132}
133
134impl<'a, 'tcx> TyEncoder<'tcx> for OffloadManifestEncoder<'a, 'tcx> {
135    const CLEAR_CROSS_CRATE: bool = true;
136
137    fn position(&self) -> usize {
138        self.encoder.position()
139    }
140
141    fn type_shorthands(&mut self) -> &mut FxHashMap<Ty<'tcx>, usize> {
142        &mut self.type_shorthands
143    }
144
145    fn predicate_shorthands(&mut self) -> &mut FxHashMap<ty::PredicateKind<'tcx>, usize> {
146        &mut self.predicate_shorthands
147    }
148
149    fn encode_alloc_id(&mut self, _alloc_id: &rustc_middle::mir::interpret::AllocId) {
150        // AllocIds are not expected in the manifest.
151    }
152}
153
154const UNRESOLVED_DEF_ID: DefId = DefId {
155    krate: rustc_span::def_id::CrateNum::MAX,
156    index: rustc_span::def_id::DefIndex::from_u32(0),
157};
158
159/// Decoder used to read the offload monomorphization manifest.
160pub(crate) struct OffloadManifestDecoder<'a, 'tcx> {
161    decoder: MemDecoder<'a>,
162    type_shorthands: Lock<FxHashMap<usize, Ty<'tcx>>>,
163    #[allow(dead_code)]
164    predicate_shorthands: Lock<FxHashMap<usize, ty::PredicateKind<'tcx>>>,
165    tcx: TyCtxt<'tcx>,
166    /// Map from (crate name, DefPath string) to DefId, used to resolve DefIds
167    /// across compilation sessions where StableCrateId differs.
168    def_path_map: Lock<Option<FxHashMap<(Symbol, String), DefId>>>,
169}
170
171impl<'a, 'tcx> OffloadManifestDecoder<'a, 'tcx> {
172    pub(crate) fn new(data: &'a [u8], tcx: TyCtxt<'tcx>) -> Result<Self, ()> {
173        let decoder = MemDecoder::new(data, 0)?;
174        Ok(OffloadManifestDecoder {
175            decoder,
176            type_shorthands: Lock::new(FxHashMap::default()),
177            predicate_shorthands: Lock::new(FxHashMap::default()),
178            tcx,
179            def_path_map: Lock::new(None),
180        })
181    }
182
183    /// (crate name, DefPath) -> DefId map for resolving cross-session DefIds.
184    fn get_or_build_def_path_map(&self) -> FxHashMap<(Symbol, String), DefId> {
185        let mut guard = self.def_path_map.lock();
186        if let Some(map) = guard.as_ref() {
187            return map.clone();
188        }
189        let map = Self::build_def_path_map(self.tcx);
190        *guard = Some(map.clone());
191        map
192    }
193
194    /// Build a (crate name, DefPath) -> DefId map. Owns the format details.
195    fn build_def_path_map(tcx: TyCtxt<'tcx>) -> FxHashMap<(Symbol, String), DefId> {
196        let mut map: FxHashMap<(Symbol, String), DefId> = FxHashMap::default();
197
198        let local_crate_name = tcx.crate_name(LOCAL_CRATE);
199        let krate_items = tcx.hir_crate_items(());
200        let local_def_ids = krate_items
201            .free_items()
202            .map(|id| id.owner_id.to_def_id())
203            .chain(krate_items.trait_items().map(|id| id.owner_id.to_def_id()))
204            .chain(krate_items.impl_items().map(|id| id.owner_id.to_def_id()))
205            .chain(krate_items.foreign_items().map(|id| id.owner_id.to_def_id()));
206        for item_id in local_def_ids {
207            let def_id = item_id;
208            let def_path = tcx.def_path(def_id);
209            map.insert((local_crate_name, def_path.to_string_no_crate_verbose()), def_id);
210        }
211
212        for &cnum in tcx.crates(()) {
213            if cnum == LOCAL_CRATE {
214                continue;
215            }
216            let crate_name = tcx.crate_name(cnum);
217            let num_defs = tcx.num_extern_def_ids(cnum);
218            for i in 0..num_defs {
219                let def_id = DefId { krate: cnum, index: DefIndex::from_usize(i) };
220                let def_path = tcx.def_path(def_id);
221                map.entry((crate_name, def_path.to_string_no_crate_verbose())).or_insert(def_id);
222            }
223        }
224
225        map
226    }
227}
228
229impl<'a, 'tcx> Decoder for OffloadManifestDecoder<'a, 'tcx> {
230    fn read_usize(&mut self) -> usize {
231        self.decoder.read_usize()
232    }
233    fn read_u128(&mut self) -> u128 {
234        self.decoder.read_u128()
235    }
236    fn read_u64(&mut self) -> u64 {
237        self.decoder.read_u64()
238    }
239    fn read_u32(&mut self) -> u32 {
240        self.decoder.read_u32()
241    }
242    fn read_u16(&mut self) -> u16 {
243        self.decoder.read_u16()
244    }
245    fn read_u8(&mut self) -> u8 {
246        self.decoder.read_u8()
247    }
248    fn read_isize(&mut self) -> isize {
249        self.decoder.read_isize()
250    }
251    fn read_i128(&mut self) -> i128 {
252        self.decoder.read_i128()
253    }
254    fn read_i64(&mut self) -> i64 {
255        self.decoder.read_i64()
256    }
257    fn read_i32(&mut self) -> i32 {
258        self.decoder.read_i32()
259    }
260    fn read_i16(&mut self) -> i16 {
261        self.decoder.read_i16()
262    }
263    fn read_i8(&mut self) -> i8 {
264        self.decoder.read_i8()
265    }
266    fn read_raw_bytes(&mut self, len: usize) -> &[u8] {
267        self.decoder.read_raw_bytes(len)
268    }
269    fn peek_byte(&self) -> u8 {
270        self.decoder.peek_byte()
271    }
272    fn position(&self) -> usize {
273        self.decoder.position()
274    }
275}
276
277impl<'a, 'tcx> BlobDecoder for OffloadManifestDecoder<'a, 'tcx> {
278    fn decode_symbol(&mut self) -> rustc_span::Symbol {
279        let s: String = Decodable::decode(self);
280        rustc_span::Symbol::intern(&s)
281    }
282
283    fn decode_byte_symbol(&mut self) -> ByteSymbol {
284        let len = self.read_usize();
285        let bytes = self.read_raw_bytes(len);
286        ByteSymbol::intern(bytes)
287    }
288
289    fn decode_def_index(&mut self) -> rustc_span::def_id::DefIndex {
290        let v = self.read_u32();
291        rustc_span::def_id::DefIndex::from_u32(v)
292    }
293}
294
295impl<'a, 'tcx> SpanDecoder for OffloadManifestDecoder<'a, 'tcx> {
296    fn decode_span(&mut self) -> Span {
297        let lo = self.read_usize();
298        let hi = self.read_usize();
299        let _ctxt = self.read_u32();
300        Span::new(BytePos::from_usize(lo), BytePos::from_usize(hi), SyntaxContext::root(), None)
301    }
302
303    fn decode_expn_id(&mut self) -> rustc_span::ExpnId {
304        let _ = self.read_u32();
305        rustc_span::ExpnId::root()
306    }
307
308    fn decode_syntax_context(&mut self) -> SyntaxContext {
309        let _ = self.read_u32();
310        SyntaxContext::root()
311    }
312
313    fn decode_crate_num(&mut self) -> rustc_span::def_id::CrateNum {
314        let stable_id: StableCrateId = Decodable::decode(self);
315        self.tcx.stable_crate_id_to_crate_num(stable_id)
316    }
317
318    fn decode_def_id(&mut self) -> rustc_span::def_id::DefId {
319        let crate_name: String = Decodable::decode(self);
320        let crate_name = Symbol::intern(&crate_name);
321        let def_path_str: String = Decodable::decode(self);
322        let map = self.get_or_build_def_path_map();
323        map.get(&(crate_name, def_path_str)).copied().unwrap_or(UNRESOLVED_DEF_ID)
324    }
325
326    fn decode_attr_id(&mut self) -> rustc_ast::AttrId {
327        self.tcx.dcx().fatal("AttrIds are not expected in offload manifests");
328    }
329}
330
331impl<'a, 'tcx> TyDecoder<'tcx> for OffloadManifestDecoder<'a, 'tcx> {
332    const CLEAR_CROSS_CRATE: bool = true;
333
334    fn cached_ty_for_shorthand<F>(&mut self, shorthand: usize, or_insert_with: F) -> Ty<'tcx>
335    where
336        F: FnOnce(&mut Self) -> Ty<'tcx>,
337    {
338        if let Some(ty) = self.type_shorthands.lock().get(&shorthand) {
339            return *ty;
340        }
341        let ty = or_insert_with(self);
342        self.type_shorthands.lock().insert(shorthand, ty);
343        ty
344    }
345
346    fn with_position<F, R>(&mut self, pos: usize, f: F) -> R
347    where
348        F: FnOnce(&mut Self) -> R,
349    {
350        let new_decoder = self.decoder.split_at(pos);
351        let old_decoder = std::mem::replace(&mut self.decoder, new_decoder);
352        let result = f(self);
353        self.decoder = old_decoder;
354        result
355    }
356
357    fn decode_alloc_id(&mut self) -> rustc_middle::mir::interpret::AllocId {
358        self.tcx.dcx().fatal("AllocIds are not expected in offload manifests");
359    }
360}
361
362impl<'a, 'tcx> rustc_middle::ty::InternerDecoder for OffloadManifestDecoder<'a, 'tcx> {
363    type Interner = TyCtxt<'tcx>;
364
365    #[inline]
366    fn interner(&self) -> Self::Interner {
367        self.tcx
368    }
369}
370
371/// Write a list of offload kernel instances to the manifest file.
372pub(crate) fn write_manifest<'tcx>(
373    path: &std::path::Path,
374    tcx: TyCtxt<'tcx>,
375    instances: &[ty::Instance<'tcx>],
376) -> std::io::Result<()> {
377    let mut encoder = OffloadManifestEncoder::new(path, tcx)?;
378    instances.encode(&mut encoder);
379    encoder.finish()
380}
381
382/// Write out the offload host-metadata manifest for `mono_items`. No-op unless
383/// the session was invoked with `-Zoffload=HostMetadata=<path>`.
384pub fn write_host_metadata_offload_manifest<'tcx>(tcx: TyCtxt<'tcx>) {
385    let Some(path) = tcx.sess.opts.unstable_opts.offload.iter().find_map(|o| {
386        if let rustc_session::config::Offload::HostMetadata(p) = o { Some(p) } else { None }
387    }) else {
388        ::rustc_middle::util::bug::bug_fmt(format_args!("HostMetadata path not found; caller should have checked"));bug!("HostMetadata path not found; caller should have checked");
389    };
390
391    let partitions = tcx.collect_and_partition_mono_items(());
392    let mono_items: Vec<MonoItem<'_>> = partitions
393        .codegen_units
394        .iter()
395        .flat_map(|cgu| cgu.items().iter())
396        .map(|(item, _)| *item)
397        .collect();
398
399    let instances: Vec<ty::Instance<'tcx>> = mono_items
400        .iter()
401        .filter_map(|item| {
402            if let MonoItem::Fn(instance) = item {
403                if tcx
404                    .codegen_fn_attrs(instance.def_id())
405                    .flags
406                    .contains(CodegenFnAttrFlags::OFFLOAD_KERNEL)
407                {
408                    Some(*instance)
409                } else {
410                    None
411                }
412            } else {
413                None
414            }
415        })
416        .collect();
417
418    if let Err(e) = write_manifest(std::path::Path::new(path), tcx, &instances) {
419        tcx.dcx().emit_fatal(crate::diagnostics::OffloadManifestWriteError {
420            path: path.clone(),
421            err: e.to_string(),
422        });
423    }
424}
425
426/// Read a list of offload kernel instances from the manifest file.
427pub(crate) fn read_manifest<'tcx>(
428    path: &std::path::Path,
429    tcx: TyCtxt<'tcx>,
430) -> std::io::Result<Vec<ty::Instance<'tcx>>> {
431    let data = fs::read(path)?;
432    let mut decoder = OffloadManifestDecoder::new(&data, tcx)
433        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid manifest"))?;
434
435    let instances: Vec<ty::Instance<'tcx>> = Decodable::decode(&mut decoder);
436
437    Ok(instances)
438}
rustc_monomorphize/offload/manifest.rs

rustc_monomorphize/offload/
manifest.rs
