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

rustc_monomorphize/offload/
manifest.rs
