[llvm] [MLGO] Model selection for models lowered through EmitC (PR #212650)

[Paul Kirth via llvm-commits]

================
@@ -0,0 +1,123 @@
+# Compile MLGO models expressed as MLIR to C++ headers via the EmitC pipeline.
+#
+# Each entry of ${models} has the form
+# "cli-flag,path/to/model.mlir,type". Entries whose type matches
+# ${target_type} are lowered with ${mlir_opt} and translated with
+# ${mlir_translate} to a C++ header defining Model<N>, placed in
+# ${LLVM_INCLUDE_DIR}/${include_subdir}. In the same directory, generate
+# ${def_file}, with one MLGO_MODEL(ClassName, "cli-flag") entry per model, and
+# ${umbrella_header}, including all the generated model headers.
+# Append the target driving generation to MLDeps in the caller's scope, and
+# define LLVM_HAVE_EMITC_COMPILE_<TARGET_TYPE>.
+function(mlgo_compile_models models mlir_opt mlir_translate target_type
+         include_subdir def_file umbrella_header)
+  if ("${mlir_opt}" MATCHES "/")
+    get_filename_component(mlir_opt "${mlir_opt}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+  endif()
+  if ("${mlir_translate}" MATCHES "/")
+    get_filename_component(mlir_translate "${mlir_translate}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+  endif()
+
+  set(DEF_CONTENT "/* Auto-generated by CMake */\n")
+  set(DEF_CONTENT "${DEF_CONTENT}#ifndef MLGO_MODEL\n#define MLGO_MODEL(CLASS_NAME, CLI_FLAG)\n#endif\n\n")
+
+  set(HEADERS_CONTENT "/* Auto-generated by CMake */\nnamespace llvm {\n")
+  set(MLGO_GEN_TARGETS "")
+  set(MODEL_INDEX 1)
+
+  foreach(MODEL_INFO IN LISTS models)
+    # Parse comma-separated fields: cli-flag,path/to/model.mlir,type
+    string(REPLACE "," ";" MODEL_FIELDS "${MODEL_INFO}")
+    list(GET MODEL_FIELDS 0 CLI_FLAG)
+    list(GET MODEL_FIELDS 1 MODEL_PATH)
+    list(GET MODEL_FIELDS 2 MODEL_TYPE)
+
+    # Only process models matching target_type
+    if (NOT MODEL_TYPE STREQUAL target_type)
+      continue()
+    endif()
+
+    set(CLASS_NAME "Model${MODEL_INDEX}")
+    math(EXPR MODEL_INDEX "${MODEL_INDEX} + 1")
+
+    if (NOT IS_ABSOLUTE "${MODEL_PATH}")
+      get_filename_component(MODEL_PATH "${MODEL_PATH}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+    endif()
+
+    set(EMITC_MLIR "${CMAKE_CURRENT_BINARY_DIR}/${CLASS_NAME}_emitc.mlir")
+    set(HEADER_FILE "${LLVM_INCLUDE_DIR}/${include_subdir}/${CLASS_NAME}.h")
+
+    # Pass pipeline to lower MLIR models to EmitC dialect
+    set(EMITC_PASSES
+      "func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor)"
+      "symbol-privatize"
+      "scalarize-single-element-tensor-return"
+      "one-shot-bufferize{bufferize-function-boundaries=true function-boundary-type-conversion=identity-layout-map buffer-alignment=0}"
+      "buffer-results-to-out-params{hoist-static-allocs=true}"
+      "func.func(promote-buffers-to-stack)"
+      "buffer-deallocation-pipeline"
+      "func.func(convert-linalg-to-loops)"
+      "expand-strided-metadata"
+      "canonicalize"
+      "memref-elide-reinterpret-cast"
+      "convert-to-emitc"
+      "wrap-emitc-func-in-class{class-name-format=${CLASS_NAME}}"
+      "mlgo-add-reflection-map{included-field-attrs=tf_saved_model.index_path}"
+      "math-expand-ops{ops=rsqrt}"
+      "arith-expand"
+      "convert-math-to-emitc"
+      "convert-arith-to-emitc"
+    )
+    string(JOIN "," PASS_PIPELINE ${EMITC_PASSES})
+    set(PASS_PIPELINE "builtin.module(${PASS_PIPELINE})")
+
+    # 1. Run MLIR pipeline to compile MODEL_PATH to EmitC MLIR
+    add_custom_command(OUTPUT ${EMITC_MLIR}
+      COMMAND ${mlir_opt} "--pass-pipeline=${PASS_PIPELINE}"
+      ${MODEL_PATH} -o ${EMITC_MLIR}
+      DEPENDS ${MODEL_PATH}
+      VERBATIM
+    )
+
+    # 2. Translate EmitC MLIR to C++ header
+    add_custom_command(OUTPUT ${HEADER_FILE}
+      COMMAND ${mlir_translate} -mlir-to-cpp ${EMITC_MLIR} -o ${HEADER_FILE} --mlir-print-stacktrace-on-diagnostic
+      DEPENDS ${EMITC_MLIR}
+      VERBATIM
+    )
+
+    # Set properties so CMake knows these are generated during the build
+    set_source_files_properties(${EMITC_MLIR} PROPERTIES GENERATED 1)
+    set_source_files_properties(${HEADER_FILE} PROPERTIES GENERATED 1)
+
+    # Custom target to force generation of this header
+    add_custom_target(mlgo_model_gen_${CLASS_NAME} DEPENDS ${HEADER_FILE})
+    list(APPEND MLGO_GEN_TARGETS mlgo_model_gen_${CLASS_NAME})
+
+    # Append the model metadata to the .def file
+    string(APPEND DEF_CONTENT "MLGO_MODEL(${CLASS_NAME}, \"${CLI_FLAG}\")\n")
+
+    # Append to the master header content
+    string(APPEND HEADERS_CONTENT "namespace ${CLASS_NAME}_ns {\n#include \"${include_subdir}/${CLASS_NAME}.h\"\n}\nusing ${CLASS_NAME}_ns::${CLASS_NAME};\n")
+  endforeach()
+
+  string(APPEND DEF_CONTENT "\n#undef MLGO_MODEL\n")
+  string(APPEND HEADERS_CONTENT "}\n")
----------------
ilovepi wrote:

Am I reading this right that this closes the namespace? maybe have a comment to that, similar to the thing clang-format adds.

also it feels weird to undef before the namespace ends when you defined before the namespace began ...

https://github.com/llvm/llvm-project/pull/212650