tvm-graph-rt 0.1.0-alpha

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

use std::{
    collections::HashMap, convert::AsRef, ffi::CStr, os::raw::c_char, string::String, sync::RwLock,
};

use lazy_static::lazy_static;

use tvm_sys::{ffi::BackendPackedCFunc, packed_func::PackedFunc};

use super::Module;

pub struct SystemLibModule;

#[cfg(target_env = "sgx")]
extern "C" {
    fn __tvm_module_startup();
}

lazy_static! {
    static ref SYSTEM_LIB_FUNCTIONS: RwLock<HashMap<String, &'static (dyn PackedFunc)>> =
        RwLock::new(HashMap::new());
}

impl Module for SystemLibModule {
    fn get_function<S: AsRef<str>>(&self, name: S) -> Option<&(dyn PackedFunc)> {
        SYSTEM_LIB_FUNCTIONS
            .read()
            .unwrap()
            .get(name.as_ref())
            .copied()
    }
}

impl Default for SystemLibModule {
    fn default() -> Self {
        #[cfg(target_env = "sgx")]
        unsafe {
            __tvm_module_startup();
        }
        SystemLibModule {}
    }
}

#[no_mangle]
pub extern "C" fn TVMBackendRegisterSystemLibSymbol(
    cname: *const c_char,
    func: BackendPackedCFunc,
) -> i32 {
    let name = unsafe { CStr::from_ptr(cname).to_str().unwrap() };
    SYSTEM_LIB_FUNCTIONS.write().unwrap().insert(
        name.to_string(),
        &*Box::leak(super::wrap_backend_packed_func(name.to_string(), func)),
    );
    0
}