diff --git a/src/runtime/mpz.h b/src/runtime/mpz.h index 1496d132d69c..bfc7e68c37fb 100644 --- a/src/runtime/mpz.h +++ b/src/runtime/mpz.h @@ -65,6 +65,17 @@ class LEAN_EXPORT mpz { mpz(mpz && s) noexcept; ~mpz(); + /** Whether allocated limb capacity exceeds twice the used size. */ + bool has_excess_capacity() const { +#ifdef LEAN_USE_GMP + size_t used = mpz_size(m_val); + size_t capacity = static_cast(m_val[0]._mp_alloc); + return capacity > used && capacity - used > used; +#else + return false; +#endif + } + #ifdef LEAN_USE_GMP void set(mpz_t r) const; #endif diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp index ddaa25677007..0ed8b44ad4bc 100644 --- a/src/runtime/object.cpp +++ b/src/runtime/object.cpp @@ -1343,13 +1343,14 @@ void deactivate_promise(lean_promise_object * promise) { // ======================================= // Natural numbers -object * alloc_mpz(mpz const & m) { +template +static object * alloc_mpz_core(M && m) { void * mem = lean_alloc_small_object(sizeof(mpz_object)); #ifdef LEAN_MIMALLOC // placement new is not guaranteed to preserve this field so store and restore it unsigned sz = ((lean_object *)mem)->m_cs_sz; #endif - mpz_object * o = new (mem) mpz_object(m); + mpz_object * o = new (mem) mpz_object(std::forward(m)); #ifdef LEAN_MIMALLOC o->m_header.m_cs_sz = sz; #endif @@ -1357,6 +1358,13 @@ object * alloc_mpz(mpz const & m) { return (lean_object*)o; } +object * alloc_mpz(mpz const & m) { return alloc_mpz_core(m); } +object * alloc_mpz(mpz && m) { + if (m.has_excess_capacity()) + return alloc_mpz(static_cast(m)); + return alloc_mpz_core(std::move(m)); +} + #ifdef LEAN_USE_GMP extern "C" LEAN_EXPORT lean_object * lean_alloc_mpz(mpz_t v) { return alloc_mpz(mpz(v)); @@ -1379,6 +1387,13 @@ static inline obj_res mpz_to_nat(mpz const & m) { return mpz_to_nat_core(m); } +static inline obj_res mpz_to_nat(mpz && m) { + if (m.is_size_t() && m.get_size_t() <= LEAN_MAX_SMALL_NAT) + return lean_box(m.get_size_t()); + else + return alloc_mpz(std::move(m)); +} + extern "C" LEAN_EXPORT object * lean_cstr_to_nat(char const * n) { return mpz_to_nat(mpz(n)); } @@ -1649,9 +1664,9 @@ inline object * mpz_to_int_core(mpz const & m) { return alloc_mpz(m); } -static object * mpz_to_int(mpz const & m) { +static object * mpz_to_int(mpz && m) { if (m < LEAN_MIN_SMALL_INT || m > LEAN_MAX_SMALL_INT) - return mpz_to_int_core(m); + return alloc_mpz(std::move(m)); else return lean_box(static_cast(m.get_int())); } diff --git a/src/runtime/object.h b/src/runtime/object.h index b0d57f5e70da..52c724bf4101 100644 --- a/src/runtime/object.h +++ b/src/runtime/object.h @@ -6,6 +6,7 @@ Author: Leonardo de Moura */ #pragma once #include +#include #include #include "runtime/mpz.h" @@ -23,6 +24,7 @@ struct mpz_object { mpz m_value; mpz_object() {} explicit mpz_object(mpz const & m):m_value(m) {} + explicit mpz_object(mpz && m):m_value(std::move(m)) {} }; typedef lean_external_class external_object_class; @@ -175,6 +177,8 @@ inline object* apply_m(object* f, unsigned n, object** args) { return lean_apply // MPZ LEAN_EXPORT object * alloc_mpz(mpz const &); +// Consume private, owned limb storage; copy instead if its capacity is excessive. +LEAN_EXPORT object * alloc_mpz(mpz &&); inline mpz_object * to_mpz(object * o) { lean_assert(is_mpz(o)); return (mpz_object*)o; } // ======================================= diff --git a/tests/bench/mpz_results.cpp b/tests/bench/mpz_results.cpp new file mode 100644 index 000000000000..cc3ed72934e3 --- /dev/null +++ b/tests/bench/mpz_results.cpp @@ -0,0 +1,232 @@ +/* +Copyright (c) 2026 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Author: Kim Morrison +*/ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" void lean_initialize_runtime_module(); +extern "C" lean_object * lean_alloc_mpz(mpz_t); +extern "C" void lean_extract_mpz_value(lean_object *, mpz_t); + +using Clock = std::chrono::steady_clock; +static volatile size_t sink; +static size_t allocations, reallocations, frees, bytes; +static size_t live_bytes, retained_bytes, peak_bytes; +static bool counting; +static void * counted_alloc(size_t n) { + if (counting) { + ++allocations; bytes += n; live_bytes += n; + peak_bytes = std::max(peak_bytes, live_bytes); + } + return std::malloc(n); +} +static void * counted_realloc(void * p, size_t old_n, size_t n) { + if (counting) { + if (old_n > live_bytes) std::abort(); + ++reallocations; bytes += n; live_bytes = live_bytes - old_n + n; + peak_bytes = std::max(peak_bytes, live_bytes); + } + return std::realloc(p, n); +} +static void counted_free(void * p, size_t n) { + if (counting) { + if (n > live_bytes) std::abort(); + ++frees; live_bytes -= n; + } + std::free(p); +} + +static lean_object * to_lean(mpz_t n, bool integer) { + if (integer) { + if (mpz_cmp_si(n, LEAN_MIN_SMALL_INT) >= 0 && mpz_cmp_si(n, LEAN_MAX_SMALL_INT) <= 0) + return lean_int_to_int(static_cast(mpz_get_si(n))); + } else if (mpz_cmp_ui(n, LEAN_MAX_SMALL_NAT) <= 0) { + return lean_box(mpz_get_ui(n)); + } + return lean_alloc_mpz(n); +} + +static void from_lean(mpz_t n, lean_object * o, bool integer) { + if (!lean_is_scalar(o)) lean_extract_mpz_value(o, n); + else if (integer) mpz_set_si(n, lean_scalar_to_int(o)); + else mpz_set_ui(n, lean_unbox(o)); +} + +enum Op { Bridge, IntAdd, IntCancel, IntMul, IntNeg, IntDiv, IntMod, + NatAdd, NatSub, NatMul, NatDiv, NatMod, ParseInt }; +static const char * names[] = {"bridge", "int_add", "int_cancel", "int_mul", "int_neg", + "int_ediv", "int_emod", "nat_add_control", "nat_sub", "nat_mul_mixed", + "nat_div", "nat_mod", "parse_int"}; + +struct Pair { + mpz_t a, b, sum, dividend, negative, remainder; + lean_object * ia, * ib, * na, * nb, * ns, * nd, * id, * minus_a; + std::string decimal; + Pair(gmp_randstate_t state, unsigned bits) { + mpz_inits(a, b, sum, dividend, negative, remainder, nullptr); + mpz_urandomb(a, state, bits); mpz_setbit(a, bits - 1); + mpz_urandomb(b, state, bits); mpz_setbit(b, bits - 1); + mpz_add(sum, a, b); + mpz_fdiv_q_2exp(remainder, b, 1); + mpz_mul(dividend, a, b); mpz_add(dividend, dividend, remainder); + mpz_neg(negative, dividend); + ia = to_lean(a, true); ib = to_lean(b, true); + na = to_lean(a, false); nb = to_lean(b, false); + ns = to_lean(sum, false); nd = to_lean(dividend, false); id = to_lean(negative, true); + mpz_t tmp; mpz_init(tmp); mpz_neg(tmp, a); minus_a = to_lean(tmp, true); mpz_clear(tmp); + std::vector buffer(mpz_sizeinbase(a, 10) + 2); + mpz_get_str(buffer.data(), 10, a); decimal = buffer.data(); + } + ~Pair() { + for (auto * o : {ia, ib, na, nb, ns, nd, id, minus_a}) lean_dec(o); + mpz_clears(a, b, sum, dividend, negative, remainder, nullptr); + } +}; + +static lean_object * run(Op op, Pair & p) { + switch (op) { + case Bridge: return lean_alloc_mpz(p.a); + case IntAdd: return lean_int_add(p.ia, p.ib); + case IntCancel: return lean_int_add(p.ia, p.minus_a); + case IntMul: return lean_int_mul(p.ia, p.ib); + case IntNeg: return lean_int_neg(p.ia); + case IntDiv: return lean_int_ediv(p.id, p.ib); + case IntMod: return lean_int_emod(p.id, p.ib); + case NatAdd: return lean_nat_add(p.na, p.nb); + case NatSub: return lean_nat_sub(p.ns, p.nb); + case NatMul: return lean_nat_mul(p.na, lean_box(3)); + case NatDiv: return lean_nat_div(p.nd, p.nb); + case NatMod: return lean_nat_mod(p.nd, p.nb); + case ParseInt: return lean_cstr_to_int(p.decimal.c_str()); + } + std::abort(); +} + +static void check(Op op, Pair & p) { + mpz_t expected, actual; mpz_inits(expected, actual, nullptr); + switch (op) { + case Bridge: case NatSub: case NatDiv: case ParseInt: mpz_set(expected, p.a); break; + case IntAdd: case NatAdd: mpz_set(expected, p.sum); break; + case IntCancel: mpz_set_ui(expected, 0); break; + case IntMul: mpz_mul(expected, p.a, p.b); break; + case IntNeg: mpz_neg(expected, p.a); break; + case IntDiv: mpz_fdiv_q(expected, p.negative, p.b); break; + case IntMod: mpz_fdiv_r(expected, p.negative, p.b); break; + case NatMul: mpz_mul_ui(expected, p.a, 3); break; + case NatMod: mpz_set(expected, p.remainder); break; + } + auto * o = run(op, p); + from_lean(actual, o, (op >= IntAdd && op <= IntMod) || op == ParseInt); + if (mpz_cmp(expected, actual)) { + std::fprintf(stderr, "incorrect result for %s\n", names[op]); std::abort(); + } + lean_dec(o); + // Check the borrowed inputs again after allocating and freeing the result. + from_lean(actual, p.ia, true); + if (mpz_cmp(actual, p.a)) std::abort(); + from_lean(actual, p.ib, true); + if (mpz_cmp(actual, p.b)) std::abort(); + mpz_clears(expected, actual, nullptr); +} + +static double timed(Op op, std::vector> & pairs, size_t rounds) { + size_t checksum = 0; + auto start = Clock::now(); + for (size_t i = 0; i < rounds; ++i) { + for (auto & p : pairs) { + auto * o = run(op, *p); + if (counting) retained_bytes += live_bytes; + checksum += lean_is_scalar(o) ? lean_unbox(o) : lean_ptr_tag(o); + lean_dec(o); + } + } + sink = checksum; + return std::chrono::duration(Clock::now() - start).count(); +} + +static void capacity_cases() { + mpz_t huge, small, near, dividend, ga, gb; + mpz_inits(huge, small, near, dividend, ga, gb, nullptr); + mpz_setbit(huge, 4096); mpz_add_ui(huge, huge, 1); + mpz_setbit(small, 64); + mpz_add(near, huge, small); + mpz_mul(dividend, huge, huge); mpz_add(dividend, dividend, small); + mpz_mul(ga, huge, small); + mpz_add_ui(gb, huge, 2); mpz_mul(gb, gb, small); + auto * a = to_lean(huge, false); + auto * b = to_lean(near, false); + auto * d = to_lean(dividend, false); + auto * x = to_lean(ga, false); + auto * y = to_lean(gb, false); + std::puts("op,retained_bytes"); + auto measure = [&](char const * name, lean_object * lhs, lean_object * rhs, + lean_object * (*f)(lean_object *, lean_object *)) { + live_bytes = peak_bytes = 0; + counting = true; + auto * result = f(lhs, rhs); + counting = false; + std::printf("%s,%zu\n", name, live_bytes); + mpz_t actual; mpz_init(actual); + from_lean(actual, result, false); + if (mpz_cmp(actual, small)) std::abort(); + mpz_clear(actual); + counting = true; lean_dec(result); counting = false; + if (live_bytes) std::abort(); + }; + measure("int_small_sub", b, a, lean_int_sub); + measure("nat_small_sub", b, a, lean_nat_sub); + measure("int_small_emod", d, a, lean_int_emod); + measure("nat_small_mod", d, a, lean_nat_mod); + measure("nat_small_gcd", x, y, lean_nat_gcd); + for (auto * o : {a, b, d, x, y}) lean_dec(o); + mpz_clears(huge, small, near, dividend, ga, gb, nullptr); +} + +int main(int argc, char ** argv) { + bool capacity = argc > 1 && std::string(argv[1]) == "--capacity"; + bool count = capacity || (argc > 1 && std::string(argv[1]) == "--allocations"); + if (count) mp_set_memory_functions(counted_alloc, counted_realloc, counted_free); + lean_initialize_runtime_module(); + if (capacity) { capacity_cases(); return 0; } + gmp_randstate_t state; gmp_randinit_mt(state); gmp_randseed_ui(state, 15160); + std::puts(count ? "op,bits,allocs,reallocs,frees,requested_bytes,retained_bytes,peak_bytes" : "op,bits,ns"); + for (unsigned bits : {16, 32, 64, 256, 1024, 4096}) { + std::vector> pairs; + for (unsigned i = 0; i < 16; ++i) pairs.push_back(std::make_unique(state, bits)); + for (unsigned code = 0; code <= ParseInt; ++code) { + Op op = static_cast(code); + for (auto & p : pairs) check(op, *p); + if (count) { + allocations = reallocations = frees = bytes = 0; + live_bytes = retained_bytes = peak_bytes = 0; + counting = true; + timed(op, pairs, 1); + counting = false; + if (live_bytes != 0) std::abort(); + std::printf("%s,%u,%.4f,%.4f,%.4f,%.4f,%.4f,%zu\n", names[op], bits, + allocations / 16.0, reallocations / 16.0, frees / 16.0, bytes / 16.0, + retained_bytes / 16.0, peak_bytes); + } else { + size_t rounds = 1; + double seconds; + do { + seconds = timed(op, pairs, rounds); + if (seconds < 0.003) rounds *= 2; + } while (seconds < 0.003); + rounds = std::max(1, static_cast(rounds * 0.025 / seconds)); + seconds = timed(op, pairs, rounds); + std::printf("%s,%u,%.4f\n", names[op], bits, seconds * 1e9 / (rounds * pairs.size())); + } + } + } + gmp_randclear(state); +} diff --git a/tests/bench/mpz_results.md b/tests/bench/mpz_results.md new file mode 100644 index 000000000000..877a2d61d5b1 --- /dev/null +++ b/tests/bench/mpz_results.md @@ -0,0 +1,22 @@ +# Bignum allocation benchmark + +EPYC 9455, GMP 6.3.0, Clang 22, release/mimalloc; medians of nine paired runs on CPU 5 against `dc34e5f5cf9c42dd783471b525abbe66c0194270`. + +| Operation | Bits | Copying ns/op | Moving ns/op | GMP allocations, before → after | +| --- | ---: | ---: | ---: | ---: | +| GMP-to-Lean bridge | 4096 | 42.3 | 32.1 | 2 → 1 | +| Int negation | 4096 | 45.6 | 36.2 | 2 → 1 | +| Int multiplication | 4096 | 1428.1 | 1561.7 | 4 → 3 | +| Nat addition (unchanged control) | 256 | 46.7 | 45.0 | 3 → 3 | + +Moving saves one limb allocation and copy. Conversion gains reproduced on a second core; arithmetic timings were noisy on this shared machine. Copying results with capacity above twice the used limb count bounds retained spare storage. + +Build matching base and optimized release worktrees. Set `GMP_INCLUDE` and `GMP_LIBRARY` from the CMake cache, then compile this harness once and run it against each runtime: + +```sh +c++ -O3 -DNDEBUG -std=c++17 -I"$BASE_WORKTREE/build/release/stage1/include" -I"$GMP_INCLUDE" tests/bench/mpz_results.cpp -L"$BASE_WORKTREE/build/release/stage1/lib/lean" -lleanshared "$GMP_LIBRARY" -o /tmp/mpz-results-bench +LD_LIBRARY_PATH="$BASE_WORKTREE/build/release/stage1/lib/lean" taskset -c 5 /tmp/mpz-results-bench +LD_LIBRARY_PATH="$MOVE_WORKTREE/build/release/stage1/lib/lean" taskset -c 5 /tmp/mpz-results-bench +``` + +Alternate run order across nine pairs. Use `--allocations` for allocation/retention counts and `--capacity` for small-result/large-operand cases. diff --git a/tests/compile/mpz_results.lean b/tests/compile/mpz_results.lean new file mode 100644 index 000000000000..c82550136d7b --- /dev/null +++ b/tests/compile/mpz_results.lean @@ -0,0 +1,45 @@ +module + +/-! +Tests bignum result allocation from compiled and interpreted callers. Retain inputs and multiple +results across later allocations, exercise shared operands, cancellation, signed division, and +conversion across scalar/bignum boundaries, and round-trip decimal representations. +-/ + +private def values (seed : Nat) : Array Nat := + #[0, 1, 2, 3, 2 ^ 30 - 1, 2 ^ 30, 2 ^ 31 - 1, 2 ^ 31, 2 ^ 31 + 1, + 2 ^ 32 - 1, 2 ^ 32, 2 ^ 63 - 1, 2 ^ 63, 2 ^ 64, 2 ^ 128 + 1, + 2 ^ 256 - 1, 2 ^ 1024 + 3].map (· + seed) + +public def main (args : List String) : IO Unit := do + let inputs := values args.length + for a in inputs do + unless a.repr.toNat? = some a do + throw <| IO.userError "Nat decimal round-trip failed" + for b in inputs do + let x : Int := a + let y : Int := b + for s in #[x, -x] do + for t in #[y, -y] do + let results := #[s + t, s - t, s * t, -s, s + s, s - s, s * s] + -- Later allocations must not invalidate either retained inputs or results. + let copies := results.map fun r => r + (2 ^ 2048 : Int) + for r in results do + unless r.repr.toInt? = some r do + throw <| IO.userError "Int decimal round-trip failed" + unless copies.map (· - (2 ^ 2048 : Int)) = results && + results[0]! - t = s && results[1]! + t = s && + results[2]! = t * s && results[3]! = -s && + results[4]! = s * 2 && results[5]! = 0 && results[6]! = s * s && + s.natAbs = a && t.natAbs = b do + throw <| IO.userError "retained Int arithmetic failed" + unless s.ediv t * t + s.emod t = s && s.tdiv t * t + s.tmod t = s do + throw <| IO.userError "signed division failed" + let q := a / b + let r := a % b + let product := a * b + unless q * b + r = a && product = b * a && + (if b = 0 then product = 0 else product / b = a) && + (a - b) + min a b = a && (a ^ 3) = a * a * a && + (a &&& b) + (a ||| b) = a + b do + throw <| IO.userError "retained Nat arithmetic failed" diff --git a/tests/elab/mpz_result_test_extern.lean b/tests/elab/mpz_result_test_extern.lean new file mode 100644 index 000000000000..d77158f682bb --- /dev/null +++ b/tests/elab/mpz_result_test_extern.lean @@ -0,0 +1,13 @@ +module + +import Lean.Util.TestExtern + +/-! Tests Int and Nat bignum results and heap-to-scalar normalization. -/ + +test_extern Int.add (2 ^ 128) (2 ^ 128 + 1) +test_extern Int.sub (-(2 ^ 128)) (2 ^ 128 + 1) +test_extern Int.add (2 ^ 128) (-(2 ^ 128)) +test_extern Int.sub (2 ^ 128) (2 ^ 128 + 1) + +test_extern Nat.sub (2 ^ 128 + 1) 3 +test_extern Nat.div (3 * 2 ^ 128 + 7) (2 ^ 128 + 3)