TVM SGX (#919)

* Update DMLC core most recent version

* Modify runtime to minimize io when building for SGX

* Add SGX example app

* Prefer streaming versions of packed_func function

apps/sgx/.gitignore

+lib/
+bin/

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
+SGX_MODE ?= SIM
+SGX_ARCH ?= x64
+SGX_DEBUG ?= 1
+
+sgx_edger8r := $(SGX_SDK)/bin/x64/sgx_edger8r
+sgx_enclave_signer := $(SGX_SDK)/bin/x64/sgx_sign
+
+ifneq ($(SGX_MODE), HW)
+	sgx_sim := _sim
+endif
+urts_library_name := sgx_urts$(sgx_sim)
+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}/dlpack/include\
+	-I.\
+	-DDMLC_LOG_STACK_TRACE=0\
+
+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++\
+
+enclave_cflags := -static -nostdinc\
+	-fvisibility=hidden -fpie -fstack-protector-strong\
+	-ffunction-sections -fdata-sections\
+	-DDMLC_CXX11_THREAD_LOCAL=0\
+	$(enclave_include_paths)\
+
+enclave_cxxflags := -nostdinc++ $(enclave_cflags)
+
+enclave_ldflags :=\
+	-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)\
+	-Wl,--end-group\
+	-Wl,-Bstatic -Wl,-Bsymbolic -Wl,--no-undefined\
+	-Wl,-pie,-eenclave_entry -Wl,--export-dynamic\
+	-Wl,--defsym,__ImageBase=0 -Wl,--gc-sections
+
+app_cflags := -I$(SGX_SDK)/include -Ilib
+
+app_ldflags := -L$(SGX_SDK)/lib64\
+	-l$(urts_library_name) -l$(uservice_library_name) -lpthread\
+
+.PHONY: clean all
+
+all: lib/test_addone.signed.so bin/test_addone
+
+# Build rule for all-in-one TVM package library
+lib/tvm_runtime_pack.o: tvm_runtime_pack.cc
+	@mkdir -p $(@D)
+	$(CXX) -c $< -o $@ $(pkg_cflags) $(pkg_ldflags) $(enclave_cxxflags) -g
+
+# The code library built by TVM
+lib/test_addone_sys.o: prepare_test_libs.py
+	python prepare_test_libs.py
+
+# EDL files
+lib/test_addone_u.c: $(sgx_edger8r) test_addone.edl
+	$(sgx_edger8r) --untrusted test_addone.edl --untrusted-dir lib --search-path $(SGX_SDK)/include
+
+lib/test_addone_u.o: lib/test_addone_u.c
+	$(CC) $(enclave_cflags) -c $< -o $@
+
+lib/test_addone_t.c: test_addone.edl
+	$(sgx_edger8r) --trusted $< --trusted-dir lib --search-path $(SGX_SDK)/include
+
+lib/test_addone_t.o: lib/test_addone_t.c
+	$(CC) $(enclave_cflags) -c $< -o $@
+
+# The enclave library
+lib/test_addone.so: enclave.cc lib/tvm_runtime_pack.o lib/test_addone_t.o lib/test_addone_sys.o
+	$(CXX) $^ -o $@ $(pkg_cflags) $(pkg_ldflags) $(enclave_cxxflags) $(enclave_ldflags) -g
+
+# The signed enclave
+lib/test_addone.signed.so: lib/test_addone.so enclave_config.xml
+	$(sgx_enclave_signer) sign -key enclave_private.pem -enclave $< -out $@ -config enclave_config.xml
+
+# An app that runs the enclave
+bin/test_addone: app.cc lib/test_addone_u.o
+	@mkdir -p $(@D)
+	$(CXX) $^ -o $@ $(app_cflags) $(app_ldflags)
+
+# Debugging runtime pack built without SGX (c.f. howto_deploy/tvm_runtime_pack.cc)
+lib/tvm_runtime_pack_nosgx.o: tvm_runtime_pack.cc
+	@mkdir -p $(@D)
+	$(CXX) -c $< -o $@ $(pkg_cflags) $(pkg_ldflags) -g
+
+# Debugging binary that runs TVM without SGX
+bin/addone_nosgx: enclave.cc lib/tvm_runtime_pack_nosgx.o lib/test_addone_sys.o
+	@mkdir -p $(@D)
+	$(CXX) $^ -o $@ $(pkg_cflags) $(pkg_ldflags) -g
+
+clean:
+	rm -rf lib bin

