[GitHub] [tvm] areusch commented on a change in pull request #6964: [µTVM] Add platform timer and RPCTimeEvaluator to enable AutoTVM

[GitBox <gi...@apache.org>]

areusch commented on a change in pull request #6964:
URL: https://github.com/apache/tvm/pull/6964#discussion_r548737828



##########
File path: src/runtime/crt/common/crt_runtime_api.c
##########
@@ -351,10 +375,120 @@ tvm_crt_error_t TVMInitializeRuntime() {
     error = TVMFuncRegisterGlobal("tvm.rpc.server.ModuleGetFunction", &ModuleGetFunction, 0);
   }
 
+  if (error == kTvmErrorNoError) {
+    error = TVMFuncRegisterGlobal("runtime.RPCTimeEvaluator", &RPCTimeEvaluator, 0);
+  }
+
   if (error != kTvmErrorNoError) {
     TVMPlatformMemoryFree(registry_backing_memory, ctx);
     TVMPlatformMemoryFree(func_registry_memory, ctx);
   }
 
   return error;
 }
+
+typedef struct {
+  uint16_t function_index;
+  TVMFunctionHandle func_to_time;
+  TVMContext ctx;
+  int number;
+  int repeat;
+  int min_repeat_ms;
+} time_evaluator_state_t;
+
+static time_evaluator_state_t g_time_evaluator_state;
+
+int RPCTimeEvaluator(TVMValue* args, int* type_codes, int num_args, TVMValue* ret_val,
+                     int* ret_type_code) {
+  ret_val[0].v_handle = NULL;
+  ret_type_code[0] = kTVMNullptr;
+  if (num_args < 8) {
+    TVMAPIErrorf("not enough args");
+    return kTvmErrorFunctionCallNumArguments;
+  }
+  if (type_codes[0] != kTVMModuleHandle || type_codes[1] != kTVMStr ||
+      type_codes[2] != kTVMArgInt || type_codes[3] != kTVMArgInt || type_codes[4] != kTVMArgInt ||
+      type_codes[5] != kTVMArgInt || type_codes[6] != kTVMArgInt || type_codes[7] != kTVMStr) {
+    TVMAPIErrorf("one or more invalid arg types");
+    return kTvmErrorFunctionCallWrongArgType;
+  }
+
+  TVMModuleHandle mod = (TVMModuleHandle)args[0].v_handle;
+  const char* name = args[1].v_str;
+  g_time_evaluator_state.ctx.device_type = args[2].v_int64;
+  g_time_evaluator_state.ctx.device_id = args[3].v_int64;
+  g_time_evaluator_state.number = args[4].v_int64;
+  g_time_evaluator_state.repeat = args[5].v_int64;
+  g_time_evaluator_state.min_repeat_ms = args[6].v_int64;
+
+  int ret_code =
+      TVMModGetFunction(mod, name, /* query_imports */ 0, &g_time_evaluator_state.func_to_time);
+  if (ret_code != 0) {
+    return ret_code;
+  }
+
+  g_time_evaluator_state.function_index++;
+  ret_val[0].v_handle =
+      EncodeFunctionHandle(kTimeEvaluatorModuleIndex, g_time_evaluator_state.function_index);
+  ret_type_code[0] = kTVMPackedFuncHandle;
+  return kTvmErrorNoError;
+}
+
+tvm_crt_error_t RunTimeEvaluator(tvm_function_index_t function_index, TVMValue* args,
+                                 int* type_codes, int num_args, TVMValue* ret_val,
+                                 int* ret_type_code) {
+  if (function_index != g_time_evaluator_state.function_index) {
+    return kTvmErrorTimeEvaluatorBadHandle;
+  }
+
+  // TODO(areusch): should *really* rethink needing to return doubles
+  DLContext result_byte_ctx = {kDLCPU, 0};
+  TVMByteArray* result_byte_arr;
+  tvm_crt_error_t err =
+      TVMPlatformMemoryAllocate(sizeof(TVMByteArray), result_byte_ctx, (void*)&result_byte_arr);
+  if (err != kTvmErrorNoError) {
+    return err;
+  }
+  size_t data_size = sizeof(double) * g_time_evaluator_state.repeat;
+  err = TVMPlatformMemoryAllocate(data_size, result_byte_ctx, (void*)&result_byte_arr->data);
+  if (err != kTvmErrorNoError) {
+    return err;

Review comment:
       good suggestion, done.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org