Skip to content

Commit e2d0751

Browse files
committed
Include C cache preloading in profile output
Also don't do it if there's no C compilation
1 parent 7ecb9f8 commit e2d0751

File tree

2 files changed

+14
-4
lines changed

2 files changed

+14
-4
lines changed

pytensor/compile/function/types.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1474,15 +1474,20 @@ def __init__(
14741474
if getattr(mode, "profile", None):
14751475
raise TypeError("profile passed via 'mode'. This isn't supported anymore")
14761476
self.profile = profile
1477-
if profile:
1477+
if profile and config.cxx:
14781478
# This is very important:
14791479
# 1) We preload the cache here to not have its timing
14801480
# included with the rewrites.
14811481
# 2) Do not refresh the cache here by default. It cause
14821482
# too much execution time during testing as we compile
14831483
# much more functions then the number of compile c
14841484
# module.
1485+
start_get_cache = time.perf_counter()
14851486
pytensor.link.c.basic.get_module_cache().refresh()
1487+
get_cache_time = time.perf_counter() - start_get_cache
1488+
self.profile.linker_time += get_cache_time
1489+
self.profile.preload_cache_time += get_cache_time
1490+
14861491
# Handle the case where inputs and/or outputs is a single
14871492
# Variable (not in a list)
14881493
unpack_single = False
@@ -1722,7 +1727,8 @@ def orig_function(
17221727
17231728
"""
17241729

1725-
t1 = time.perf_counter()
1730+
if profile:
1731+
t1 = time.perf_counter()
17261732
mode = pytensor.compile.mode.get_mode(mode)
17271733

17281734
inputs = list(map(convert_function_input, inputs))
@@ -1755,8 +1761,8 @@ def orig_function(
17551761
with config.change_flags(compute_test_value="off"):
17561762
fn = m.create(defaults)
17571763
finally:
1758-
t2 = time.perf_counter()
1759-
if fn and profile:
1764+
if profile and fn:
1765+
t2 = time.perf_counter()
17601766
profile.compile_time += t2 - t1
17611767
# TODO: append
17621768
profile.nb_nodes = len(fn.maker.fgraph.apply_nodes)

pytensor/compile/profiling.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,9 @@ def reset(self):
269269
linker_time: float = 0.0
270270
# time spent linking graph (FunctionMaker.create)
271271

272+
preload_cache_time: float = 0.0
273+
# time spent preloading the cache, so it does not affect rewrites profiling
274+
272275
import_time: float = 0.0
273276
# time spent in importing compiled python module.
274277

@@ -811,6 +814,7 @@ def summary_function(self, file):
811814
),
812815
file=file,
813816
)
817+
print(f" C-cache preloading {self.preload_cache_time:e}s", file=file)
814818
print(f" Import time {self.import_time:e}s", file=file)
815819
print(
816820
f" Node make_thunk time {self.linker_node_make_thunks:e}s", file=file

0 commit comments

Comments
 (0)