Correct units for time printouts

do_bench returns ms rather than sec

stack-info: PR: #133, branch: jansel/stack/29

Author: jansel

README.md

@@ -129,11 +129,11 @@ typical autotuning session produces output similar to:
 
 ```
 [0s] Starting DifferentialEvolutionSearch with population=40, generations=20, crossover_rate=0.8
-[20s] Initial population: failed=10 min=0.9677s mid=3.0013s max=22.1430s best=Config(block_sizes=[[64, 32], [32]], loop_orders=[[1, 0]], num_warps=2, num_stages=2, indexing='pointer', l2_grouping=1, use_yz_grid=False)
-[52s] Generation 2: replaced=16 min=0.7731s mid=1.7203s max=3.1227s best=Config(block_sizes=[[32, 128], [16]], loop_orders=[[0, 1]], num_warps=4, num_stages=4, indexing='block_ptr', l2_grouping=16)
-[85s] Generation 3: replaced=19 min=0.6256s mid=1.3916s max=2.7868s best=Config(block_sizes=[[64, 128], [16]], loop_orders=[[0, 1]], num_warps=4, num_stages=4, indexing='block_ptr', l2_grouping=16)
+[20s] Initial population: failed=10 min=0.9677 mid=3.0013 max=22.1430 best=Config(block_sizes=[[64, 32], [32]], loop_orders=[[1, 0]], num_warps=2, num_stages=2, indexing='pointer', l2_grouping=1, use_yz_grid=False)
+[52s] Generation 2: replaced=16 min=0.7731 mid=1.7203 max=3.1227 best=Config(block_sizes=[[32, 128], [16]], loop_orders=[[0, 1]], num_warps=4, num_stages=4, indexing='block_ptr', l2_grouping=16)
+[85s] Generation 3: replaced=19 min=0.6256 mid=1.3916 max=2.7868 best=Config(block_sizes=[[64, 128], [16]], loop_orders=[[0, 1]], num_warps=4, num_stages=4, indexing='block_ptr', l2_grouping=16)
 ...
-[593s] Generation 19: replaced=7 min=0.6072s mid=0.6626s max=0.7496s best=Config(block_sizes=[[64, 128], [16]], loop_orders=[[1, 0]], num_warps=4, num_stages=3, indexing='block_ptr', l2_grouping=32)
+[593s] Generation 19: replaced=7 min=0.6072 mid=0.6626 max=0.7496 best=Config(block_sizes=[[64, 128], [16]], loop_orders=[[1, 0]], num_warps=4, num_stages=3, indexing='block_ptr', l2_grouping=32)
 [593s] Autotuning complete in 593.1s after searching 1520 configs.
 One can hardcode the best config and skip autotuning with:
     @helion.kernel(config=helion.Config(block_sizes=[[64, 128], [16]], loop_orders=[[1, 0]], num_warps=4, num_stages=3, indexing='block_ptr', l2_grouping=32))

examples/add.py

@@ -32,7 +32,7 @@ def check(m: int, n: int) -> None:
     sec = do_bench(lambda: add(x, y))
     baseline_sec = do_bench(lambda: torch.add(x, y))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

examples/attention.py

@@ -109,7 +109,7 @@ def test(
     flex_sec = do_bench(lambda: flex_attention(q, k, v))
     helion_sec = do_bench(lambda: attention(q, k, v))
     print(
-        f"Helion time: {helion_sec:.4f}s, flex time: {flex_sec:.4f}, torch time: {spda_sec:.4f}"
+        f"Helion time: {helion_sec:.4f}ms, flex time: {flex_sec:.4f}, torch time: {spda_sec:.4f}"
     )
 
 

examples/bmm.py

@@ -35,7 +35,7 @@ def check(b: int, m: int, k: int, n: int) -> None:
     sec = do_bench(lambda: bmm(x, y))
     baseline_sec = do_bench(lambda: torch.bmm(x, y))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

examples/concatenate.py

@@ -41,7 +41,7 @@ def main() -> None:
     sec = do_bench(lambda: concat2d_dim1(x, y))
     baseline_sec = do_bench(lambda: torch.cat([x, y], dim=1))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

examples/embedding.py

@@ -34,7 +34,7 @@ def main() -> None:
     sec = do_bench(lambda: embedding(x, weight))
     baseline_sec = do_bench(lambda: torch.nn.functional.embedding(x, weight))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

examples/long_sum.py

@@ -95,7 +95,7 @@ def check(m: int, n: int) -> None:
     manual_loop_sec = do_bench(lambda: longsum_manual(x))
     baseline_sec = do_bench(lambda: baseline_sum(x))
     print(
-        f"Helion Naive time: {sec:.4f}s, Helion Looped Time: {loop_sec:.4f},  Helion Manual Loop Time: {manual_loop_sec:.4f} torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x {baseline_sec / loop_sec:.2f}x {baseline_sec / manual_loop_sec:.2f}x"
+        f"Helion Naive time: {sec:.4f}ms, Helion Looped Time: {loop_sec:.4f},  Helion Manual Loop Time: {manual_loop_sec:.4f} torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x {baseline_sec / loop_sec:.2f}x {baseline_sec / manual_loop_sec:.2f}x"
     )
 
 

examples/matmul.py

@@ -33,7 +33,7 @@ def check(m: int, k: int, n: int) -> None:
     sec = do_bench(lambda: matmul(x, y))
     baseline_sec = do_bench(lambda: torch.matmul(x, y))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

examples/moe_matmul_ogs.py

@@ -172,7 +172,7 @@ def check(T: int, K: int, N: int, n_experts: int) -> None:
         lambda: moe_matmul_ogs_reference(A, W, top1_expert_per_token)
     )
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}s, speed-up: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speed-up: {baseline_sec / sec:.2f}x"
     )
 
 

examples/softmax.py

@@ -64,7 +64,7 @@ def check(m: int, n: int) -> None:
     sec = do_bench(lambda: softmax(x))
     baseline_sec = do_bench(lambda: torch.nn.functional.softmax(x, dim=1))
     print(
-        f"Helion time: {sec:.4f}s, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
+        f"Helion time: {sec:.4f}ms, torch time: {baseline_sec:.4f}, speedup: {baseline_sec / sec:.2f}x"
     )
 
 

helion/autotuner/base_search.py

@@ -110,7 +110,7 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
             )
             t2 = time.perf_counter()
             self.log.debug(
-                lambda: f"result: {res:.4f}s (took {t1 - t0:.1f}s + {t2 - t1:.1f}s)",
+                lambda: f"result: {res:.4f}ms (took {t1 - t0:.1f}s + {t2 - t1:.1f}s)",
             )
             return res
         except OutOfResources:
@@ -336,15 +336,15 @@ def population_statistics(population: list[PopulationMember]) -> str:
         working = [x for x in population if not math.isinf(x.perf)]
         return (
             f"failed={len(population) - len(working)} "
-            f"min={working[0].perf:.4f}s "
-            f"mid={working[len(working) // 2].perf:.4f}s "
-            f"max={working[-1].perf:.4f}s "
+            f"min={working[0].perf:.4f} "
+            f"mid={working[len(working) // 2].perf:.4f} "
+            f"max={working[-1].perf:.4f} "
             f"best={population[0].config!s}"
         )
     return (
-        f"min={population[0].perf:.4f}s "
-        f"mid={population[len(population) // 2].perf:.4f}s "
-        f"max={population[-1].perf:.4f}s "
+        f"min={population[0].perf:.4f} "
+        f"mid={population[len(population) // 2].perf:.4f} "
+        f"max={population[-1].perf:.4f} "
         f"best={population[0].config!s}"
     )