Update SGX example (#1825)

.gitignore

@@ -198,3 +198,11 @@ tvm_t.*
 
 # tmp file
 .nfs*
+
+# keys
+*.pem
+*.p12
+*.pfx
+*.cer
+*.crt
+*.der

apps/sgx/Makefile

-# 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 <tvm/runtime/c_runtime_api.h>"}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

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`.  

apps/sgx/enclave/.rustfmt.toml

+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

apps/sgx/enclave/Cargo.toml

+[package]
+name = "model-enclave"
+version = "0.1.0"
+authors = ["Nick Hynes <nhynes@berkeley.edu>"]
+
+[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

apps/sgx/enclave/Makefile

+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

apps/sgx/enclave/Xargo.toml

+[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

apps/sgx/enclave/enclave.lds

+enclave.so
+{
+    global:
+        g_global_data_sim;
+        g_global_data;
+        enclave_entry;
+    local:
+        *;
+};

apps/sgx/enclave_config.xml → apps/sgx/enclave/enclave_config.xml.in

 <EnclaveConfiguration>
   <ProdID>0</ProdID>
   <ISVSVN>0</ISVSVN>
-  <StackMaxSize>0x2000</StackMaxSize>
-  <HeapMaxSize>0x2000</HeapMaxSize>
-  <TCSNum>5</TCSNum>
-  <TCSPolicy>1</TCSPolicy>
+  <StackMaxSize>0x100000</StackMaxSize>
+  <HeapMaxSize>0xf0000000</HeapMaxSize>
+  <TCSNum>NUM_THREADS</TCSNum>
+  <TCSPolicy>0</TCSPolicy> <!-- must be "bound" to use thread_local -->
   <DisableDebug>0</DisableDebug>
   <MiscSelect>0</MiscSelect>
   <MiscMask>0xFFFFFFFF</MiscMask>

apps/sgx/enclave/src/lib.rs

+#![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<GraphExecutor<'static, 'static>> = {
+    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<PackedFunc> = {
+      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::<Vec<TVMArgValue>>()
+    };
+    let (rv, tc) = ECALL_FUNCS[func_id as usize](&args).into_tvm_value();
+    unsafe {
+      *ret_val = rv;
+      *ret_type_code = tc;
+    }
+  }
+}

apps/sgx/prepare_test_libs.py

-"""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()

apps/sgx/run_example.sh

 #!/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

apps/sgx/run_model.py

+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()

apps/sgx/test_addone.py

-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!")

cmake/modules/SGX.cmake

 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