add benchmark for lcm and gcd

Signed-off-by: Brian Grenier <[email protected]>
modularml · Jun 17, 2024 · e1202b4 · e1202b4
1 parent 4df58b1
commit e1202b4
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/stdlib/benchmarks/math/bench_math.mojo b/stdlib/benchmarks/math/bench_math.mojo
@@ -15,6 15,7 @@
 from benchmark import Bench, Bencher, BenchId, keep, BenchConfig, Unit, run
 from random import *
 from math import *
 from builtin._math import lcm, gcd
 
 # ===----------------------------------------------------------------------===#
 # Benchmark Data
@@ -36,7 37,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
@@ -74,6 88,22 @@ fn bench_math3[
  b.iter[call_fn]()
 
 
 # ===----------------------------------------------------------------------===#
 # Benchmark lcm/gcd
 # ===----------------------------------------------------------------------===#
 @parameter
 fn bench_math2[
  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
 # ===----------------------------------------------------------------------===#
@@ -93,4 123,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[lcm]](BenchId("bench_math_lcm"))
  m.bench_function[bench_math2[gcd]](BenchId("bench_math_gcd"))
  m.dump_report()