diff --git a/tests/compile_bench/io_compute.lean b/tests/compile_bench/io_compute.lean new file mode 100644 index 0000000000..2aeae1d3d3 --- /dev/null +++ b/tests/compile_bench/io_compute.lean @@ -0,0 +1,82 @@ +/-! Mixing basic IO actions and heavy unboxed compute -/ + +def N : Nat := 12500000 / 2 -- ~50 MB + +-- xorshift64* PRNG (same algorithm as Rust version) +@[inline] def xorshift64 (state : UInt64) : UInt64 × UInt64 := + let x := state ^^^ (state >>> 12) + let x := x ^^^ (x <<< 25) + let x := x ^^^ (x >>> 27) + (x, x * 0x2545F4914F6CDD1D) + +-- Push UInt64 as 8 little-endian bytes +@[inline] def pushLE (ba : ByteArray) (v : UInt64) : ByteArray := + let ba := ba.push (v &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 8) &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 16) &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 24) &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 32) &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 40) &&& 0xFF).toUInt8 + let ba := ba.push ((v >>> 48) &&& 0xFF).toUInt8 + ba.push ((v >>> 56) &&& 0xFF).toUInt8 + +-- Read UInt64 from 8 little-endian bytes at offset +@[inline] def readLE (ba : ByteArray) (off : Nat) : UInt64 := + (ba.get! off).toUInt64 + + ((ba.get! (off + 1)).toUInt64 <<< 8) + + ((ba.get! (off + 2)).toUInt64 <<< 16) + + ((ba.get! (off + 3)).toUInt64 <<< 24) + + ((ba.get! (off + 4)).toUInt64 <<< 32) + + ((ba.get! (off + 5)).toUInt64 <<< 40) + + ((ba.get! (off + 6)).toUInt64 <<< 48) + + ((ba.get! (off + 7)).toUInt64 <<< 56) + +def timeS {α : Type} (act : IO α) : IO (α × Float) := do + let t0 ← IO.monoMsNow + let r ← act + let t1 ← IO.monoMsNow + return (r, (t1 - t0).toFloat / 1000.0) + +def main : IO Unit := do + -- Phase 1: Generate & Sort + let (data, elapsed) ← timeS do + let mut state : UInt64 := 42 + let mut arr : Array UInt64 := Array.mkEmpty N + for _ in [:N] do + let (s, v) := xorshift64 state + state := s + arr := arr.push v + return arr.qsortOrd + IO.println s!"measurement: generate_sort {elapsed} s" + + let (_, file) ← IO.FS.createTempFile + -- Phase 2: Write + let (_, elapsed) ← timeS do + let mut ba := ByteArray.emptyWithCapacity (N * 8) + for i in [:data.size] do + ba := pushLE ba data[i]! + IO.FS.writeBinFile file ba + IO.println s!"measurement: write {elapsed} s" + + -- Phase 3: Read + let (data, elapsed) ← timeS do + let ba ← IO.FS.readBinFile file + let mut arr : Array UInt64 := Array.mkEmpty N + for i in [:N] do + arr := arr.push (readLE ba (i * 8)) + return arr + IO.println s!"measurement: read {elapsed} s" + + -- Phase 4: Fisher-Yates shuffle + let (_, elapsed) ← timeS do + let mut arr := data + let mut state : UInt64 := 42 + for i' in [:N - 1] do + let i := N - 1 - i' + let (s, v) := xorshift64 state + state := s + let j := (v % (i + 1).toUInt64).toNat + arr := arr.swapIfInBounds i j + return arr + IO.println s!"measurement: shuffle {elapsed} s" + diff --git a/tests/compile_bench/io_compute.lean.no_interpret b/tests/compile_bench/io_compute.lean.no_interpret new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/compile_bench/io_compute.lean.out.expected b/tests/compile_bench/io_compute.lean.out.expected new file mode 100644 index 0000000000..d58673d342 --- /dev/null +++ b/tests/compile_bench/io_compute.lean.out.expected @@ -0,0 +1,4 @@ +measurement: generate_sort ... +measurement: write ... +measurement: read ... +measurement: shuffle ...