add methods to count allocated bytes; restrict memory usage in zeta sum sieving

acb.h

@@ -1140,6 +1140,24 @@ void acb_nth_root(acb_t res, ulong order, slong prec);
 void _acb_vec_nth_roots_pe(acb_ptr z, slong p, slong e, slong len, slong step, slong prec);
 void _acb_vec_nth_roots(acb_ptr z, slong len, slong prec);
 
+ACB_INLINE slong
+acb_allocated_bytes(const acb_t x)
+{
+    return arb_allocated_bytes(acb_realref(x)) + arb_allocated_bytes(acb_imagref(x));
+}
+
+ACB_INLINE slong
+_acb_vec_allocated_bytes(acb_srcptr vec, slong len)
+{
+    return _arb_vec_allocated_bytes((arb_srcptr) vec, 2 * len);
+}
+
+ACB_INLINE double
+_acb_vec_estimate_allocated_bytes(slong len, slong prec)
+{
+    return 2 * _arb_vec_estimate_allocated_bytes(len, prec);
+}
+
 #ifdef __cplusplus
 }
 #endif

acb_mat.h

@@ -337,6 +337,12 @@ void acb_mat_charpoly(acb_poly_t cp, const acb_mat_t mat, slong prec);
 
 void acb_mat_trace(acb_t trace, const acb_mat_t mat, slong prec);
 
+ACB_MAT_INLINE slong
+acb_mat_allocated_bytes(const acb_mat_t x)
+{
+    return _acb_vec_allocated_bytes(x->entries, x->r * x->c) + x->r * sizeof(acb_ptr);
+}
+
 #ifdef __cplusplus
 }
 #endif

acb_poly.h

@@ -753,6 +753,12 @@ void acb_poly_elliptic_p_series(acb_poly_t res, const acb_poly_t z, const acb_t
 void _acb_poly_erf_series(acb_ptr g, acb_srcptr h, slong hlen, slong n, slong prec);
 void acb_poly_erf_series(acb_poly_t g, const acb_poly_t h, slong n, slong prec);
 
+ACB_POLY_INLINE slong
+acb_poly_allocated_bytes(const acb_poly_t x)
+{
+    return _acb_vec_allocated_bytes(x->coeffs, x->alloc);
+}
+
 #ifdef __cplusplus
 }
 #endif

acb_poly/zeta_em_sum.c

@@ -28,6 +28,8 @@ void _acb_poly_mullow_cpx(acb_ptr res, acb_srcptr src, slong len, const acb_t c,
     acb_mul(res, src, c, prec);
 }
 
+/* todo: don't hardcode this */
+#define SIEVE_ALLOC_LIMIT 4e9   /* 4 GB */
 
 void
 _acb_poly_zeta_em_sum(acb_ptr z, const acb_t s, const acb_t a, int deflate, ulong N, ulong M, slong d, slong prec)
@@ -48,7 +50,7 @@ _acb_poly_zeta_em_sum(acb_ptr z, const acb_t s, const acb_t a, int deflate, ulon
     acb_one(one);
 
     /* sum 1/(k+a)^(s+x) */
-    if (acb_is_one(a) && d <= 3)
+    if (acb_is_one(a) && d <= 3 && _acb_vec_estimate_allocated_bytes(d * N / 6, prec) < SIEVE_ALLOC_LIMIT)
         _acb_poly_powsum_one_series_sieved(sum, s, N, d, prec);
     else if (N > 50 && flint_get_num_threads() > 1)
         _acb_poly_powsum_series_naive_threaded(sum, s, a, one, N, d, prec);

arb.h

@@ -1129,6 +1129,38 @@ void _arb_atan_sum_bs_powtab(fmpz_t T, fmpz_t Q, mp_bitcnt_t * Qexp,
 
 void arb_atan_arf_bb(arb_t z, const arf_t x, slong prec);
 
+ARB_INLINE slong
+arb_allocated_bytes(const arb_t x)
+{
+    return arf_allocated_bytes(arb_midref(x)) + mag_allocated_bytes(arb_radref(x));
+}
+
+ARB_INLINE slong
+_arb_vec_allocated_bytes(arb_srcptr vec, slong len)
+{
+    slong i, size;
+
+    size = len * sizeof(arb_struct);
+
+    for (i = 0; i < len; i++)
+        size += arb_allocated_bytes(vec + i);
+
+    return size;
+}
+
+ARB_INLINE double
+_arb_vec_estimate_allocated_bytes(slong len, slong prec)
+{
+    double size;
+
+    size = len * (double) sizeof(arb_struct);
+
+    if (prec > ARF_NOPTR_LIMBS * FLINT_BITS)
+        size += len * (double) ((prec + FLINT_BITS - 1) / FLINT_BITS) * sizeof(mp_limb_t);
+
+    return size;
+}
+
 #ifdef __cplusplus
 }
 #endif

arb_mat.h

@@ -345,6 +345,11 @@ arb_mat_count_not_is_zero(const arb_mat_t mat)
     return size - arb_mat_count_is_zero(mat);
 }
 
