From de420de9574c799549dddcf686f03d28d435ace6 Mon Sep 17 00:00:00 2001
From: Yukio Siraichi <yukio.siraichi@gmail.com>
Date: Fri, 9 Feb 2024 17:01:14 -0300
Subject: [PATCH] [benchmarks] Increase compilation cache. (#6509)

---
 benchmarks/torchbench_model.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/benchmarks/torchbench_model.py b/benchmarks/torchbench_model.py
index fd30ee359871..6fa64b0329ac 100644
--- a/benchmarks/torchbench_model.py
+++ b/benchmarks/torchbench_model.py
@@ -76,6 +76,7 @@
         {
             "test": "eval",
             "xla": "PJRT",
+            "dynamo": None,
         },  # TIMEOUT
     ],
     "hf_T5_generate": [
@@ -85,6 +86,7 @@
         {
             "test": "eval",
             "xla": "PJRT",
+            "dynamo": None,
         },  # TIMEOUT
     ],
     "doctr_det_predictor": [{
@@ -135,6 +137,13 @@
     ],
 }
 
+# Models that had more graphs to be compiled than the actual size of
+# the cache.
+NEED_LARGER_CACHE = {
+    "cm3leon_generate",
+    "hf_T5_generate",
+}
+
 
 class TorchBenchModelLoader(ModelLoader):
 
@@ -401,6 +410,9 @@ def update_process_env(self, process_env):
     if precision_flag is not None:
       process_env[precision_flag] = '1'
 
+    if self.model_name in NEED_LARGER_CACHE:
+      process_env["XLA_COMPILATION_CACHE_SIZE"] = "2048"
+
   def pick_grad(self):
     # special case
     if self.model_name in ("maml",):