[SERVE][OPENCL] Pre compiler cache enable for opencl (#2954)

Pre compiler cache enable for opencl to avoid kernel compile affecting benchmark
mlc-ai · Sep 30, 2024 · d2cd68e · d2cd68e
1 parent d1ebaa5
commit d2cd68e
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/cpp/serve/function_table.cc b/cpp/serve/function_table.cc
@@ -123,6 +123,14 @@ void FunctionTable::Init(String reload_lib_path, Device device, picojson::object
     } else {
       executable = tvm::runtime::Module::LoadFromFile(reload_lib_path);
       fload_exec = executable->GetFunction("vm_load_executable");
+      /* precompile opencl kernel programs */
+      if (device.device_type == kDLOpenCL) {
+        auto f_get = executable->GetFunction("opencl.GetPreCompiledPrograms", true);
+        std::string ss = f_get();
+        auto bytes = tvm::String(ss);
+        auto f_set = executable->GetFunction("opencl.SetPreCompiledPrograms", true);
+        f_set(tvm::String(bytes));
+      }
       ICHECK(fload_exec.defined()) << "TVM runtime cannot find vm_load_executable";
     }
     this->use_disco = false;