+ARB_MAT_INLINE slong
+arb_mat_allocated_bytes(const arb_mat_t x)
+{
+    return _arb_vec_allocated_bytes(x->entries, x->r * x->c) + x->r * sizeof(arb_ptr);
+}
 
 #ifdef __cplusplus
 }

arb_poly.h

@@ -704,6 +704,12 @@ void _arb_poly_root_bound_fujiwara(mag_t bound, arb_srcptr poly, slong len);
 
 void arb_poly_root_bound_fujiwara(mag_t bound, arb_poly_t poly);
 
+ARB_POLY_INLINE slong
+arb_poly_allocated_bytes(const arb_poly_t x)
+{
+    return _arb_vec_allocated_bytes(x->coeffs, x->alloc);
+}
+
 /* Macros */
 
 

arf.h

@@ -1269,6 +1269,17 @@ int arf_sum(arf_t s, arf_srcptr terms, slong len, slong prec, arf_rnd_t rnd);
 double arf_get_d(const arf_t x, arf_rnd_t rnd);
 void arf_set_d(arf_t x, double v);
 
+ARF_INLINE slong
+arf_allocated_bytes(const arf_t x)
+{
+    slong size = fmpz_allocated_bytes(ARF_EXPREF(x));
+
+    if (ARF_HAS_PTR(x))
+        size += ARF_PTR_ALLOC(x) * sizeof(mp_limb_t);
+
+    return size;
+}
+
 #ifdef __cplusplus
 }
 #endif

doc/source/acb.rst

@@ -71,6 +71,25 @@ Memory management
 
     Clears an array of *n* initialized *acb_struct*:s.
 