apps/sgx/README.md

+# TVM in Intel SGX Example
+
+This application demonstrates the use of a simple TVM model in the [Intel SGX](https://software.intel.com/en-us/blogs/2013/09/26/protecting-application-secrets-with-intel-sgx) trusted computing environment.
+
+## Prerequisites
+
+1. A GNU/Linux environment
+2. TVM compiled with LLVM 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)
+
+## Running the example
+
+`SGX_SDK=/path/to/sgxsdk bash run_example.sh`
+
+If everything goes well, you should see a lot of build messages and below them
+the text `It works!`.
+
+## High-level overview
+
+First of all, it helps to think of an SGX enclave as a library that can be called
+to perform trusted computation.
+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.
+
+For more information on building, please refer to the `Makefile`.  
+For more information on the TVM module, please refer to `../howto_deploy`.  
+For more in formation on SGX enclaves, please refer to the [SGX Enclave Demo](https://github.com/intel/linux-sgx/tree/master/SampleCode/SampleEnclave/)

apps/sgx/app.cc

+#include <cstdio>
+
+#include "sgx_urts.h"
+#include "sgx_eid.h"
+#include "test_addone_u.h"
+
+#define TOKEN_FILENAME   "bin/test_addone.token"
+#define ENCLAVE_FILENAME "lib/test_addone.signed.so"
+
+sgx_enclave_id_t global_eid = 0;  // global EID shared by multiple threads
+
+typedef struct _sgx_errlist_t {
+  sgx_status_t err;
+  const char *msg;
+} sgx_errlist_t;
+
+/* Error code returned by sgx_create_enclave */
+static sgx_errlist_t sgx_errlist[] = {
+  { SGX_ERROR_DEVICE_BUSY, "SGX device was busy." },
+  { SGX_ERROR_ENCLAVE_FILE_ACCESS, "Can't open enclave file." },
+  { SGX_ERROR_ENCLAVE_LOST, "Power transition occurred." },
+  { SGX_ERROR_INVALID_ATTRIBUTE, "Enclave was not authorized." },
+  { SGX_ERROR_INVALID_ENCLAVE, "Invalid enclave image." },
+  { SGX_ERROR_INVALID_ENCLAVE_ID, "Invalid enclave identification." },
+  { SGX_ERROR_INVALID_METADATA, "Invalid enclave metadata." },
+  { SGX_ERROR_INVALID_PARAMETER, "Invalid parameter." },
+  { SGX_ERROR_INVALID_SIGNATURE, "Invalid enclave signature." },
+  { SGX_ERROR_INVALID_VERSION, "Enclave version was invalid." },
+  { SGX_ERROR_MEMORY_MAP_CONFLICT, "Memory map conflicted." },
+  { SGX_ERROR_NO_DEVICE, "Invalid SGX device." },
+  { SGX_ERROR_OUT_OF_EPC, "Out of EPC memory." },
+  { SGX_ERROR_OUT_OF_MEMORY, "Out of memory." },
+  { SGX_ERROR_UNEXPECTED, "Unexpected error occurred." },
+};
+
+/* Check error conditions for loading enclave */
+void print_error_message(sgx_status_t status)
+{
+  size_t idx = 0;
+  size_t ttl = sizeof sgx_errlist/sizeof sgx_errlist[0];
+
+  for (idx = 0; idx < ttl; idx++) {
+    if(status == sgx_errlist[idx].err) {
+      printf("Error: %s\n", sgx_errlist[idx].msg);
+      break;
+    }
+  }
+
+  if (idx == ttl)
+    printf("Error code is 0x%X. Please refer to the \"Intel SGX SDK Developer Reference\" for more details.\n", status);
+}
+
+/* Initialize the enclave:
+ *   Step 1: try to retrieve the launch token saved by last transaction
+ *   Step 2: call sgx_create_enclave to initialize an enclave instance
+ *   Step 3: save the launch token if it is updated
+ */
+int initialize_enclave(void)
+{
+  sgx_launch_token_t token = {0};
+  sgx_status_t sgx_status = SGX_ERROR_UNEXPECTED;
+  int updated = 0;
+
+  /* Step 1: try to retrieve the launch token saved by last transaction
+   *     if there is no token, then create a new one.
+   */
+  FILE *fp = fopen(TOKEN_FILENAME, "rb");
+  if (fp == NULL && (fp = fopen(TOKEN_FILENAME, "wb")) == NULL) {
+    printf("Warning: Failed to create/open the launch token file \"%s\".\n", TOKEN_FILENAME);
+    return -1;
+  }
+
+  /* read the token from saved file */
+  size_t read_num = fread(token, 1, sizeof(sgx_launch_token_t), fp);
+  if (read_num != 0 && read_num != sizeof(sgx_launch_token_t)) {
+    /* if token is invalid, clear the buffer */
+    memset(&token, 0x0, sizeof(sgx_launch_token_t));
+    printf("Warning: Invalid launch token read from \"%s\".\n", TOKEN_FILENAME);
+  }
+
+  /* Step 2: call sgx_create_enclave to initialize an enclave instance */
+  /* Debug Support: set 2nd parameter to 1 */
+  sgx_status = sgx_create_enclave(ENCLAVE_FILENAME, SGX_DEBUG_FLAG, &token, &updated, &global_eid, NULL);
+  if (sgx_status != SGX_SUCCESS) {
+    print_error_message(sgx_status);
+    if (fp != NULL) fclose(fp);
+    return -1;
+  }
+
+  /* Step 3: save the launch token if it is updated */
+  if (updated == 0 || fp == NULL) {
+    /* if the token is not updated, or file handler is invalid, do not perform saving */
+    if (fp != NULL) fclose(fp);
+    return 0;
+  }
+
+  /* reopen the file with write capablity */
+  fp = freopen(TOKEN_FILENAME, "wb", fp);
+  if (fp == NULL) return 0;
+  size_t write_num = fwrite(token, 1, sizeof(sgx_launch_token_t), fp);
+  if (write_num != sizeof(sgx_launch_token_t))
+    printf("Warning: Failed to save launch token to \"%s\".\n", TOKEN_FILENAME);
+  fclose(fp);
+  return 0;
+}
+
+int SGX_CDECL main(int argc, char *argv[]) {
+  if(initialize_enclave() < 0){
+    printf("Failed to initialize enclave.\n");
+    return -1;
+  }
+
+  /* Run TVM within the enclave */
+  int addone_status;
+  sgx_status_t sgx_status = SGX_ERROR_UNEXPECTED;
+  sgx_status = enclave_main(global_eid, &addone_status);
+  if (sgx_status != SGX_SUCCESS) {
+    print_error_message(sgx_status);
+  }
+
+  sgx_destroy_enclave(global_eid);
+
+  if (addone_status == 1) {
+    printf("It works!");
+    return 0;
+  }
+  printf("It doesn't work.");
+  return -1;
+}

apps/sgx/enclave.cc

+#include <dlpack/dlpack.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/runtime/packed_func.h>
+#ifndef _LIBCPP_SGX_CONFIG
+#include <iostream>
+#endif
+
+/* This function mirrors the one in howto_deploy except without the iostream */
+int Verify(tvm::runtime::Module mod, std::string fname) {
+  // Get the function from the module.
+  tvm::runtime::PackedFunc f = mod.GetFunction(fname);
+
+  // Allocate the DLPack data structures.
+  DLTensor* x;
+  DLTensor* y;
+  int ndim = 1;
+  int dtype_code = kDLFloat;
+  int dtype_bits = 32;
+  int dtype_lanes = 1;
+  int device_type = kDLCPU;
+  int device_id = 0;
+  int64_t shape[1] = {10};
+  TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes,
+                device_type, device_id, &x);
+  TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes,
+                device_type, device_id, &y);
+  for (int i = 0; i < shape[0]; ++i) {
+    static_cast<float*>(x->data)[i] = i;
+  }
+
+  // Invoke the function
+  f(x, y);
+
+  // check the output
+  bool all_eq = true;
+  for (int i = 0; i < shape[0]; ++i) {
+    all_eq = all_eq && static_cast<float*>(y->data)[i] == i + 1.0f;
+  }
+
+  return all_eq;
+}
+
+
+extern "C" {
+int enclave_main() {
+  tvm::runtime::Module mod_syslib = (*tvm::runtime::Registry::Get("module._GetSystemLib"))();
+  return Verify(mod_syslib, "addonesys");
+}
+}
+
+#ifndef _LIBCPP_SGX_CONFIG
+int main(void) {
+  tvm::runtime::Module mod_syslib = (*tvm::runtime::Registry::Get("module._GetSystemLib"))();
+  if (Verify(mod_syslib, "addonesys")) {
+    std::cout << "It works!" << std::endl;
+    return 0;
+  }
+  std::cerr << "It doesn't work." << std::endl;
+  return -1;
+}
+#endif

