diff --git a/.gitignore b/.gitignore index 833eee1a0774..f4dd65e256a3 100644 --- a/.gitignore +++ b/.gitignore @@ -201,3 +201,11 @@ build* # tmp file .nfs* + +# keys +*.pem +*.p12 +*.pfx +*.cer +*.crt +*.der diff --git a/apps/sgx/Makefile b/apps/sgx/Makefile index 1038f57c3ba1..875897b82d23 100644 --- a/apps/sgx/Makefile +++ b/apps/sgx/Makefile @@ -1,13 +1,12 @@ -# Makefile for example to deploy TVM modules in SGX. - -TVM_ROOT := $(shell cd ../..; pwd) -NNVM_PATH := nnvm -DMLC_CORE := ${TVM_ROOT}/dmlc-core - SGX_SDK ?= /opt/sgxsdk +RUST_SGX_SDK ?= /opt/rust-sgx-sdk SGX_MODE ?= SIM -SGX_ARCH ?= x64 -SGX_DEBUG ?= 1 +DEBUG ?= true +NUM_THREADS ?= 4 + +TVM_DIR ?= ../.. + +export sgx_edger8r := $(SGX_SDK)/bin/x64/sgx_edger8r sgx_enclave_signer := $(SGX_SDK)/bin/x64/sgx_sign @@ -20,69 +19,71 @@ trts_library_name := sgx_trts$(sgx_sim) tservice_library_name := sgx_tservice$(sgx_sim) uservice_library_name := sgx_uae_service$(sgx_sim) -pkg_cflags := -std=c++11 -O2 -fPIC\ - -I${TVM_ROOT}/include\ - -I${DMLC_CORE}/include\ - -I${TVM_ROOT}/3rdparty/dlpack/include\ - -I.\ - -DDMLC_LOG_STACK_TRACE=0\ - -fmax-errors=4 - -pkg_ldflags := -L${TVM_ROOT}/lib - -enclave_include_paths := -I$(SGX_SDK)/include\ - -I$(SGX_SDK)/include/tlibc\ - -I$(SGX_SDK)/include/libcxx\ - -I$(SGX_SDK)/include/stdc++\ +pkg_cflags := -std=c++11 -fPIC \ + -I$(SGX_SDK)/include \ + -I$(TVM_DIR)/include \ + -I$(TVM_DIR)/dlpack/include \ + -I$(TVM_DIR)/dmlc-core/include + +pkg_ldflags := -L$(TVM_DIR)/build -ltvm_runtime + +ifneq ($(DEBUG), false) + debug := debug + enclave_cflags += -Og -g + pkg_cflags += -Og -g +else + debug := release + enclave_cflags += -O2 + pkg_cflags += -O2 +endif -enclave_cflags := -static -nostdinc\ - -fvisibility=hidden -fpie -fstack-protector-strong\ - -ffunction-sections -fdata-sections\ - -DDMLC_CXX11_THREAD_LOCAL=0\ - -include "lib/tvm_t.h"\ - $(enclave_include_paths)\ +build_dir := build -enclave_cxxflags := -nostdinc++ $(enclave_cflags) -DTVM_SGX_MAX_CONCURRENCY=4 +enclave_cflags := \ + -I$(SGX_SDK)/include \ + -I$(SGX_SDK)/include/tlibc \ + -I$(SGX_SDK)/include/stdport \ + -I$(SGX_SDK)/include/epid \ + -I$(TVM_DIR)/include \ + -I$(TVM_DIR)/dlpack/include \ + -I$(TVM_DIR)/dmlc-core/include enclave_ldflags :=\ + -L$(build_dir) -L$(TVM_DIR)/build \ -Wl,--no-undefined -nostdlib -nodefaultlibs -nostartfiles -L$(SGX_SDK)/lib64\ -Wl,--whole-archive -l$(trts_library_name) -Wl,--no-whole-archive\ -Wl,--start-group\ -lsgx_tstdc -lsgx_tstdcxx -lsgx_tcxx -lsgx_tcrypto -lsgx_tkey_exchange -l$(tservice_library_name)\ + -lenclave -ltvm_t\ -Wl,--end-group\ -Wl,-Bstatic -Wl,-Bsymbolic -Wl,--no-undefined\ -Wl,-pie,-eenclave_entry -Wl,--export-dynamic\ - -Wl,--defsym,__ImageBase=0 -Wl,--gc-sections - -.PHONY: clean all + -Wl,--defsym,__ImageBase=0 -Wl,--gc-sections\ + -Wl,--version-script=enclave/enclave.lds -all: lib/test_addone.signed.so +.PHONY: enclave clean -# The code library built by TVM -lib/test_addone_sys.o: prepare_test_libs.py - python prepare_test_libs.py +enclave: $(build_dir)/enclave.signed.so -lib/tvm_t.h: ../../src/runtime/sgx/tvm.edl - $(sgx_edger8r) --trusted $< --trusted-dir lib --search-path $(SGX_SDK)/include - mv $@ $@.in - awk 'NR==4{print "#include "}1' $@.in > $@ +$(build_dir)/enclave.signed.so: $(build_dir)/enclave.so build/enclave_config.xml enclave/enclave.pem + $(sgx_enclave_signer) sign -key enclave/enclave.pem -enclave $< -out $@ -config build/enclave_config.xml -lib/tvm_t.c: lib/tvm_t.h +enclave/enclave.pem: + curl -sSo $@ 'https://gist.githubusercontent.com/nhynes/8a2d80068a92e672f8b0b7d710ceb404/raw/2d5ae5fbe83198ede49465fdc6535065e093543b/tvm_sgx_demo.pem' -lib/tvm_t.o: lib/tvm_t.c - $(CC) $(enclave_cflags) $(pkg_cflags) -c $< -o $@ -include $(TVM_ROOT)/include/tvm/runtime/c_runtime_api.h +build/enclave_config.xml: enclave/enclave_config.xml.in + cpp $^ -P -o $@ -DNUM_THREADS=$$(( $(NUM_THREADS) + 1 )) -# The enclave library -lib/test_addone.so: $(TVM_ROOT)/src/runtime/sgx/trusted/runtime.cc lib/tvm_t.o lib/test_addone_sys.o - $(CXX) $^ -o $@ $(pkg_cflags) $(pkg_ldflags) $(enclave_cxxflags) $(enclave_ldflags) -g +$(build_dir)/enclave.so: $(build_dir)/libenclave.a $(TVM_DIR)/build/libtvm_t.a + $(CXX) $< -o $@ $(enclave_ldflags) $(enclave_cflags) -ltvm_t -# The demo enclave signing key -lib/enclave.pem: - curl -Lso $@ https://gist.githubusercontent.com/nhynes/8a2d80068a92e672f8b0b7d710ceb404/raw/2d5ae5fbe83198ede49465fdc6535065e093543b/tvm_sgx_demo.pem +$(build_dir)/libenclave.a: enclave/target/x86_64-unknown-linux-sgx/$(debug)/libmodel_enclave.a + @mkdir -p $(@D) + @cp $< $@ -# The signed enclave -lib/test_addone.signed.so: lib/test_addone.so enclave_config.xml lib/enclave.pem - $(sgx_enclave_signer) sign -key lib/enclave.pem -enclave $< -out $@ -config enclave_config.xml +enclave/target/x86_64-unknown-linux-sgx/$(debug)/libmodel_enclave.a: enclave/**/* + $(MAKE) -C enclave clean: - rm -rf lib + $(MAKE) -s -C enclave clean + rm -rf build diff --git a/apps/sgx/README.md b/apps/sgx/README.md index 565519d457ce..dd21cff02f80 100644 --- a/apps/sgx/README.md +++ b/apps/sgx/README.md @@ -4,13 +4,22 @@ This application demonstrates the use of a simple TVM model in the [Intel SGX](h ## Prerequisites +1. The TVM premade Docker image + +or + 1. A GNU/Linux environment 2. TVM compiled with LLVM and SGX; and the `tvm` Python module 3. The [Linux SGX SDK](https://github.com/intel/linux-sgx) [link to pre-built libraries](https://01.org/intel-software-guard-extensions/downloads) +4. [Rust](https://rustup.sh) +5. The [rust-sgx-sdk](https://github.com/baidu/rust-sgx-sdk) +6. [xargo](https://github.com/japaric/xargo) + +Check out the `/tvm/install/ubuntu_install_sgx.sh` for the commands to get these dependencies. ## Running the example -`SGX_SDK=/path/to/sgxsdk bash run_example.sh` +`bash run_example.sh` If everything goes well, you should see a lot of build messages and below them the text `It works!`. @@ -24,10 +33,9 @@ In this library, one can use other libraries like TVM. Building this example performs the following steps: 1. Creates a simple TVM module that computes `x + 1` and save it as a system library. -2. Builds a minimal TVM runtime pack that can load the module. -3. Links the TVM module into an SGX enclave along with some code that runs the module. -4. Compiles and runs an executable that loads the enclave and calls a function - which invokes the TVM module. +2. Builds a TVM runtime that links the module and allows running it using the TVM Python runtime. +3. Packages the bundle into an SGX enclave +4. Runs the enclave using the usual TVM Python `module` API For more information on building, please refer to the `Makefile`. For more information on the TVM module, please refer to `../howto_deploy`. diff --git a/apps/sgx/enclave/.rustfmt.toml b/apps/sgx/enclave/.rustfmt.toml new file mode 100644 index 000000000000..9ae87cc6bfcf --- /dev/null +++ b/apps/sgx/enclave/.rustfmt.toml @@ -0,0 +1,59 @@ +max_width = 100 +hard_tabs = false +tab_spaces = 2 +newline_style = "Auto" +use_small_heuristics = "Default" +indent_style = "Block" +wrap_comments = false +comment_width = 80 +normalize_comments = false +format_strings = false +format_macro_matchers = false +format_macro_bodies = true +empty_item_single_line = true +struct_lit_single_line = true +fn_single_line = false +where_single_line = false +imports_indent = "Block" +imports_layout = "Mixed" +merge_imports = true +reorder_imports = true +reorder_modules = true +reorder_impl_items = false +type_punctuation_density = "Wide" +space_before_colon = false +space_after_colon = true +spaces_around_ranges = false +binop_separator = "Front" +remove_nested_parens = true +combine_control_expr = true +struct_field_align_threshold = 0 +match_arm_blocks = true +force_multiline_blocks = false +fn_args_density = "Tall" +brace_style = "SameLineWhere" +control_brace_style = "AlwaysSameLine" +trailing_semicolon = true +trailing_comma = "Vertical" +match_block_trailing_comma = true +blank_lines_upper_bound = 1 +blank_lines_lower_bound = 0 +edition = "2015" +merge_derives = true +use_try_shorthand = true +use_field_init_shorthand = false +force_explicit_abi = true +condense_wildcard_suffixes = true +color = "Auto" +required_version = "0.99.5" +unstable_features = false +disable_all_formatting = false +skip_children = false +hide_parse_errors = false +error_on_line_overflow = false +error_on_unformatted = false +report_todo = "Never" +report_fixme = "Never" +ignore = [] +emit_mode = "Files" +make_backup = false diff --git a/apps/sgx/enclave/Cargo.toml b/apps/sgx/enclave/Cargo.toml new file mode 100644 index 000000000000..9a14c76c5897 --- /dev/null +++ b/apps/sgx/enclave/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "model-enclave" +version = "0.1.0" +authors = ["Nick Hynes "] + +[lib] +crate-type = ["staticlib"] + +[dependencies] +lazy_static = "1.1.0" +# tvm = { path = "../../../rust", default-features = false, features = ["sgx"] } +tvm = { path = "/home/nhynes/myelin/deps/tvm-rs", default-features = false, features = ["sgx"] } + +[profile.release] +lto = true +opt-level = 3 diff --git a/apps/sgx/enclave/Makefile b/apps/sgx/enclave/Makefile new file mode 100644 index 000000000000..e8515356238a --- /dev/null +++ b/apps/sgx/enclave/Makefile @@ -0,0 +1,35 @@ +MODEL ?= resnet +NUM_THREADS ?= 4 +BATCH_SIZE ?= 64 +TRAINING ?= true +DEBUG ?= false + +build_dir := ../build + +ifeq ($(DEBUG), false) + debug := release + xargo_args := --release +else + debug := debug +endif + +target/x86_64-unknown-linux-sgx/$(debug)/libmodel-enclave.a: $(build_dir)/libmodel.a **/* + RUST_TARGET_PATH=$(shell pwd) \ + RUST_TARGET_DIR=$(shell pwd)/target \ + RUSTFLAGS="-Z force-unstable-if-unmarked" \ + TVM_NUM_THREADS=$(NUM_THREADS) \ + BUILD_DIR=../build \ + xargo build --target x86_64-unknown-linux-sgx $(xargo_args) -q + +$(build_dir)/libmodel.a: $(build_dir)/model.o + llvm-ar cr $@ $^ + +$(build_dir)/model.o: $(build_dir)/model.bc + clang -c $< -o $@ -fPIC -O3 + objcopy --globalize-symbol __tvm_module_startup $@ + +$(build_dir)/model.bc: src/build_model.py + python3 $< -o $(build_dir) + +clean: + xargo clean diff --git a/apps/sgx/enclave/Xargo.toml b/apps/sgx/enclave/Xargo.toml new file mode 100644 index 000000000000..1fd50d699264 --- /dev/null +++ b/apps/sgx/enclave/Xargo.toml @@ -0,0 +1,17 @@ +[dependencies] +alloc = {} +panic_unwind = {} +panic_abort = {} + +[dependencies.std] +features = ["backtrace", "stdio", "untrusted_time"] +path = "/home/nhynes/myelin/deps/rust-sgx-sdk/xargo/sgx_tstd" +# git = "https://github.com/oasislabs/rust-sgx-sdk" +# rev = "7334c30d85cb1752577998705110b7b27c69b570" +stage = 2 + +[dependencies.xargo_sgx_rand] +# git = "https://github.com/oasislabs/rust-sgx-sdk" +path = "/home/nhynes/myelin/deps/rust-sgx-sdk/xargo/sgx_rand" +# rev = "7334c30d85cb1752577998705110b7b27c69b570" +stage = 3 diff --git a/apps/sgx/enclave/enclave.lds b/apps/sgx/enclave/enclave.lds new file mode 100644 index 000000000000..e3d9d0ee0d90 --- /dev/null +++ b/apps/sgx/enclave/enclave.lds @@ -0,0 +1,9 @@ +enclave.so +{ + global: + g_global_data_sim; + g_global_data; + enclave_entry; + local: + *; +}; diff --git a/apps/sgx/enclave_config.xml b/apps/sgx/enclave/enclave_config.xml.in similarity index 50% rename from apps/sgx/enclave_config.xml rename to apps/sgx/enclave/enclave_config.xml.in index 07be0d7a7ad2..d49b6693f231 100644 --- a/apps/sgx/enclave_config.xml +++ b/apps/sgx/enclave/enclave_config.xml.in @@ -1,10 +1,10 @@ 0 0 - 0x2000 - 0x2000 - 5 - 1 + 0x100000 + 0xf0000000 + NUM_THREADS + 0 0 0 0xFFFFFFFF diff --git a/apps/sgx/enclave/src/lib.rs b/apps/sgx/enclave/src/lib.rs new file mode 100644 index 000000000000..d74015a92510 --- /dev/null +++ b/apps/sgx/enclave/src/lib.rs @@ -0,0 +1,119 @@ +#![feature(try_from)] + +#[macro_use] +extern crate lazy_static; +extern crate tvm; + +use std::{convert::TryFrom, sync::Mutex}; + +use tvm::runtime::{sgx, Graph, GraphExecutor, SystemLibModule, TVMArgValue, TVMRetValue}; + +lazy_static! { + static ref SYSLIB: SystemLibModule = { SystemLibModule::default() }; + static ref MODEL: Mutex> = { + let _params = include_bytes!(concat!("../", env!("BUILD_DIR"), "/params.bin")); + let graph_json = include_str!(concat!("../", env!("BUILD_DIR"), "/graph.json")); + + let graph = Graph::try_from(graph_json).unwrap(); + Mutex::new(GraphExecutor::new(graph, &*SYSLIB).unwrap()) + }; +} + +fn ecall_init(_args: &[TVMArgValue]) -> TVMRetValue { + lazy_static::initialize(&MODEL); + TVMRetValue::from(0) +} + +fn ecall_main(_args: &[TVMArgValue]) -> TVMRetValue { + let model = MODEL.lock().unwrap(); + // model.set_input("data", args[0]); + model.run(); + sgx::shutdown(); + // model.get_output(0).into() + TVMRetValue::from(42) +} + +pub mod ecalls { + //! todo: generate this using proc_macros + + use super::*; + + use std::{ + ffi::CString, + os::raw::{c_char, c_int}, + slice, + }; + + use tvm::{ + ffi::runtime::{TVMRetValueHandle, TVMValue}, + runtime::{ + sgx::{run_worker, SgxStatus}, + PackedFunc, + }, + }; + + macro_rules! tvm_ocall { + ($func: expr) => { + match $func { + 0 => Ok(()), + err => Err(err), + } + }; + } + + const ECALLS: &'static [&'static str] = &["__tvm_run_worker__", "__tvm_main__", "init"]; + + lazy_static! { + static ref ECALL_FUNCS: Vec = { + vec![ + Box::new(run_worker), + Box::new(ecall_main), + Box::new(ecall_init), + ] + }; + } + + extern "C" { + fn __tvm_module_startup() -> (); + fn tvm_ocall_register_export(name: *const c_char, func_id: c_int) -> SgxStatus; + } + + #[no_mangle] + pub extern "C" fn tvm_ecall_init(_ret: TVMRetValueHandle) { + unsafe { + __tvm_module_startup(); + + ECALLS.into_iter().enumerate().for_each(|(i, ecall)| { + tvm_ocall!(tvm_ocall_register_export( + CString::new(*ecall).unwrap().as_ptr(), + i as i32 + )).expect(&format!("Error registering `{}`", ecall)); + }); + } + } + + #[no_mangle] + pub extern "C" fn tvm_ecall_packed_func( + func_id: c_int, + arg_values: *const TVMValue, + type_codes: *const c_int, + num_args: c_int, + ret_val: *mut TVMValue, + ret_type_code: *mut i64, + ) { + let args = unsafe { + let values = slice::from_raw_parts(arg_values, num_args as usize); + let type_codes = slice::from_raw_parts(type_codes, num_args as usize); + values + .into_iter() + .zip(type_codes.into_iter()) + .map(|(v, t)| TVMArgValue::new(*v, *t as i64)) + .collect::>() + }; + let (rv, tc) = ECALL_FUNCS[func_id as usize](&args).into_tvm_value(); + unsafe { + *ret_val = rv; + *ret_type_code = tc; + } + } +} diff --git a/apps/sgx/prepare_test_libs.py b/apps/sgx/prepare_test_libs.py deleted file mode 100644 index f676f46b7ff0..000000000000 --- a/apps/sgx/prepare_test_libs.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Script to prepare test_addone_sys.o""" - -from os import path as osp - -import tvm - -CWD = osp.dirname(osp.abspath(osp.expanduser(__file__))) - - -def main(): - out_dir = osp.join(CWD, 'lib') - - n = tvm.var('n') - A = tvm.placeholder((n,), name='A') - B = tvm.compute(A.shape, lambda *i: A(*i) + 1, name='B') - s = tvm.create_schedule(B.op) - s[B].parallel(s[B].op.axis[0]) - print(tvm.lower(s, [A, B], simple_mode=True)) - - # Compile library in system library mode - fadd_syslib = tvm.build(s, [A, B], 'llvm --system-lib') - fadd_syslib.save(osp.join(out_dir, 'test_addone_sys.o')) - - -if __name__ == '__main__': - main() diff --git a/apps/sgx/run_example.sh b/apps/sgx/run_example.sh index 9334b260cbf3..cc6f22f24e00 100755 --- a/apps/sgx/run_example.sh +++ b/apps/sgx/run_example.sh @@ -1,6 +1,6 @@ #!/bin/bash sgx_sdk=${SGX_SDK:=/opt/sgxsdk} -make -echo "=========================" -LD_LIBRARY_PATH="$sgx_sdk/lib64":${LD_LIBRARY_PATH} TVM_CACHE_DIR=/tmp python test_addone.py +LD_LIBRARY_PATH="$sgx_sdk/lib64":${LD_LIBRARY_PATH} make +printf "\n" +LD_LIBRARY_PATH="$sgx_sdk/lib64":${LD_LIBRARY_PATH} TVM_CACHE_DIR=/tmp python3 run_model.py diff --git a/apps/sgx/run_model.py b/apps/sgx/run_model.py new file mode 100644 index 000000000000..491a5ccbda3c --- /dev/null +++ b/apps/sgx/run_model.py @@ -0,0 +1,20 @@ +import os.path as osp +import numpy as np +import tvm + +CWD = osp.abspath(osp.dirname(__file__)) + + +def main(): + ctx = tvm.context('cpu', 0) + model = tvm.module.load(osp.join(CWD, 'build', 'enclave.signed.so')) + out = model() + if out == 42: + print('It works!') + else: + print('It doesn\'t work!') + exit(1) + + +if __name__ == '__main__': + main() diff --git a/apps/sgx/test_addone.py b/apps/sgx/test_addone.py deleted file mode 100644 index 5ddccfa425cc..000000000000 --- a/apps/sgx/test_addone.py +++ /dev/null @@ -1,13 +0,0 @@ -import tvm -import numpy as np - -ctx = tvm.context('cpu', 0) -fadd1 = tvm.module.load('lib/test_addone.signed.so') - -n = 10 -x = tvm.nd.array(np.random.uniform(size=n).astype('float32'), ctx) -y = tvm.nd.array(np.zeros(n, dtype='float32'), ctx) -fadd1(x, y) - -np.testing.assert_allclose(y.asnumpy(), x.asnumpy() + 1) -print("It works!") diff --git a/cmake/modules/SGX.cmake b/cmake/modules/SGX.cmake index c9894de11f8b..608d6ff5a4bd 100644 --- a/cmake/modules/SGX.cmake +++ b/cmake/modules/SGX.cmake @@ -1,5 +1,4 @@ if(NOT USE_SGX STREQUAL "OFF") - message(STATUS "Build with SGX support") set(_sgx_src ${CMAKE_CURRENT_SOURCE_DIR}/src/runtime/sgx) set(_tvm_u_h ${_sgx_src}/untrusted/tvm_u.h) @@ -9,8 +8,11 @@ if(NOT USE_SGX STREQUAL "OFF") set(_sgx_ustdc ${RUST_SGX_SDK}/sgx_ustdc) set(_urts_lib "sgx_urts") - if(SGX_MODE STREQUAL "SIM") + if(NOT SGX_MODE STREQUAL "HW") + message(STATUS "Build with SGX support (SIM)") set(_urts_lib "${_urts_lib}_sim") + else() + message(STATUS "Build with SGX support (HW)") endif() # build edge routines