diff --git a/.gitignore b/.gitignore index 3c968eb3ed47..01f91d69874f 100644 --- a/.gitignore +++ b/.gitignore @@ -182,6 +182,7 @@ cat.jpg docs.tgz cat.png *.mlmodel +tvm_u.* # Mac OS X .DS_Store build* diff --git a/CMakeLists.txt b/CMakeLists.txt index 65a7d9e36e2d..1f03b9f64ab9 100644 --- a/CMakeLists.txt +++ b/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}) diff --git a/cmake/config.cmake b/cmake/config.cmake index c364a88cce11..a92be7ce3008 100644 --- a/cmake/config.cmake +++ b/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) diff --git a/cmake/modules/SGX.cmake b/cmake/modules/SGX.cmake new file mode 100644 index 000000000000..d7b8546d5d41 --- /dev/null +++ b/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_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() diff --git a/src/runtime/sgx/tvm.edl b/src/runtime/sgx/tvm.edl index b4d9852f8499..55c8a878d766 100644 --- a/src/runtime/sgx/tvm.edl +++ b/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); diff --git a/src/runtime/sgx/untrusted/sgx_module.cc b/src/runtime/sgx/untrusted/sgx_module.cc index 8dd696349b05..b1c1692de398 100644 --- a/src/runtime/sgx/untrusted/sgx_module.cc +++ b/src/runtime/sgx/untrusted/sgx_module.cc @@ -4,11 +4,11 @@ * \brief SGX enclave module. */ #include +#include #include #include #include #include -#include #include #include #include @@ -18,6 +18,7 @@ #include #include "../common.h" #include "../../file_util.h" +#include "./tvm_u.h" namespace tvm { namespace runtime {