From dc62760e97d7950cb1fec2b24c0e678e895e9fd1 Mon Sep 17 00:00:00 2001
From: nhynes <nhynes@berkeley.edu>
Date: Fri, 13 Apr 2018 14:24:07 -0700
Subject: [PATCH] Improve SGXModule (#1104)
* Improve SGXModule
* Address code review comments
---
src/runtime/sgx/trusted/ecall_registry.h | 2 +-
src/runtime/sgx/trusted/runtime.cc | 2 +-
src/runtime/sgx/tvm.edl | 2 +-
src/runtime/sgx/untrusted/sgx_module.cc | 39 +++++++++++++++++++++---
4 files changed, 38 insertions(+), 7 deletions(-)
diff --git a/src/runtime/sgx/trusted/ecall_registry.h b/src/runtime/sgx/trusted/ecall_registry.h
index 841b9fb7c..4bc476aae 100644
--- a/src/runtime/sgx/trusted/ecall_registry.h
+++ b/src/runtime/sgx/trusted/ecall_registry.h
@@ -82,7 +82,7 @@ std::vector<ECallRegistry> ECallRegistry::exports_;
*/
#define TVM_REGISTER_ENCLAVE_FUNC(OpName) \
TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = \
- ::tvm::runtime::sgx::ECallRegistry::Register(OpName)
+ ::tvm::runtime::sgx::ECallRegistry::Register(OpName, true)
} // namespace sgx
} // namespace runtime
diff --git a/src/runtime/sgx/trusted/runtime.cc b/src/runtime/sgx/trusted/runtime.cc
index 314f2625e..a863327f9 100644
--- a/src/runtime/sgx/trusted/runtime.cc
+++ b/src/runtime/sgx/trusted/runtime.cc
@@ -45,7 +45,7 @@ void tvm_ecall_packed_func(int func_id,
void* ret_buf;
TVM_SGX_CHECKED_CALL(tvm_ocall_reserve_space(
- &ret_buf, bytes.size() + sizeof(TVMByteArray)));
+ &ret_buf, bytes.size() + sizeof(TVMByteArray), sizeof(uint64_t)));
char* data_buf = static_cast<char*>(ret_buf) + sizeof(TVMByteArray);
memcpy(data_buf, bytes.data(), bytes.size());
diff --git a/src/runtime/sgx/tvm.edl b/src/runtime/sgx/tvm.edl
index 1593e9186..b4d9852f8 100644
--- a/src/runtime/sgx/tvm.edl
+++ b/src/runtime/sgx/tvm.edl
@@ -22,7 +22,7 @@ enclave {
[in, count=num_ret] const int* type_code,
int num_ret);
void tvm_ocall_register_export([in, string] const char* name, int func_id);
- void* tvm_ocall_reserve_space(size_t num_bytes);
+ void* tvm_ocall_reserve_space(size_t num_bytes, size_t alignment);
};
};
diff --git a/src/runtime/sgx/untrusted/sgx_module.cc b/src/runtime/sgx/untrusted/sgx_module.cc
index 68bfa6d4c..8dd696349 100644
--- a/src/runtime/sgx/untrusted/sgx_module.cc
+++ b/src/runtime/sgx/untrusted/sgx_module.cc
@@ -5,6 +5,7 @@
*/
#include <dmlc/logging.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>
@@ -156,6 +157,26 @@ TVM_REGISTER_GLOBAL("__sgx_set_last_error__")
TVMAPISetLastError(err.c_str());
});
+TVM_REGISTER_GLOBAL("__sgx_println__")
+.set_body([](TVMArgs args, TVMRetValue* rv) {
+ std::ostringstream msg;
+ for (int i = 0; i < args.num_args; ++i) {
+ switch (args.type_codes[i]) {
+ case kDLInt: msg << static_cast<int64_t>(args[i]); break;
+ case kDLUInt: msg << static_cast<uint64_t>(args[i]); break;
+ case kDLFloat: msg << static_cast<double>(args[i]); break;
+ case kStr:
+ case kBytes: {
+ std::string val = args[i];
+ msg << val;
+ }
+ break;
+ }
+ msg << " ";
+ }
+ LOG(INFO) << msg.str();
+});
+
extern "C" {
void tvm_ocall_register_export(const char* name, int func_id) {
@@ -177,10 +198,20 @@ void tvm_ocall_packed_func(const char* name,
// Allocates space for return values. The returned pointer is only valid between
// successive calls to `tvm_ocall_reserve_space`.
-void* tvm_ocall_reserve_space(size_t num_bytes) {
- static thread_local std::vector<uint64_t> buf;
- buf.reserve(num_bytes);
- return buf.data();
+void* tvm_ocall_reserve_space(size_t num_bytes, size_t alignment) {
+ static TVMContext ctx = { kDLCPU, 0 };
+ static thread_local void* buf = nullptr;
+ static thread_local size_t buf_size = 0;
+ static thread_local size_t buf_align = 0;
+
+ if (buf_size >= num_bytes && buf_align >= alignment) return buf;
+
+ DeviceAPI::Get(ctx)->FreeDataSpace(ctx, buf);
+ buf = DeviceAPI::Get(ctx)->AllocDataSpace(ctx, num_bytes, alignment, {});
+ buf_size = num_bytes;
+ buf_align = alignment;
+
+ return buf;
}
void tvm_ocall_set_return(TVMRetValueHandle ret,
--
GitLab