[ANN] ObjectPools.jl

Package Announcements
1

ObjectPools.jl is a small package that enables convenient semi-manual memory management. It has proven to be very useful in my own work hence I’ve spun it out into a separate package which I just registered in General. I’d be interested in feedback - is this even a viable approach to these problems? What might alternatives be? I’m obviously also interested in issues and PRs!

The package tries to solve the following problems:

  • conveniently constructing temporary work arrays or return arrays at runtime without allocation
  • the element types and sizes of the arrays must be flexible so they can be changed at will during runtime

My own use-case is for constructing highly specialized “small” ML models where BLAS3 cannot be (easily) exploited and memory management can become an issue.

The package offers several array types that try to solve the problems outlined above, with slightly different performance characteristics. Here is a benchmark showing the capability (v0.2.0, M1, Julia 1.8.4)

Evaluate a Chebyshev basis of length nB, for nX inputs
Timings are min/man in ns
┌─────────────────────────────┬───────────┬───────────┬───────────┬────────────┐
│                     nB / nX │   10 / 16 │   10 / 32 │   30 / 16 │    30 / 32 │
├─────────────────────────────┼───────────┼───────────┼───────────┼────────────┤
│                       Array │ 147 / 259 │ 163 / 566 │ 377 / 876 │ 412 / 1286 │
│               pre-allocated │   89 / 97 │   65 / 66 │ 253 / 263 │  213 / 214 │
│                   FlexArray │  95 / 100 │   63 / 63 │ 264 / 273 │  207 / 213 │
│        ArrayPool(FlexArray) │   91 / 93 │   68 / 70 │ 264 / 271 │  216 / 223 │
│              FlexArrayCache │ 104 / 106 │   88 / 94 │ 280 / 287 │  270 / 283 │
│   ArrayPool(FlexArrayCache) │ 111 / 112 │   93 / 98 │ 285 / 292 │  275 / 287 │
│            TSafe(FlexArray) │   87 / 89 │   67 / 68 │ 262 / 269 │  212 / 219 │
│ TSafe(ArrayPool(FlexArray)) │   96 / 97 │   74 / 77 │ 262 / 271 │  224 / 232 │
└─────────────────────────────┴───────────┴───────────┴───────────┴────────────┘

I finish with a basic example how I use this package in my own work. This is not meant to be run, just indicative.

struct Ylms
   L::Int 
   tmpP::FlexArray
   outY::FlexArrayCache
end

Ylms(L::Integer) = Ylms(L, FlexArray(), FlexCacheArray())

function (ylms::Ylms)(r::SVector{3, T}) where {T <: Real}
   # not shown: lenP, lenY, eval_alp!, eval_ylm!
   L = ylms.L
   P = acquire!(ylms.tmpP, (lenP(L),), T)
   eval_alp!(P, r)
   Y = acquire!(ylms.outY, (lenY(L),), Complex{T})
   eval_ylm!(Y, P, r)
   return Y 
end

ylms = Ylms(L)

for i = 1:niter
   r = @SVector randn(3)  # generate an input somehow 
   Y = ylms(r)            # evaluate the Ylms 
   # ....                   do something with Y 
   release!(Y)            # return it to the pool
end

In this case the return type is not hard to infer, but I may then want to forward-diff it, and suddenly I need to work with Duals. In other much more complicated cases I’ve encountered it becomes very tedious (never really impossible though) to predict the array element types at the time when I’m constructing the objects. With FlexArray I don’t have to worry about this anymore at all.

5 Likes