adjust cache size

Signed-off-by: Anatoly Myachev <[email protected]>
intel · Sep 20, 2024 · a59c94e · a59c94e
1 parent 49fc392
commit a59c94e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/benchmarks/triton_kernels_benchmark/benchmark_testing.py b/benchmarks/triton_kernels_benchmark/benchmark_testing.py
@@ -64,7 +64,10 @@ def do_bench_ipex(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, fas
     # We maintain a buffer of 256 MB that we clear
     # before each kernel call to make sure that the L2
     # doesn't contain any input data before the run
-    cache_size = 256 * 1024 * 1024
+    factor = 1
+    if os.getenv("ZE_FLAT_DEVICE_HIERARCHY", "FLAT") == "COMPOSITE":
+        factor = 2
+    cache_size = factor * 256 * 1024 * 1024
     if fast_flush:
         cache = torch.empty(int(cache_size // 4), dtype=torch.int, device=device)
     else:

diff --git a/python/triton/testing.py b/python/triton/testing.py
@@ -169,7 +169,10 @@ def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, fast_flu
     # We maintain a buffer of 256 MB that we clear
     # before each kernel call to make sure that the L2 cache
     # doesn't contain any input data before the run
-    cache_size = 256 * 1024 * 1024
+    factor = 1
+    if os.getenv("ZE_FLAT_DEVICE_HIERARCHY", "FLAT") == "COMPOSITE":
+        factor = 2
+    cache_size = factor * 256 * 1024 * 1024
     if fast_flush:
         cache = torch.empty(int(cache_size // 4), dtype=torch.int, device=device_type)
     else: