[PATCH] D96796: [mlgo] Fetch models from path / URL

[Mircea Trofin via Phabricator via llvm-commits]

mtrofin created this revision.
mtrofin added a reviewer: phosek.
Herald added subscribers: hiraditya, mgorny.
mtrofin requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

Allow custom location for pre-trained models used when AOT-compiling
policies. The location may be a directory (relative or absolute paths)
or a http/s URL. In the latter case, it is assumed the URL points at an
archive which, when uncompressed, expands to a saved model (i.e. the
root of the archive contains saved_model.pb[txt] and the variables
directory)


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D96796

Files:
  llvm/cmake/modules/TensorFlowCompile.cmake
  llvm/lib/Analysis/CMakeLists.txt


Index: llvm/lib/Analysis/CMakeLists.txt
===================================================================
--- llvm/lib/Analysis/CMakeLists.txt
+++ llvm/lib/Analysis/CMakeLists.txt
@@ -1,7 +1,10 @@
 if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
   if (DEFINED LLVM_HAVE_TF_AOT)
+    set(LLVM_INLINER_MODEL_PATH "models/inliner" 
+      CACHE STRING 
+      "ML-driven inliner policy location - URL to archive; or path to saved model (relative to this directory, or absolute)")
     include(TensorFlowCompile)
-    tfcompile(models/inliner serve action InlinerSizeModel llvm::InlinerSizeModel)
+    tfcompile(${LLVM_INLINER_MODEL_PATH} serve action InlinerSizeModel llvm::InlinerSizeModel)
     list(APPEND GeneratedMLSources
       $<TARGET_OBJECTS:tf_xla_runtime_objects>
       ${GENERATED_OBJS}
Index: llvm/cmake/modules/TensorFlowCompile.cmake
===================================================================
--- llvm/cmake/modules/TensorFlowCompile.cmake
+++ llvm/cmake/modules/TensorFlowCompile.cmake
@@ -1,3 +1,36 @@
+# Ensure the ${model} is available at ${final_path}.
+#
+# If the model is a url (starts with http/https), then it's assumed to reference
+# an archive. The last path component is assumed to be the name of the archive.
+# The archive is downloaded and deflated under the ${CMAKE_CURRENT_BINARY_DIR}
+# directory, and the resulting directory is the ${final_path}.
+#
+# If the model is not a url, we assume it's a directory, resolve the absolute
+# path and return it.
+function(tfgetmodel model final_path)
+  string(FIND ${model} "http:" pos_http)
+  string(FIND ${model} "https:" pos_https)
+  if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
+    message("Downloading model " ${model})
+    string(FIND ${model} "/" fname_start REVERSE)
+    MATH(EXPR fname_start "${fname_start}+1")
+    string(SUBSTRING ${model} ${fname_start}+1 -1 fname)
+    message("Model archive: " ${fname})
+    file(DOWNLOAD ${model} ${CMAKE_CURRENT_BINARY_DIR}/${fname})
+    file(ARCHIVE_EXTRACT INPUT 
+      ${CMAKE_CURRENT_BINARY_DIR}/${fname}
+      DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model)
+    set(${final_path} ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model PARENT_SCOPE)
+  else()
+    if (IS_ABSOLUTE ${model})
+      set(${final_path} ${model} PARENT_SCOPE)
+    else()
+      set(${final_path}
+        ${CMAKE_CURRENT_SOURCE_DIR}/${model} PARENT_SCOPE)
+    endif()
+  endif()
+endfunction()
+
 # Run the tensorflow compiler (saved_model_cli) on the saved model in the 
 # ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
 # ${signature_def_key}.
@@ -5,13 +38,8 @@
 # ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
 # called ${cpp_class} - which may be a namespace-qualified class name.
 function(tfcompile model tag_set signature_def_key fname cpp_class)
-  if (IS_ABSOLUTE ${model})
-    set(LLVM_ML_MODELS_ABSOLUTE ${model})
-  else()
-    set(LLVM_ML_MODELS_ABSOLUTE
-      ${CMAKE_CURRENT_SOURCE_DIR}/${model})
-  endif()
-
+  tfgetmodel(${model} LLVM_ML_MODELS_ABSOLUTE)
+  message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
   set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
   set(obj_file ${prefix}.o)
   set(hdr_file ${prefix}.h)


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D96796.324030.patch
Type: text/x-patch
Size: 3245 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210216/23216d85/attachment.bin>