add benchmark for lcm and gcd

Signed-off-by: Brian Grenier <[email protected]>
modularml · Jul 17, 2024 · 397aed6 · 397aed6
1 parent e2a3587
commit 397aed6
Showing 1 changed file with 43 additions and 0 deletions.
diff --git a/stdlib/benchmarks/math/bench_math.mojo b/stdlib/benchmarks/math/bench_math.mojo
@@ -16,6 16,7 @@ from math import *
 from random import *
 
 from benchmark import Bench, BenchConfig, Bencher, BenchId, Unit, keep, run
 from math.math import lcm, gcd
 
 # ===----------------------------------------------------------------------===#
 # Benchmark Data
@@ -37,7 38,20 @@ fn make_inputs(
  return result
 
 
 fn make_int_inputs(begin: Int, end: Int, num: Int) -> List[Int]:
  if num == 1:
  return List[Int](begin)
 
  var step = (end - begin) // (num - 1)
 
  var result: List[Int] = List[Int]()
  for i in range(num):
  result.append(begin   step * i)
  return result
 
 
 var inputs = make_inputs(0, 10_000, 1_000_000)
 var int_inputs = make_int_inputs(0, 10_000_000, 1_000_000)
 
 # ===----------------------------------------------------------------------===#
 # Benchmark math_func
@@ -75,6 89,33 @@ fn bench_math3[
  b.iter[call_fn]()
 
 
 # ===----------------------------------------------------------------------===#
 # Benchmark lcm/gcd
 # ===----------------------------------------------------------------------===#
 @parameter
 fn bench_math2[math_f2p: fn (Int, Int, /) -> Int](inout b: Bencher) raises:
  @always_inline
  @parameter
  fn call_fn() raises:
  for i in range(len(int_inputs) // 2):
  _ = math_f2p(int_inputs[i], int_inputs[-(i   1)])
 
  b.iter[call_fn]()
 
 
 @parameter
 fn bench_math2_owned[
  math_f2p: fn (owned Int, owned Int, /) -> Int
 ](inout b: Bencher) raises:
  @always_inline
  @parameter
  fn call_fn() raises:
  for i in range(len(int_inputs) // 2):
  _ = math_f2p(int_inputs[i], int_inputs[-(i   1)])
 
  b.iter[call_fn]()
 
 
 # ===----------------------------------------------------------------------===#
 # Benchmark Main
 # ===----------------------------------------------------------------------===#
@@ -94,4 135,6 @@ def main():
  m.bench_function[bench_math[exp]](BenchId("bench_math_exp"))
  m.bench_function[bench_math[erf]](BenchId("bench_math_erf"))
  m.bench_function[bench_math3[fma]](BenchId("bench_math_fma"))
  m.bench_function[bench_math2_owned[lcm]](BenchId("bench_math_lcm"))
  m.bench_function[bench_math2[gcd]](BenchId("bench_math_gcd"))
  m.dump_report()