apps/sgx/enclave_config.xml

+<EnclaveConfiguration>
+  <ProdID>0</ProdID>
+  <ISVSVN>0</ISVSVN>
+  <StackMaxSize>0x2000</StackMaxSize>
+  <HeapMaxSize>0x1000</HeapMaxSize>
+  <TCSNum>1</TCSNum>
+  <TCSPolicy>1</TCSPolicy>
+  <DisableDebug>0</DisableDebug>
+  <MiscSelect>0</MiscSelect>
+  <MiscMask>0xFFFFFFFF</MiscMask>
+</EnclaveConfiguration>

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 prepare_test_libs(base_path):
+    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)
+
+    # Compile library in system library mode
+    fadd_syslib = tvm.build(s, [A, B], 'llvm --system-lib', name='addonesys')
+    syslib_path = osp.join(base_path, 'test_addone_sys.o')
+    fadd_syslib.save(syslib_path)
+
+def main():
+    prepare_test_libs(osp.join(CWD, 'lib'))
+
+if __name__ == '__main__':
+    main()

apps/sgx/run_example.sh

+#!/bin/bash
+
+sgx_sdk=${SGX_SDK:=/opt/sgxsdk}
+
+mkdir -p bin lib
+make
+echo "========================="
+LD_LIBRARY_PATH="$sgx_sdk/lib64":${LD_LIBRARY_PATH} bin/test_addone

