llvm/cmake/modules/MLGOLower.cmake

# Lower MLGO models to C++ headers via the MLIR-based pipeline.
#
# Each entry in ${models} has the form "cli-flag,path/to/model.mlir,type".
# For entries whose type matches ${target_type}, this function:
#   1. Lowers model MLIR with ${mlir_opt} and translates it with ${mlir_translate}
#      to a C++ header defining Model<N> in ${LLVM_INCLUDE_DIR}/${include_subdir}.
#   2. Generates ${generated_file_basename}.def in ${LLVM_INCLUDE_DIR}/${include_subdir}
#      with one MLGO_MODEL(ClassName, "cli-flag") entry per model.
#   3. Generates ${generated_file_basename}.h in ${LLVM_INCLUDE_DIR}/${include_subdir}
#      including all generated model headers.
#   4. Appends the target driving generation to ${out_var} in the caller's scope.
#   5. Defines LLVM_HAVE_MLIR_LOWERING_<TARGET_TYPE>.
function(mlgo_lower_models models mlir_opt mlir_translate target_type
         include_subdir generated_file_basename out_var)
  set(def_file "${generated_file_basename}.def")
  set(umbrella_header "${generated_file_basename}.h")

  if(NOT IS_ABSOLUTE "${mlir_opt}")
    cmake_path(ABSOLUTE_PATH mlir_opt BASE_DIRECTORY "${CMAKE_BINARY_DIR}")
  endif()

  if(NOT IS_ABSOLUTE "${mlir_translate}")
    cmake_path(ABSOLUTE_PATH mlir_translate BASE_DIRECTORY "${CMAKE_BINARY_DIR}")
  endif()

  string(APPEND def_content
    "/* Auto-generated by CMake */\n"
    "#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 (from the LLVM_MLGO_MODELS flag)
    # 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 "${target_type}_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
    # TODO: Simplify with builtin pipeline for translation.
    # TODO: It isn't ideal to hardcode this pass pipeline here. It would be better to use the transform dialect.
    # TODO: This pipeline is tested independently in the following files. Make sure that tests and this CMake logic share the definition.
    # mlir/test/Integration/Dialect/EmitC/inline-oz-test-model-tosa.mlir
    # mlir/test/Integration/Dialect/EmitC/regalloc-eviction-test-model-tosa.mlir
    set(mlir_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"
      "math-expand-ops{ops=rsqrt}"
      "arith-expand{include-min-max-f=true include-min-max-i=true}"
      "convert-math-to-emitc"
      "convert-arith-to-emitc"
      "wrap-emitc-func-in-class{class-name-format=${class_name}}"
      "mlgo-add-reflection-map{included-field-attrs=tf_saved_model.index_path}"
    )
    string(JOIN "," pass_pipeline ${mlir_passes})
    set(pass_pipeline "builtin.module(${pass_pipeline})")

    # 1. Run MLIR pipeline to lower 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_${model_type}_${class_name} DEPENDS ${header_file})
    list(APPEND mlgo_gen_targets mlgo_model_gen_${model_type}_${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"
      "} // namespace ${class_name}_ns\n"
      "using ${class_name}_ns::${class_name};\n"
    )
  endforeach()

  string(APPEND def_content "\n#undef MLGO_MODEL\n")
  string(APPEND headers_content "} // namespace llvm\n")

  # Stage the generated files, then only update the copies under
  # LLVM_INCLUDE_DIR when their content changed, so that reconfiguring does not
  # trigger rebuilds of their consumers.
  file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/${def_file}.tmp" "${def_content}")
  configure_file("${CMAKE_CURRENT_BINARY_DIR}/${def_file}.tmp"
    "${LLVM_INCLUDE_DIR}/${include_subdir}/${def_file}" COPYONLY)
  file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/${umbrella_header}.tmp" "${headers_content}")
  configure_file("${CMAKE_CURRENT_BINARY_DIR}/${umbrella_header}.tmp"
    "${LLVM_INCLUDE_DIR}/${include_subdir}/${umbrella_header}" COPYONLY)

  if (mlgo_gen_targets)
    add_custom_target(mlgo_models_gen_${target_type} DEPENDS ${mlgo_gen_targets})
    set(${out_var} ${${out_var}} mlgo_models_gen_${target_type} PARENT_SCOPE)
    string(TOUPPER ${target_type} target_type_allcaps)
    add_compile_definitions(LLVM_HAVE_MLIR_LOWERING_${target_type_allcaps})
  endif()
endfunction()