1function(tf_get_absolute_path path base final_path)
2  if (IS_ABSOLUTE ${path})
3    set(${final_path} ${path} PARENT_SCOPE)
4  else()
5    set(${final_path} ${base}/${path} PARENT_SCOPE)
6  endif()
7endfunction()
8
9function(tf_get_model model final_path)
10  string(FIND ${model} "http:" pos_http)
11  string(FIND ${model} "https:" pos_https)
12  if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
13    message("Downloading model " ${model})
14    string(FIND ${model} "/" fname_start REVERSE)
15    math(EXPR fname_start "${fname_start}+1")
16    string(SUBSTRING ${model} ${fname_start}+1 -1 fname)
17    message("Model archive: " ${fname})
18    file(DOWNLOAD ${model} ${CMAKE_CURRENT_BINARY_DIR}/${fname})
19    file(ARCHIVE_EXTRACT INPUT
20      ${CMAKE_CURRENT_BINARY_DIR}/${fname}
21      DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model)
22    set(${final_path} ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model/model PARENT_SCOPE)
23  else()
24    tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
25    set(${final_path} ${model_path} PARENT_SCOPE)
26  endif()
27endfunction()
28
29# Generate a mock model for tests.
30function(generate_mock_model generator output)
31  tf_get_absolute_path(${generator} ${CMAKE_CURRENT_SOURCE_DIR} generator_absolute_path)
32  tf_get_absolute_path(${output} ${CMAKE_CURRENT_BINARY_DIR} output_absolute_path)
33  message(WARNING "Autogenerated mock models should not be used in production builds.")
34  execute_process(COMMAND python3
35    ${generator_absolute_path}
36    ${output_absolute_path}
37    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
38  )
39endfunction()
40
41# Run the tensorflow compiler (saved_model_cli) on the saved model in the
42# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
43# ${signature_def_key}.
44# Produce a pair of files called ${fname}.h and  ${fname}.o in the
45# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
46# called ${cpp_class} - which may be a namespace-qualified class name.
47function(tfcompile model tag_set signature_def_key fname cpp_class)
48  set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
49  set(obj_file ${prefix}.o)
50  set(hdr_file ${prefix}.h)
51  string(TOUPPER ${fname} fname_allcaps)
52  set(override_header ${LLVM_OVERRIDE_MODEL_HEADER_${fname_allcaps}})
53  set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
54  if (EXISTS "${override_header}" AND EXISTS "${override_object}")
55    configure_file(${override_header} ${hdr_file} COPYONLY)
56    configure_file(${override_object} ${obj_file} COPYONLY)
57    message("Using provided header "
58      ${hdr_file} " and object "   ${obj_file}
59      " files for model " ${model})
60  else()
61    tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
62    message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
63    add_custom_command(OUTPUT ${obj_file} ${hdr_file}
64      COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
65            --multithreading false
66            --dir ${LLVM_ML_MODELS_ABSOLUTE}
67            --tag_set ${tag_set}
68            --signature_def_key ${signature_def_key}
69            --output_prefix ${prefix}
70            --cpp_class ${cpp_class}
71            --target_triple ${LLVM_HOST_TRIPLE}
72    )
73  endif()
74
75  # Aggregate the objects so that results of different tfcompile calls may be
76  # grouped into one target.
77  set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
78  set_source_files_properties(${obj_file} PROPERTIES
79    GENERATED 1 EXTERNAL_OBJECT 1)
80
81  set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
82  set_source_files_properties(${hdr_file} PROPERTIES
83    GENERATED 1)
84
85endfunction()
86
87function(tf_find_and_compile model default_url default_path test_model_generator tag_set signature_def_key fname cpp_class)
88  if ("${model}" STREQUAL "download")
89    # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
90    if ("${default_url}" STREQUAL "TO_BE_UPDATED")
91        message(FATAL_ERROR "Default URL was set to 'download' but there is no model url currently specified in cmake - likely, the model interface recently changed, and so there is not a released model available.")
92    endif()
93
94    set(model ${default_url})
95  endif()
96
97  if ("${model}" STREQUAL "autogenerate")
98    set(model ${default_path}-autogenerated)
99    generate_mock_model(${test_model_generator} ${model})
100  endif()
101
102  tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
103  tfcompile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class})
104
105  set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
106  set_source_files_properties(${obj_file} PROPERTIES
107    GENERATED 1 EXTERNAL_OBJECT 1)
108
109  set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
110  set_source_files_properties(${hdr_file} PROPERTIES
111    GENERATED 1)
112
113  set(GeneratedMLSources ${GeneratedMLSources} ${GENERATED_HEADERS} PARENT_SCOPE)
114  set(MLDeps ${MLDeps} tf_xla_runtime PARENT_SCOPE)
115  set(MLLinkDeps ${MLLinkDeps} tf_xla_runtime ${GENERATED_OBJS} PARENT_SCOPE)
116
117endfunction()
118