Update SGX example

.gitignore

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

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 $@ [email protected]
-	awk 'NR==4{print "#include <tvm/runtime/c_runtime_api.h>"}1' [email protected] > $@
+$(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

@@ -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

apps/sgx/enclave/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "model-enclave"
+version = "0.1.0"
+authors = ["Nick Hynes <[email protected]>"]
+
+[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

@@ -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

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

apps/sgx/enclave/enclave.lds

@@ -0,0 +1,9 @@
+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

@@ -1,10 +1,10 @@
 <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>