Add tf_xla_force_compile_on_miss flag to compile every shape/signature unconditionally

[Copilot]

DeviceCompilationProfiler::ShouldCompileCluster uses several heuristics that silently skip compilation on cache misses: megamorphic avoidance, lazy/async thresholds, and an async concurrency cap. This makes it impossible to guarantee every new shape/signature gets compiled via the JIT path.

Changes

• flags.h / flags.cc — New bool tf_xla_force_compile_on_miss field in XlaOpsCommonFlags (default false), exposed as --tf_xla_force_compile_on_miss via TF_XLA_FLAGS.

• device_compilation_profiler.cc — Early-return true from ShouldCompileCluster when the flag is set, bypassing:

  • is_megamorphic gate
  • kDefaultCompilationThreshold (lazy mode)
  • kMaxNumOngoingCompilations (async mode)

• device_compilation_profiler_test.cc — New ForceCompileOnMissFlag test: drives a cluster to megamorphic state, saturates the async concurrency counter, and asserts ShouldCompileCluster returns true for all compile modes (kAsync, kLazy, kStrict) with the flag enabled.

Usage

TF_XLA_FLAGS="--tf_xla_force_compile_on_miss=true" python your_model.py

Original prompt

Create a change on branch cxf-fix-2.20 in repo fenxcc/tensorflow to force XLA DeviceCompiler to compile on every cache miss for every new shape/signature (option 2 behavior).

Goal/behavior:

• For DeviceCompiler JIT path, on cache miss (DeviceCompileState::kUncompiled) always proceed to compilation immediately instead of skipping due to DeviceCompilationProfiler heuristics.
• Disable megamorphic avoidance and compilation thresholds/limits so that DeviceCompilationProfiler::ShouldCompileCluster(...) always returns true (or equivalent) for non-strict modes as well.
• Ensure async mode does not skip compilation due to kMaxNumOngoingCompilations limit; either remove the limit in the heuristic or force synchronous compile in DeviceCompiler when compile_mode==kAsync.
• Keep code changes minimal and well-scoped behind an opt-in flag if feasible (preferred), e.g. a new TF flag in tensorflow/compiler/jit/flags.* like --tf_xla_force_compile_on_miss (name up to you) defaulting to false, and when true enforces the above behavior.
• Add tests if there are existing unit tests for DeviceCompilationProfiler/DeviceCompiler; otherwise add a small targeted test validating that ShouldCompileCluster returns true even when megamorphic/threshold conditions would normally prevent compilation when the new flag is enabled.
• Update any relevant documentation/comments.

Context:

• User wants option 2 behavior: compile every shape/signature.
• Target branch: https://github.com/fenxcc/tensorflow/tree/cxf-fix-2.20

Reference files discussed:

• tensorflow/compiler/jit/device_compiler.h
• tensorflow/compiler/jit/device_compilation_profiler.cc

The following is the prior conversation context from the user's chat exploration (may be truncated):

User: ShouldBeMegamorphic函数的作用是什么？

Assistant: ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) 的作用是：用一个简单的启发式规则判断某个 XLA 编译 cluster（按 function/cluster 统计）是否“可能是 megamorphic（高度形状/签名多态、动态性过强）”，从而在后续决定是否继续触发编译时作为依据。

结合 device_compilation_profiler.cc 里的调用链，它的具体语义可以概括为：

1. 输入含义

   • compile_count：这个 cluster 到目前为止被编译了多少次。
   • execution_count：这个 cluster 到目前为止被执行了多少次。

2. 判定逻辑（启发式）

   const int64_t kCompileThreshold = 10;
   const int64_t kMinExecutionsPerCompile = 50;
   
   return compile_count > kCompileThreshold &&
          execution_count < kMinExecutionsPerCompile * compile_count;

   也就是说，当满足以下两点时返回 true（认为“应当被视为 megamorphic”）：

   • 编译次数已经很多：compile_count > 10
   • 但“每次编译带来的执行回报”偏低：execution_count / compile_count < 50

   直觉上：如果一个 cluster 反复编译（通常意味着输入形状/签名变化太多导致缓存命中差），但执行次数又没有相应增长，那么继续为它做 XLA 编译往往不划算。

3. 它在系统中的实际效果

   • 该函数在 RegisterExecutionForCluster() 中被使用：
     • 每次执行会 execution_count++
     • 如果此前还没标记 megamorphic，则调用 ShouldBeMegamorphic(...)
     • 一旦判定为 megamorphic，会把 stats->is_megamorphic = true（“sticky”，永久生效）
   • 在 DeviceCompilationProfiler::ShouldCompileCluster() 里，如果 is_megamorphic 为 true，会直接返回 false，从而阻止后续继续编译该 cluster（并广播一个 optimization remark：MEGAMORPHIC_FUNCTION）。

总结一句话：ShouldBeMegamorphic 用来识别“编译开销很大但收益很低、形状/签名变化过多导致反复编译”的 cluster，并触发“以后不要再编译它”的策略。

User: 如果我将 ShouldBeMegamorphic 始终设置为返回 false，那么 xla 能否编译所有cluster并处理所有形状？我想知道是否存在某种驱逐机制？例如，当缓存内容变得过大时，移除一些不常使用的cluster。

Assistant: 把 ShouldBeMegamorphic() 永远改成返回 false，只能移除“因为被判定为 megamorphic 而不再编译”的那条刹车；它并不意味着：

1. XLA 就“能够”编译所有 cluster、并“处理所有形状”；
2. 也不意味着缓存会无限安全增长或具备自动驱逐（eviction）机制。

下面分开说明。

---

1) 改成恒 false 后，XLA 会不会编译“所有 cluster 并处理所有形状”？

不会得到这个保证。

(A) 仍然可能“不编译”的路径

即使 megamorphic 逻辑被禁用，DeviceCompilationProfiler::ShouldCompileCluster() 里仍有其它 gating：

• DeviceCompileMode::kLazy：要满足阈值（默认 kDefaultCompilationThreshold = 2）才会编译（但注意：代码里还有“第一次执行一定编译”的逻辑，见下文）。
• DeviceCompileMode::kAsync：有并发上限（kMaxNumOngoingCompilations = kNumAsyncDeviceCompilerThreads），ongoing 太多会拒绝继续异步编译。
• FailOnXlaCompilation()：如果配置为禁止编译，DeviceCompiler::CompileImpl() 在 entry 还未编译时会直接报错（你贴的 device_compiler.h 里就有这段逻辑）。

另外还有更“硬”的原因：XLA 编译本身可能失败（不支持的 op、布局/设备限制、shape 推导失败等），这不是 megamorphic 逻辑能解决的。

(B) “处理所有形状”的含义：不是“一个可执行文件覆盖所有形状”

你引用的注释已经说明了关键点：

Since XLA computations must have static shapes, DeviceCompiler generates a new XLA computation for each new set of input shapes.

也就是：XLA 需要静态 shape；如果你的输入 shape 不断变化，系统往往会为“每一种 shape signature”生成/缓存一个新的编译结果/可执行文件。把 ShouldBeMegamorphic 设为 false，只会让系统更愿意“继续为新 shape 编译”，但不代表“一个 cluster 就能涵盖无限 shape”。

因此在高度动态 shape 的场景，结果很可能是：编译次数暴涨 + cache 变大 + 可能 OOM 或性能变差，而不是“神奇地处理所有形状”。

---

2) 你问的“驱逐机制”：缓存太大时是否会移除不常用 cluster？

就你贴出来的 device_compiler.h 代码而言，有非常明确的信号：目前并没有做 eviction 的结构性支持，至少这段实现仍是“TODO”。

在 CompileImpl() 里有注释（你贴的原文）：

// TODO(phawkins): this locking will need to be restructured when we implement
// cache eviction.

这通常意味着：当前 DeviceCompilationCache 的并发/锁设计还没为 eviction 准备好，也暗示“现在没有真正的 eviction”。

另外你还能看到一个重要行为：

• CompileStrict() 结束时会调用 cache_->Finalize();，注释说的是释放 cache 持有的 XlaComputation 引用（即释放 HLO/计算图层面的对...

This pull request was created from Copilot chat.