Add tf_xla_always_compile_on_cache_miss flag to bypass XLA compilation heuristics

[Copilot]

XLA's DeviceCompilationProfiler uses several heuristics to skip compilation on cache misses: megamorphic detection, lazy compilation thresholds, and async ongoing compilation limits. These heuristics can prevent XLA from compiling clusters with highly dynamic shapes, even when full compilation coverage is desired.

Changes

• flags.h / flags.cc: Add tf_xla_always_compile_on_cache_miss bool to BuildXlaOpsPassFlags (default false). Registered as a TF_XLA_FLAGS entry.

• device_compilation_profiler.cc: When the flag is set, ShouldCompileCluster bypasses all three skip gates:

  • Megamorphic gate (is_megamorphic == true → normally returns false)
  • Async ongoing compilations limit (num_ongoing_compilations_ >= kMaxNumOngoingCompilations)
  • Lazy compilation threshold (current_request_count < compile_threshold)

• device_compilation_profiler_test.cc: Add ShouldCompileClusterAlwaysCompileOnCacheMiss test covering all three bypass paths.

Usage

# Via environment variable
TF_XLA_FLAGS="--tf_xla_always_compile_on_cache_miss=true" ./my_program

// Programmatically
GetBuildXlaOpsPassFlags()->tf_xla_always_compile_on_cache_miss = true;

Note: Enabling this flag on graphs with highly dynamic shapes will cause compilation to trigger on every new shape/signature, potentially leading to increased memory usage and compilation overhead proportional to the number of distinct input shapes seen at runtime.

Original prompt

Implement the requested code changes in fenxcc/tensorflow on ref master, since the current commit has no code modifications yet.

Requirements:

• Apply the previously discussed behavior: option 2 (force XLA to compile on every cache miss for every new shape/signature).
• Ensure actual code modifications are present (not an empty/no-op commit).
• Prefer an opt-in flag (default false) to avoid changing default behavior unexpectedly.
• Update relevant files (likely tensorflow/compiler/jit/device_compilation_profiler.cc and flags in tensorflow/compiler/jit/flags.*) so when the flag is enabled, the profiler/compiler never skips compilation due to megamorphic heuristics, thresholds, or async ongoing compilation limit.
• Add or update tests if feasible.

Repo: fenxcc/tensorflow
Branch/ref: master

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/计算图层面的对象引用），但它不是“把 cache entry 驱逐掉”。entry（含编译结果、可执行文件指针等）仍然会留在 cache 里。

所以回答你的问题：
**在你提供的这些代码上下文中，看不到“缓存过大自动驱逐不常用 cluster”的机制。**更像是“增长型缓存”，配合 Finalize() 释放部分中间对象占用，但不是 LRU/size-based eviction。

---

3) 如果强行让它编译所有 shape，会发生什么风险？

• 编译风暴：动态 shape 多时，每个新 signature 都触发编译。
• 内存增长：缓存里会累积越来越多的 executable / metadata（哪怕 computation 引用被 finalize，executable 仍可能占大头）。
• 性能退化：大量编译会抢占 CPU/线程池；异步模式还可能因为 kMaxNumOngoingCompilations 限制而频繁拒绝/排队。
• 持久化缓存（磁盘）也可能膨胀：你这里还有 DeviceExecutablePersistor，如果 persist 打开，磁盘也可能越来越大。

---

4) 我建议你下一步怎么验证“是否真的没有 eviction”（在 TF 当前 master 上）

你给的两个文件不足以 100% 排除“别处有 eviction”，但可以很快确认：

1. 在 tensorflow/compiler/jit/device_compilation_cache.h/.cc 里搜索：
   • evict, eviction, LRU, max_size, capacity, prune, Trim
2. 搜索 DeviceCompilationCache::Store / LookupOrCreate 是否有 size 上限逻辑
3. 搜索是...

This pull request was created from Copilot chat.