apps/sgx/test_addone.edl

+enclave {
+    from "sgx_tstdc.edl" import sgx_thread_wait_untrusted_event_ocall, sgx_thread_set_untrusted_event_ocall, sgx_thread_setwait_untrusted_events_ocall, sgx_thread_set_multiple_untrusted_events_ocall;
+
+    trusted {
+        public int enclave_main();
+    };
+};

apps/sgx/tvm_runtime_pack.cc

+/*!
+ * \brief This is an all in one TVM runtime file for use in an SGX enclave.
+ *
+ *   The files included here will be statically linked into the enclave.
+ *   Please refer to the Makefile (rule lib/tvm_runtime_pack.o) for how to build.
+ *
+ */
+#include "../../src/runtime/c_runtime_api.cc"
+#include "../../src/runtime/cpu_device_api.cc"
+#include "../../src/runtime/workspace_pool.cc"
+#include "../../src/runtime/module_util.cc"
+#include "../../src/runtime/module.cc"
+#include "../../src/runtime/registry.cc"
+#include "../../src/runtime/system_lib_module.cc"
+#ifndef _LIBCPP_SGX_CONFIG
+#include "../../src/runtime/file_util.cc"
+#endif

dmlc-core @ 7e84e8b0

-Subproject commit c0871823b518093a0d04d6cba0a3291bc7b31401
+Subproject commit 7e84e8b036a3ff5c0104a3da1f4c7eebf94396ec

include/tvm/runtime/packed_func.h

@@ -640,6 +640,7 @@ inline const char* TypeCode2Str(int type_code) {
   }
 }
 
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
 inline std::ostream& operator<<(std::ostream& os, TVMType t) {  // NOLINT(*)
   os << TypeCode2Str(t.code);
   if (t.code == kHandle) return os;
@@ -649,11 +650,23 @@ inline std::ostream& operator<<(std::ostream& os, TVMType t) {  // NOLINT(*)
   }
   return os;
 }
+#endif
 
 inline std::string TVMType2String(TVMType t) {
+#ifndef _LIBCPP_SGX_NO_IOSTREAMS
   std::ostringstream os;
   os << t;
   return os.str();
+#else
+  std::string repr = "";
+  repr += TypeCode2Str(t.code);
+  if (t.code == kHandle) return repr;
+  repr += std::to_string(static_cast<int>(t.bits));
+  if (t.lanes != 1) {
+    repr += "x" + std::to_string(static_cast<int>(t.lanes));
+  }
+  return repr;
+#endif
 }
 
 inline TVMType String2TVMType(std::string s) {
@@ -674,10 +687,13 @@ inline TVMType String2TVMType(std::string s) {
     scan = s.c_str();
     LOG(FATAL) << "unknown type " << s;
   }
-  unsigned bits = t.bits, lanes = t.lanes;
-  sscanf(scan, "%ux%u", &bits, &lanes);
-  t.bits = static_cast<uint8_t>(bits);
-  t.lanes = static_cast<uint16_t>(lanes);
+  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
+  unsigned bits = strtoul(scan, &xdelim, 10);
+  if (bits != 0) t.bits = static_cast<uint8_t>(bits);
+  if (*xdelim == 'x') {
+    unsigned lanes = strtoul(xdelim + 1, nullptr, 10);
+    t.lanes = static_cast<uint16_t>(lanes);
+  }
   return t;
 }
 

