[all-commits] [llvm/llvm-project] 313b1a: [mlgo] Support composite AOT-ed models (#96276)

[Mircea Trofin via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 313b1a8250d3cd742356cc71c36402b10d8b2e16
      https://github.com/llvm/llvm-project/commit/313b1a8250d3cd742356cc71c36402b10d8b2e16
  Author: Mircea Trofin <mtrofin at google.com>
  Date:   2024-06-24 (Mon, 24 Jun 2024)

  Changed paths:
    M llvm/include/llvm/Analysis/ReleaseModeModelRunner.h
    M llvm/lib/Analysis/MLInlineAdvisor.cpp
    M llvm/unittests/Analysis/MLModelRunnerTest.cpp

  Log Message:
  -----------
  [mlgo] Support composite AOT-ed models (#96276)

This applies to the AOT case where we embed models in the compiler. The
change adds support for multiple models for the same agent, and allows
the user select one via a command line flag. "agent" refers to e.g. the
inline advisor or the register allocator eviction advisor.

To avoid build setup complexity, the support is delegated to the saved
model. Since saved models define computational graphs, we can generate a
composite model (this happens prior to building and embedding it in LLVM
and is not shown in this change) that exposes an extra feature with a
predefined name: `_model_selector`. The model, then, delegates
internally to contained models based on that feature value.

Model selection is expected to happen at model instantiation, there is
no current scenario for switching them afterwards.

If the model doesn't expose such a feature but the user passes one, we
report error.

If the model exposes such a feature but the user doesn't pass one, we
also report an error.

Invalid model selector values are expected to be handled by the saved
model.

Internally, the model uses a pair of uint64 values - the high and low of
the MD5 hash of the name.

A tool composing models would, then, need to:
- expose the extra feature, `_model_selector`, shape (2,), uint64 data
type
- test its value (`tf.cond` or `tf.case` in Tensorflow) against the MD5
hash, in the [high, low] order, of contained models based on a
user-specified name (which the user will then use as flag value to the
compiler)

Agents just need to add a flag to capture the name of a model and pass
it to `ReleaseModeModelRunner` at construction. This can be passed in
all cases without checking - the case where the model is not composite
and we pass an empty name, everything works as before.

This change also factors out the string flags we pass to the
`ReleaseModeModelRunner` for better maintainability (we risk confusing
parameters that are strings otherwise)



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications