diff --git a/mace/core/BUILD b/mace/core/BUILD index 48d746169832be86f5b477a88aca2d0377ce1723..1fbeed2f2da0ba71dc8c763b88f4802420941424 100644 --- a/mace/core/BUILD +++ b/mace/core/BUILD @@ -96,6 +96,7 @@ cc_library( deps = [ ":opencl_headers", "//mace/codegen:generated_opencl_dev", + "//mace/utils:utils_hdrs", ], ) diff --git a/mace/core/runtime/opencl/opencl_development.cc b/mace/core/runtime/opencl/opencl_development.cc index a13bb9a34b078d6deaebd29e555bb454e7cf92ea..864a4712e28bf80e33ed48fc11b0520957885203 100644 --- a/mace/core/runtime/opencl/opencl_development.cc +++ b/mace/core/runtime/opencl/opencl_development.cc @@ -5,24 +5,11 @@ #include #include "mace/core/runtime/opencl/cl2_header.h" +#include "mace/utils/utils.h" namespace mace { -namespace { -inline void DecryptOpenCLSource(const std::vector &src, - std::vector *dst) { - dst->reserve(src.size()); - - // Keep consistent with encrypt in python tool - const std::string decrypt_lookup_table = "Xiaomi-AI-Platform-Mace"; - size_t lookup_table_size = decrypt_lookup_table.size(); - for (int i = 0; i < src.size(); i++) { - dst->push_back(src[i] ^ decrypt_lookup_table[i % lookup_table_size]); - } -} -} // namespace - bool GetSourceOrBinaryProgram(const std::string &program_name, const std::string &binary_file_name_prefix, cl::Context &context, @@ -36,9 +23,9 @@ bool GetSourceOrBinaryProgram(const std::string &program_name, return false; } cl::Program::Sources sources; - std::vector decrypt_source; - DecryptOpenCLSource(it_source->second, &decrypt_source); - sources.push_back(std::string(decrypt_source.begin(), decrypt_source.end())); + std::string kernel_source(it_source->second.begin(), it_source->second.end()); + ObfuscateString(&kernel_source); + sources.push_back(kernel_source); *program = cl::Program(context, sources); return true; diff --git a/mace/python/tools/binary_codegen.py b/mace/python/tools/binary_codegen.py index cb1b1280df96b9974c967659afb929f04d537b9e..0deadf6c60d903fc7f4665b8341115a18fe341c6 100644 --- a/mace/python/tools/binary_codegen.py +++ b/mace/python/tools/binary_codegen.py @@ -27,6 +27,7 @@ def generate_cpp_source(): key_size, = struct.unpack("i", binary_array[idx:idx+4]) idx += 4 key, = struct.unpack(str(key_size) + "s", binary_array[idx:idx+key_size]) + key = ''.join(['\\x{:02x}'.format(ord(c)) for c in key]) idx += key_size params_size, = struct.unpack("i", binary_array[idx:idx+4]) idx += 4 diff --git a/mace/utils/tuner.h b/mace/utils/tuner.h index 6296934dbe310fec2baa4f79da468bd5f187a40e..34a62bdfcb5e86760b078d82556e6df27abe72e7 100644 --- a/mace/utils/tuner.h +++ b/mace/utils/tuner.h @@ -15,6 +15,7 @@ #include "mace/utils/logging.h" #include "mace/utils/timer.h" +#include "mace/utils/utils.h" namespace mace { @@ -42,17 +43,19 @@ class Tuner { ¶m_generator, const std::function &)> &func, Timer *timer) { + std::string obfucated_param_key = param_key; + ObfuscateString(&obfucated_param_key); if (IsTuning() && param_generator != nullptr) { // tune std::vector opt_param = default_param; RetType res = Tune(param_generator, func, timer, &opt_param); VLOG(1) << "Tuning result. " << param_key << ": " << internal::MakeString(opt_param); - param_table_[param_key] = opt_param; + param_table_[obfucated_param_key] = opt_param; return res; } else { // run - if (param_table_.find(param_key) != param_table_.end()) { + if (param_table_.find(obfucated_param_key) != param_table_.end()) { VLOG(1) << param_key << ": " << internal::MakeString(param_table_[param_key]); return func(param_table_[param_key]); diff --git a/mace/utils/utils.h b/mace/utils/utils.h index a8b13828de5208047292218a27d76e3f328923b7..6bd0be5291e99747a509404650cb5f72a2a50c7d 100644 --- a/mace/utils/utils.h +++ b/mace/utils/utils.h @@ -48,5 +48,15 @@ inline std::string ToString(T v) { return ss.str(); } +// ObfuscateString(ObfuscateString(str)) ==> str +inline void ObfuscateString(std::string *str) { + // Keep consistent with obfuscation in python tools + const std::string kLookupTable = "Xiaomi-AI-Platform-Mace"; + size_t lookup_table_size = kLookupTable.size(); + for (int i = 0; i < str->size(); i++) { + (*str)[i] = (*str)[i] ^ kLookupTable[i % lookup_table_size]; + } +} + } // namespace mace #endif // MACE_UTILS_UTILS_H_