1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

#![cfg(target_arch = "wasm32")]
#![no_std]
// Note: The WebAssembly bindings need to provide alloc error handling.
#![feature(core_intrinsics, lang_items, alloc_error_handler)]

extern crate alloc;

pub mod allocator;
mod buf_writer;
mod capability;
mod guards;
pub mod intrinsics;
mod logging;
mod model;
mod resources;
pub mod serial;
mod stats_allocator;
pub mod tensor_output;

use core::{alloc::Layout, fmt::Write, panic::PanicInfo};

use dlmalloc::GlobalDlmalloc;

use crate::allocator::Allocator;
pub use crate::{
    buf_writer::BufWriter,
    capability::Capability,
    guards::{PipelineGuard, SetupGuard},
    logging::Logger,
    model::Model,
    resources::{Resource, ResourceError},
    serial::Serial,
    tensor_output::TensorOutput,
};

#[global_allocator]
pub static ALLOCATOR: Allocator<GlobalDlmalloc> =
    Allocator::new(GlobalDlmalloc);

#[panic_handler]
fn on_panic(info: &PanicInfo) -> ! {
    static mut PANICKING: bool = false;

    unsafe {
        // We need to guard against the possiblity that logging a panic may
        // in turn trigger a panic (e.g. due to OOM), causing infinite
        // recursion.
        if !PANICKING {
            PANICKING = true;

            // First we try to log the panic at the ERROR level. This should
            // be translated into a runtime trap, so under most circumstances
            // the log call won't return and our user will get a nice error
            // message.
            log::error!("{}", info);
        }

        // However, some times the runtime won't receive the log message (e.g.
        // log level filtering or because an OOM in logging recursively
        // triggered the panic handler). If that is the case, we still try to
        // send *some* message to the runtime so they know the world is broken.

        // Safety: We need our own buffer for panic messages in case the
        // allocator is fubar. Runes are single-threaded, so we can
        // guarantee we'll never have aliased mutation.
        static mut DEBUG_BUFFER: [u8; 1024] = [0; 1024];
        let mut w = BufWriter::new(&mut DEBUG_BUFFER);

        if write!(w, "{}", info).is_ok() {
            let written = w.written();
            intrinsics::_debug(written.as_ptr(), written.len() as u32);
        }

        // And now we've done everything we can, we ungracefully crash.
        core::arch::wasm32::unreachable()
    }
}

#[alloc_error_handler]
fn on_alloc_error(layout: Layout) -> ! {
    panic!(
        "memory allocation of {} bytes failed ({:?})",
        layout.size(),
        ALLOCATOR.stats()
    );
}