Breaking: always represent inner array as UnsafeArray (#18)

* Breaking: always use UnsafeArrays inside

* fix typo

* rm `StrideArraysCore`

* rm `R` type; not necessary

* add no-op constructor

* more constructors

* fix test

* cleanup ci name

* update versions

* fix test again

* test `recursive_alloc_arrays`

* try withuot key

Author: ericphanson

.github/workflows/CI.yml

@@ -16,7 +16,7 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} - Bumper ${{ matrix.bumper }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -31,15 +31,15 @@ jobs:
           - x64
           - x86
     steps:
-      - uses: actions/checkout@v3
-      - uses: julia-actions/setup-julia@v1
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: julia-actions/cache@v1
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v3
+      - uses: codecov/codecov-action@v5
         with:
           files: lcov.info

.github/workflows/docs.yml

@@ -1,50 +1,25 @@
-name: Documentation
+name: Documenter
 on:
   push:
-    branches: [main]
-    tags: '*'
-    paths:
-      - '.github/workflows/docs.yml'
-      - 'src/**'
-      - 'docs/**'
-      - 'Project.toml'
+    branches: [main, master]
+    tags: [v*]
   pull_request:
-    types: [opened, synchronize, reopened, ready_for_review]
-    paths:
-      - '.github/workflows/docs.yml'
-      - 'src/**'
-      - 'docs/**'
-      - 'Project.toml'
 
 jobs:
-  Documentation:
-    # Run on push's or non-draft PRs
-    if: (github.event_name == 'push') || (github.event.pull_request.draft == false) || (github.event_name == 'workflow_dispatch')
+  Documenter:
+    permissions:
+      contents: write
+      statuses: write
+    name: Documentation
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: julia-actions/setup-julia@latest
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
         with:
-          version: "1.8"
-      - name: Cache artifacts
-        uses: actions/cache@v2
+          version: '1.10'                         # replace this with whatever version you need
+          show-versioninfo: true               # this causes versioninfo to be printed to the action log
+      - uses: julia-actions/cache@v2              # cache using https://github.com/julia-actions/cache
+      - uses: julia-actions/julia-buildpkg@v1  # if package requires Pkg.build()
+      - uses: julia-actions/julia-docdeploy@v1
         env:
-          cache-name: cache-artifacts
-        with:
-          path: |
-            ~/.julia/artifacts
-          key: ${{ runner.os }}-docs-${{ env.cache-name }}-${{ hashFiles('**/docs/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-docs-${{ env.cache-name }}-
-            ${{ runner.os }}-docs-
-            ${{ runner.os }}-
-      - name: Install dependencies
-        shell: julia --color=yes --project=docs/ {0}
-        run: |
-          using Pkg
-          Pkg.develop(PackageSpec(path=pwd()))
-          Pkg.instantiate()
-      - uses: julia-actions/julia-docdeploy@releases/v1
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
-          DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

Project.toml

@@ -1,29 +1,31 @@
 name = "AllocArrays"
 uuid = "5c00bae2-1499-4716-9206-27f63fd08a44"
 authors = ["Eric P. Hanson"]
-version = "0.2.0"
+version = "0.3.0"
 
 [deps]
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ConcurrentUtilities = "f0e56b4a-5159-44fe-b623-3e5288b988bb"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
+UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [compat]
 Aqua = "0.7"
 Bumper = "0.7"
 ConcurrentUtilities = "2.2.1"
+Functors = "0.5.2"
 PrecompileTools = "1.2"
 ScopedValues = "1"
-StrideArraysCore = "0.5.1"
+UnsafeArrays = "1.0.6"
 julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-StrideArraysCore = "7792a7ef-975c-4747-a70f-980b88e8d1da"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Flux", "Random", "StrideArraysCore", "Test"]
+test = ["Aqua", "Flux", "Functors", "Random", "Test"]

README.md

@@ -35,8 +35,8 @@ function basic_reduction(a)
 end
 
 arr = ones(Float64, 100_000)
-@time basic_reduction(arr) #  0.057540 seconds (209.39 k allocations: 16.487 MiB, 99.48% compilation time)
-@time basic_reduction(arr) #  0.000265 seconds (7 allocations: 2.289 MiB)
+@time basic_reduction(arr) # 0.068864 seconds (377.41 k allocations: 21.811 MiB, 99.47% compilation time)
+@time basic_reduction(arr) # 0.007774 seconds (10 allocations: 2.289 MiB, 94.21% gc time)
 
 
 function bumper_reduction!(b, a)
@@ -50,8 +50,8 @@ end
 b = BumperAllocator(2^24) # 16 MiB
 a = AllocArray(arr);
 
-@time bumper_reduction!(b, a) #  0.223121 seconds (748.98 k allocations: 2.045 GiB, 4.32% gc time, 97.59% compilation time)
-@time bumper_reduction!(b, a) #  0.000311 seconds (36 allocations: 1.172 KiB)
+@time bumper_reduction!(b, a) #  0.246476 seconds (1.05 M allocations: 52.495 MiB, 19.78% gc time, 99.65% compilation time)
+@time bumper_reduction!(b, a) #  0.000226 seconds (27 allocations: 832 bytes)
 ```
 
 We can see we brought allocations down from 2.289 MiB to ~1 KiB.
@@ -60,18 +60,49 @@ For a less-toy example, in `test/flux.jl` we test inference over a Flux model:
 
 ```julia
 # Baseline: Array
-infer!(b, predictions, model, data): 0.499735 seconds (8.05 k allocations: 306.796 MiB, 6.47% gc time)
-# Baseline: StrideArray
-stride_data = StrideArray.(data)
-infer!(b, predictions, model, stride_data): 0.364180 seconds (8.05 k allocations: 306.796 MiB, 8.32% gc time)
-# Using AllocArray:
+infer!(b, predictions, model, data): 0.163573 seconds (6.77 k allocations: 221.508 MiB, 16.42% gc time)# Using AllocArray:
 alloc_data = AllocArray.(data)
-infer!(b, predictions, model, alloc_data): 0.351953 seconds (13.60 k allocations: 3.221 MiB)
+infer!(b, predictions, model, alloc_data): 0.114566 seconds (8.72 k allocations: 777.547 KiB)
+# see "Usage with Flux" below for `recursive_alloc_arrays` and `infer!`
+aa_model = recursive_alloc_arrays(model)
+infer!(b, predictions, aa_model, alloc_data): 0.113104 seconds (7.69 k allocations: 688.047 KiB)
+# checked example (use for testing)
 checked_alloc_data = CheckedAllocArray.(data)
-infer!(b, predictions, model, checked_alloc_data): 15.522897 seconds (25.54 k allocations: 3.742 MiB)
+infer!(b, predictions, model, checked_alloc_data): 13.721077 seconds (22.54 k allocations: 1.354 MiB)
 ```
 
-We can see in this example, we got 100x less allocation (and no GC time), and similar runtime, for `AllocArray`s. We can see `CheckedAllocArrays` are far slower here.
+We can see in this example, we got 200x less allocation (and no GC time), and similar runtime, for `AllocArray`s. We also can reduce allocations a bit more with `aa_model` than `model`. We see `CheckedAllocArrays` are far slower.
+
+## Usage with Flux
+
+For reducing allocations as much as possible with Flux models, we can use `recursive_alloc_arrays` below to convert a model to use `AllocArray`s. This will convert all arrays in the model to `AllocArray`s, and will also convert any arrays in the model's parameters to `AllocArray`s. This way, any layers during the forward pass of the model which use `similar` calls based on the layer's parameters will use the bump allocator, when the forward pass is invoked within `with_allocator`.
+
+```julia
+using Functors
+
+function recursive_alloc_arrays(obj)
+    return fmap(x -> begin
+        x isa AbstractArray || return x
+        isbitstype(eltype(x)) || return x
+        return AllocArray(x)
+    end, obj; exclude=x -> x isa AbstractArray{<:Number} || x isa Function)
+end
+
+function infer!(b::BumpAllocator, predictions, model, data)
+    # Here we use a locked bumper for thread-safety, since NNlib multithreads
+    # some of it's functions. However we are sure to only deallocate outside of the threaded region. (All concurrency occurs within the `model` call itself).
+    with_allocator(b) do
+        for (idx, x) in enumerate(data)
+            predictions[idx] .= model(x)
+            reset!(b) # reset `b` after each batch
+        end
+    end
+    # don't escape bump-allocated memory!
+    return predictions
+end
+```
+
+The specific function to use here may need to be adjusted depending on the details of the model and Functors.jl. This one is tested in `test/flux.jl` on a simple model.
 
 ## Design notes
 
@@ -83,6 +114,8 @@ In particular, the caller must:
 - ...not allow memory allocated with a buffer to be live after the underlying buffer has been reset
 - ...reset their buffers before it runs out of memory
 
+In v0.3, AllocArrays started representing the inner array as an `UnsafeArray` always, rather than only when allocated via `similar`. This seems to avoid a Julia bug around missing vectorization: https://github.com/JuliaLang/julia/issues/57799.
+
 ## Safety
 
 Before using a bump allocator (`BumperAllocator`, or `UncheckedBumperAllocator`) it is recommended the user read the [Bumper.jl README](https://github.com/MasonProtter/Bumper.jl#bumperjl) to understand how it works and what the limitations are.

julia_57799.jl

@@ -0,0 +1,48 @@
+using AllocArrays, BenchmarkTools, UnsafeArrays
+
+function mycopy(dest, src, iter)
+    # precondition: dest and src do not alias
+    # precondition: the iterators of dest and src are equal
+    for i in iter
+        @inbounds dest[i] = src[i]
+    end
+    return dest
+end
+
+b = BumperAllocator(2^30);
+arr = rand(Float32, 50000000);
+arr2 = similar(arr);
+a = AllocArray(arr);
+
+println("With mycopy(arr2, arr, eachindex(arr2)):")
+display(@benchmark mycopy(arr2, arr, eachindex(arr2)))
+
+println("With a = AllocArray(a)")
+display(@benchmark with_allocator(b) do
+    c = similar(a)
+    mycopy(c, a, eachindex(c))
+end setup=(reset!(b)) evals=1)
+
+# julia> include("bench.jl")
+# With mycopy(arr2, arr, eachindex(arr2)):
+# BenchmarkTools.Trial: 627 samples with 1 evaluation per sample.
+#  Range (min … max):  6.662 ms … 38.670 ms  ┊ GC (min … max): 0.00% … 0.00%
+#  Time  (median):     7.151 ms              ┊ GC (median):    0.00%
+#  Time  (mean ± σ):   7.950 ms ±  2.168 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%
+
+#   █▇                                                          
+#   ███▅▄▄▃▄▄▃▃▃▃▃▃▃▃▄▆▄▃▃▃▂▃▃▃▃▃▃▂▂▁▂▁▂▂▁▁▂▁▂▁▁▂▂▁▁▁▁▁▁▂▁▁▁▁▂ ▃
+#   6.66 ms        Histogram: frequency by time          14 ms <
+
+#  Memory estimate: 16 bytes, allocs estimate: 1.
+# With a = AllocArray(a)
+# BenchmarkTools.Trial: 698 samples with 1 evaluation per sample.
+#  Range (min … max):  6.659 ms …  17.031 ms  ┊ GC (min … max): 0.00% … 0.00%
+#  Time  (median):     6.789 ms               ┊ GC (median):    0.00%
+#  Time  (mean ± σ):   7.166 ms ± 941.846 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%
+
+#   █▇▄▄▂ ▁   ▁▁    ▁                                            
+#   ██████████████▇▇██▆▇▇▆▇█▅▇▇▇▅▆▆▁▄▁▄▆▁▄▁▆▄▅▅▁▆▅▄▆▅▄▁▆▅▅▁▁▄▅▄ ▇
+#   6.66 ms      Histogram: log(frequency) by time      10.1 ms <
+
+#  Memory estimate: 384 bytes, allocs estimate: 13.

src/AllocArray.jl

@@ -1,21 +1,55 @@
 using Base: Base.Broadcast
 using Base: Dims
 using Base.Broadcast: Broadcasted, ArrayStyle
+using UnsafeArrays: UnsafeArray
 
 """
-    struct AllocArray{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
-        arr::A
+    struct AllocArray{T,N} <: DenseArray{T,N}
+        arr::UnsafeArray{T,N}
+        gcref::Any
     end
 
+    AllocArray(arr::AbstractArray)
+
 Wrapper type which forwards most array methods to the inner array `arr`,
 but dispatches `similar` to special allocation methods.
+
+The inner array `arr` is always represented as an `UnsafeArray`,
+but the field `gcref` may hold a materialized `AbstractArray` corresponding
+to the same data, which is held to preserve a reference to the data to prevent GC.
+This field is never accessed or used, so the `::Any` type does not affect type stability
+of code using AllocArrays.
+
+Use the constructor `AllocArray(arr)` to construct an `AllocArray`. Note that `arr`
+must be able to be represented as an `UnsafeArray`, meaning it must be a bits-type
+and have a pointer. To support `UnitRange` and similar, `collect` it first.
+
+Typically this constructor is only used at the entrypoint of a larger set of code
+which is expected to use `similar` based on this input for further allocations.
+When inside a `with_allocator` block, `similar` can be dispatched to a
+(dynamically-scoped) bump allocator.
 """
-struct AllocArray{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
-    arr::A
+struct AllocArray{T,N} <: DenseArray{T,N}
+    arr::UnsafeArray{T,N}
+    gcref::Any
+
+    AllocArray(gcref::AbstractArray{T,N}) where {T,N} = AllocArray{T,N}(gcref)
+    function AllocArray{T,N}(gcref::AbstractArray{T,N}) where {T,N}
+        arr = UnsafeArray(pointer(gcref), size(gcref))
+        return new{eltype(arr),ndims(arr)}(arr, gcref)
+    end
+
+    # already allocated with Bumper, no gcref needed
+    AllocArray(arr::UnsafeArray{T,N}) where {T,N} = AllocArray{T,N}(arr)
+    AllocArray{T,N}(arr::UnsafeArray{T,N}) where {T,N} = new{T,N}(arr, nothing)
+
+
+    AllocArray(a::AllocArray) = a
+    AllocArray{T,N}(a::AllocArray{T,N}) where {T,N} = a
 end
 
-AllocMatrix{T,A<:AbstractMatrix{T}} = AllocArray{T,2,A}
-AllocVector{T,A<:AbstractVector{T}} = AllocArray{T,1,A}
+AllocMatrix{T} = AllocArray{T,2}
+AllocVector{T} = AllocArray{T,1}
 
 @inline Base.parent(a::AllocArray) = getfield(a, :arr)
 
@@ -41,7 +75,7 @@ end
 
 Base.size(a::AllocArray) = size(getfield(a, :arr))
 
-Base.IndexStyle(::Type{<:AllocArray{T,N,Arr}}) where {T,N,Arr} = Base.IndexStyle(Arr)
+Base.IndexStyle(::Type{<:AllocArray{T,N}}) where {T,N} = Base.IndexStyle(UnsafeArray{T,N})
 
 # used only by broadcasting?
 function Base.similar(::Type{<:AllocArray{T}}, dims::Dims) where {T}
@@ -79,7 +113,7 @@ function Base.unsafe_convert(::Type{Ptr{T}}, a::AllocArray) where {T}
     return Base.unsafe_convert(Ptr{T}, getfield(a, :arr))
 end
 
-Base.elsize(::Type{<:AllocArray{T,N,Arr}}) where {T,N,Arr} = Base.elsize(Arr)
+Base.elsize(::Type{<:AllocArray{T,N}}) where {T,N} = Base.elsize(UnsafeArray{T,N})
 
 Base.strides(a::AllocArray) = strides(getfield(a, :arr))
 

test/checked.jl

@@ -77,5 +77,5 @@ end
 
     @test sprint(showerror, InvalidMemoryException()) == "InvalidMemoryException: Array accessed after its memory has been deallocated."
 
-    @test Base.parent(c) === inner
+    @test c.alloc_array.gcref === inner
 end