Fix timer overflow on small, fast snippets (#55200)
Summary:
- Fixes https://github.com/pytorch/pytorch/issues/54114
- Capped estimated block size to the largest multiple of ten less than C++ INT_MAX
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55200
Test Plan: unit test doesn't throw exception as expected
Reviewed By: robieta
Differential Revision: D27542652
Pulled By: naveedgol
fbshipit-source-id: 3ba68ce84d5fa1d8338cdd5c9f9e5d8c9adda51c
diff --git a/test/benchmark_utils/test_benchmark_utils.py b/test/benchmark_utils/test_benchmark_utils.py
index 1ae63f8..0925d51 100644
--- a/test/benchmark_utils/test_benchmark_utils.py
+++ b/test/benchmark_utils/test_benchmark_utils.py
@@ -168,6 +168,17 @@
@slowTest
@unittest.skipIf(IS_SANDCASTLE, "C++ timing is OSS only.")
+ def test_timer_tiny_fast_snippet(self):
+ timer = benchmark_utils.Timer(
+ 'auto x = 1;',
+ timer=timeit.default_timer,
+ language=benchmark_utils.Language.CPP,
+ )
+ median = timer.blocked_autorange().median
+ self.assertIsInstance(median, float)
+
+ @slowTest
+ @unittest.skipIf(IS_SANDCASTLE, "C++ timing is OSS only.")
def test_cpp_timer(self):
timer = benchmark_utils.Timer(
"""
diff --git a/torch/utils/benchmark/utils/timer.py b/torch/utils/benchmark/utils/timer.py
index 2749077..e0e6e90 100644
--- a/torch/utils/benchmark/utils/timer.py
+++ b/torch/utils/benchmark/utils/timer.py
@@ -310,6 +310,9 @@
break
if time_taken > min_run_time:
break
+ # Avoid overflow in C++ pybind11 interface
+ if number * 10 > 2147483647:
+ break
number *= 10
return number
