Update SGX cmake (#1763)

Author: nhynes

.gitignore

@@ -182,6 +182,7 @@ cat.jpg
 docs.tgz
 cat.png
 *.mlmodel
+tvm_u.*
 # Mac OS X
 .DS_Store
 build*

CMakeLists.txt

@@ -32,6 +32,7 @@ tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" O
 tvm_option(USE_STACKVM_RUNTIME "Include stackvm into the runtime" OFF)
 tvm_option(USE_GRAPH_RUNTIME "Build with tiny graph runtime" ON)
 tvm_option(USE_GRAPH_RUNTIME_DEBUG "Build with tiny graph runtime debug mode" OFF)
+tvm_option(USE_SGX "Build with SGX" OFF)
 tvm_option(USE_RTTI "Build with RTTI" ON)
 tvm_option(USE_MSVC_MT "Build with MT" OFF)
 tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
@@ -170,6 +171,7 @@ include(cmake/modules/OpenGL.cmake)
 include(cmake/modules/Vulkan.cmake)
 include(cmake/modules/Metal.cmake)
 include(cmake/modules/ROCM.cmake)
+include(cmake/modules/SGX.cmake)
 include(cmake/modules/LLVM.cmake)
 include(cmake/modules/contrib/BLAS.cmake)
 include(cmake/modules/contrib/Random.cmake)
@@ -179,6 +181,9 @@ include(cmake/modules/contrib/NNPack.cmake)
 add_library(tvm SHARED ${COMPILER_SRCS} ${RUNTIME_SRCS})
 add_library(tvm_topi SHARED ${TOPI_SRCS})
 add_library(tvm_runtime SHARED ${RUNTIME_SRCS})
+if(NOT USE_SGX STREQUAL "OFF")
+  add_dependencies(tvm_runtime sgx_edl)
+endif()
 add_library(nnvm_compiler SHARED ${NNVM_COMPILER_SRCS})
 
 target_link_libraries(tvm ${TVM_LINKER_LIBS} ${TVM_RUNTIME_LINKER_LIBS})

cmake/config.cmake

@@ -62,6 +62,17 @@ set(USE_VULKAN OFF)
 # Whether enable OpenGL runtime
 set(USE_OPENGL OFF)
 
+# Whether to enable SGX runtime
+#
+# Possible values for USE_SGX:
+# - /path/to/sgxsdk: path to Intel SGX SDK
+# - OFF: disable SGX
+#
+# SGX_MODE := HW|SIM
+set(USE_SGX OFF)
+set(SGX_MODE "SIM")
+set(RUST_SGX_SDK "/path/to/rust-sgx-sdk")
+
 # Whether enable RPC runtime
 set(USE_RPC ON)
 

cmake/modules/SGX.cmake

@@ -0,0 +1,37 @@
+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)
+  set(_tvm_edl ${_sgx_src}/tvm.edl)
+  set(_sgx_ustdc ${RUST_SGX_SDK}/sgx_ustdc)
+
+  set(_urts_lib "sgx_urts")
+  if(SGX_MODE STREQUAL "SIM")
+    set(_urts_lib "${_urts_lib}_sim")
+  endif()
+
+  add_custom_command(
+    OUTPUT ${_tvm_u_h}
+    COMMAND ${USE_SGX}/bin/x64/sgx_edger8r --untrusted
+      --untrusted-dir ${_sgx_src}/untrusted
+      --search-path ${USE_SGX}/include --search-path ${RUST_SGX_SDK}/edl
+      ${_tvm_edl}
+    COMMAND sed -i "4i '#include <tvm/runtime/c_runtime_api.h>'" ${_tvm_u_h}
+    DEPENDS ${_tvm_edl}
+  )
+  add_custom_command(
+    OUTPUT ${_sgx_ustdc}/libsgx_ustdc.a
+    COMMAND make
+    WORKING_DIRECTORY ${_sgx_ustdc}
+  )
+  add_custom_target(sgx_edl DEPENDS ${_tvm_u_h} ${_sgx_ustdc}/libsgx_ustdc.a)
+
+  include_directories(${USE_SGX}/include)
+  file(GLOB RUNTIME_SGX_SRCS ${_sgx_src}/untrusted/*.c*)
+  list(APPEND TVM_RUNTIME_LINKER_LIBS
+    -lpthread
+    -L${USE_SGX}/lib64 -l${_urts_lib}
+    -L${RUST_SGX_SDK}/sgx_ustdc -lsgx_ustdc)
+  list(APPEND RUNTIME_SRCS ${RUNTIME_SGX_SRCS})
+endif()

src/runtime/sgx/tvm.edl

@@ -1,5 +1,7 @@
 enclave {
     from "sgx_tstdc.edl" import *;
+    from "sgx_stdio.edl" import *;
+    from "sgx_backtrace.edl" import *;
 
     trusted {
         public void tvm_ecall_init([isptr, user_check] TVMRetValueHandle ret);

src/runtime/sgx/untrusted/sgx_module.cc

@@ -4,11 +4,11 @@
  * \brief SGX enclave module.
  */
 #include <dmlc/logging.h>
+#include <sgx_urts.h>
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/runtime/threading_backend.h>
-#include <sgx_urts.h>
 #include <algorithm>
 #include <fstream>
 #include <iostream>
@@ -18,6 +18,7 @@
 #include <unordered_map>
 #include "../common.h"
 #include "../../file_util.h"
+#include "./tvm_u.h"
 
 namespace tvm {
 namespace runtime {