src/runtime/c_runtime_api.cc

@@ -324,9 +324,9 @@ int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
           int ret = func((TVMValue*)args.values, (int*)args.type_codes, // NOLINT(*)
                          args.num_args, rv, resource_handle);
           if (ret != 0) {
-            std::ostringstream os;
-            os << "TVMCall CFunc Error:\n" << TVMGetLastError();
-            throw dmlc::Error(os.str());
+            std::string err = "TVMCall CFunc Error:\n";
+            err += TVMGetLastError();
+            throw dmlc::Error(err);
           }
         });
   } else {
@@ -338,9 +338,9 @@ int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
           int ret = func((TVMValue*)args.values, (int*)args.type_codes, // NOLINT(*)
                          args.num_args, rv, rpack.get());
           if (ret != 0) {
-            std::ostringstream os;
-            os << "TVMCall CFunc Error:\n" << TVMGetLastError();
-            throw dmlc::Error(os.str());
+            std::string err = "TVMCall CFunc Error:\n";
+            err += TVMGetLastError();
+            throw dmlc::Error(err);
           }
       });
   }

src/runtime/cpu_device_api.cc

@@ -28,6 +28,9 @@ class CPUDeviceAPI final : public DeviceAPI {
 #if _MSC_VER
     ptr = _aligned_malloc(nbytes, alignment);
     if (ptr == nullptr) throw std::bad_alloc();
+#elif defined(_LIBCPP_SGX_CONFIG)
+    ptr = memalign(alignment, nbytes);
+    if (ptr == nullptr) throw std::bad_alloc();
 #else
     int ret = posix_memalign(&ptr, alignment, nbytes);
     if (ret != 0) throw std::bad_alloc();

src/runtime/module.cc

@@ -8,7 +8,9 @@
 #include <tvm/runtime/packed_func.h>
 #include <unordered_set>
 #include <cstring>
+#ifndef _LIBCPP_SGX_CONFIG
 #include "./file_util.h"
+#endif
 
 namespace tvm {
 namespace runtime {
@@ -44,6 +46,7 @@ void Module::Import(Module other) {
 
 Module Module::LoadFromFile(const std::string& file_name,
                             const std::string& format) {
+#ifndef _LIBCPP_SGX_CONFIG
   std::string fmt = GetFileFormat(file_name, format);
   CHECK(fmt.length() != 0)
       << "Cannot deduce format of file " << file_name;
@@ -57,6 +60,9 @@ Module Module::LoadFromFile(const std::string& file_name,
       << load_f_name << ") is not presented.";
   Module m = (*f)(file_name, format);
   return m;
+#else
+  LOG(FATAL) << "SGX does not support LoadFromFile";
+#endif
 }
 
 void ModuleNode::SaveToFile(const std::string& file_name,

src/runtime/module_util.cc

@@ -3,7 +3,9 @@
  * \file module_util.cc
  * \brief Utilities for module.
  */
+#ifndef _LIBCPP_SGX_CONFIG
 #include <dmlc/memory_io.h>
+#endif
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/registry.h>
 #include "./module_util.h"
@@ -12,6 +14,7 @@ namespace tvm {
 namespace runtime {
 
 void ImportModuleBlob(const char* mblob, std::vector<Module>* mlist) {
+#ifndef _LIBCPP_SGX_CONFIG
   CHECK(mblob != nullptr);
   uint64_t nbytes = 0;
   for (size_t i = 0; i < sizeof(nbytes); ++i) {
@@ -34,6 +37,9 @@ void ImportModuleBlob(const char* mblob, std::vector<Module>* mlist) {
     Module m = (*f)(static_cast<void*>(stream));
     mlist->push_back(m);
   }
+#else
+  LOG(FATAL) << "SGX does not support ImportModuleBlob";
+#endif
 }
 
 PackedFunc WrapPackedFunc(BackendPackedCFunc faddr,