+.. function:: slong acb_allocated_bytes(const acb_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(acb_struct)`` to get the size of the object as a whole.
+
+.. function:: slong _acb_vec_allocated_bytes(acb_srcptr vec, slong len)
+
+    Returns the total number of bytes allocated for this vector, i.e. the
+    space taken up by the vector itself plus the sum of the internal heap
+    allocation sizes for all its member elements.
+
+.. function:: double _acb_vec_estimate_allocated_bytes(slong len, slong prec)
+
+    Estimates the number of bytes that need to be allocated for a vector of
+    *len* elements with *prec* bits of precision, including the space for
+    internal limb data.
+    See comments for :func:`_arb_vec_estimate_allocated_bytes`.
+
 Basic manipulation
 -------------------------------------------------------------------------------
 

doc/source/acb_mat.rst

@@ -52,6 +52,11 @@ Memory management
 
     Clears the matrix, deallocating all entries.
 
+.. function:: slong acb_mat_allocated_bytes(const acb_mat_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(acb_mat_struct)`` to get the size of the object as a whole.
 
 Conversions
 -------------------------------------------------------------------------------

doc/source/acb_poly.rst

@@ -57,6 +57,12 @@ Memory management
 
     Swaps *poly1* and *poly2* efficiently.
 
+.. function:: slong acb_poly_allocated_bytes(const acb_poly_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(acb_poly_struct)`` to get the size of the object as a whole.
+
 
 Basic properties and manipulation
 -------------------------------------------------------------------------------

doc/source/arb.rst

@@ -103,6 +103,34 @@ Memory management
 
     Swaps *x* and *y* efficiently.
 
+.. function:: slong arb_allocated_bytes(const arb_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(arb_struct)`` to get the size of the object as a whole.
+
+.. function:: slong _arb_vec_allocated_bytes(arb_srcptr vec, slong len)
+
+    Returns the total number of bytes allocated for this vector, i.e. the
+    space taken up by the vector itself plus the sum of the internal heap
+    allocation sizes for all its member elements.
+
+.. function:: double _arb_vec_estimate_allocated_bytes(slong len, slong prec)
+
+    Estimates the number of bytes that need to be allocated for a vector of
+    *len* elements with *prec* bits of precision, including the space for
+    internal limb data.
+    This function returns a *double* to avoid overflow issues when both
+    *len* and *prec* are large.
+
+    This is only an approximation of the physical memory that will be used
+    by an actual vector. In practice, the space varies with the content
+    of the numbers; for example, zeros and small integers require no
+    internal heap allocation even if the precision is huge.
+    The estimate assumes that exponents will not be bignums.
+    The actual amount may also be higher or lower due to overhead in the
+    memory allocator or overcommitment by the operating system.
+
 Assignment and rounding
 -------------------------------------------------------------------------------
 

doc/source/arb_mat.rst

@@ -52,6 +52,11 @@ Memory management
 
     Clears the matrix, deallocating all entries.
 
+.. function:: slong arb_mat_allocated_bytes(const arb_mat_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(arb_mat_struct)`` to get the size of the object as a whole.
 
 Conversions
 -------------------------------------------------------------------------------

doc/source/arb_poly.rst

@@ -54,6 +54,12 @@ Memory management
 
     Strips any trailing coefficients which are identical to zero.
 
+.. function:: slong arb_poly_allocated_bytes(const arb_poly_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(arb_poly_struct)`` to get the size of the object as a whole.
+
 Basic manipulation
 -------------------------------------------------------------------------------
 

doc/source/arf.rst

@@ -131,6 +131,12 @@ Memory management
 
     Clears the variable *x*, freeing or recycling its allocated memory.
 
+.. function:: slong arf_allocated_bytes(const arf_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(arf_struct)`` to get the size of the object as a whole.
+
 Special values
 -------------------------------------------------------------------------------
 

doc/source/fmpz_extras.rst

@@ -6,6 +6,15 @@
 This module implements a few utility methods for the FLINT
 multiprecision integer type (*fmpz_t*). It is mainly intended for internal use.
 
+Memory-related methods
+-------------------------------------------------------------------------------
+
+.. function:: slong fmpz_allocated_bytes(const fmpz_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(fmpz)`` to get the size of the object as a whole.
+
 Convenience methods
 -------------------------------------------------------------------------------
 

doc/source/mag.rst

@@ -58,6 +58,12 @@ Memory management
 
     Clears a vector of length *n*.
 
+.. function:: slong mag_allocated_bytes(const mag_t x)
+
+    Returns the total number of bytes heap-allocated internally by this object.
+    The count excludes the size of the structure itself. Add
+    ``sizeof(mag_struct)`` to get the size of the object as a whole.
+
 Special values
 -------------------------------------------------------------------------------
 

fmpz_extras.h

@@ -217,6 +217,15 @@ fmpz_min(fmpz_t z, const fmpz_t x, const fmpz_t y)
 
 void fmpz_lshift_mpn(fmpz_t z, mp_srcptr d, mp_size_t dn, int sgnbit, mp_bitcnt_t shift);
 
+static __inline__ slong
+fmpz_allocated_bytes(const fmpz_t x)
+{
+    if (COEFF_IS_MPZ(*x))
+        return sizeof(__mpz_struct) + COEFF_TO_PTR(*x)->_mp_alloc * sizeof(mp_limb_t);
+    else
+        return 0;
+}
+
 #ifdef __cplusplus
 }
 #endif

mag.h

@@ -687,6 +687,11 @@ MAG_INLINE void mag_const_pi(mag_t res)
     mag_set_ui_2exp_si(res, 843314857, -28);
 }
 
+MAG_INLINE slong mag_allocated_bytes(const mag_t x)
+{
+    return fmpz_allocated_bytes(MAG_EXPREF(x));
+}
+
 #ifdef __cplusplus
 }
 #endif