Skip to content

Commit

Permalink
[mlgo] Fetch models from path / URL
Browse files Browse the repository at this point in the history
Allow custom location for pre-trained models used when AOT-compiling
policies.

Differential Revision: https://reviews.llvm.org/D96796
  • Loading branch information
mtrofin committed Feb 17, 2021
1 parent a3c783d commit 33481c9
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 8 deletions.
20 changes: 13 additions & 7 deletions llvm/cmake/modules/TensorFlowCompile.cmake
Original file line number Diff line number Diff line change
@@ -1,17 +1,23 @@
# Ensure the ${model} is available at ${final_path}.
#
function(tfgetmodel model final_path)
if (IS_ABSOLUTE ${model})
set(${final_path} ${model} PARENT_SCOPE)
else()
set(${final_path}
${CMAKE_CURRENT_SOURCE_DIR}/${model} PARENT_SCOPE)
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}.
# Produce a pair of files called ${fname}.h and ${fname}.o in the
# ${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)
Expand Down
5 changes: 4 additions & 1 deletion llvm/lib/Analysis/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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 (path to saved model)")
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}
Expand Down

0 comments on commit 33481c9

Please sign in to comment.