add no_prealloc

Project.toml

@@ -4,15 +4,16 @@ authors = ["Lyndon White"]
 version = "0.1.0"
 
 [deps]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 
 [compat]
+BenchmarkTools = "0.5"
 Cassette = "0.2, 0.3"
 julia = "1"
 
 [extras]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "BenchmarkTools"]

README.md

@@ -13,6 +13,20 @@ The process to use AutoPreallocation.jl is two step:
 1. Generate a *record* of all allocations
 2. Use that *record* to avoid allocations when the function is called
 
+## Functions;
+### `value, record = record_allocations(f, args...; kwargs...)`
+Record the allocations from computing `f(args...; kwargs...)`.
+The returned `record` object is what holds these allocations, to be reused.
+
+### `value = avoid_allocations(record, f, args...; kwargs...)`
+Compute `f(args...; kwargs...)`, while making use of the preallocations stored in the `record`.
+
+### `@no_prealloc(expr)`
+This macro is used within code that one might use AutoPreallocation on.
+It is used to mark a section to have all its allocations ignored.
+They will neither be recorded, nor avoided.
+This is useful if internally to the function its allocations are e.g. not nondetermanistic.
+
 ## Example:
 ```julia
 julia> using AutoPreallocation, BenchmarkTools
@@ -83,6 +97,8 @@ julia> # If the type is the same and only size differs AutoPreallocation right n
 One way to deal with this is to keep track of which parameters change the allocation pattern,
 and then declare one _record_ for each of them.
 
+If one part of your function has nondetermanistic internal allocations you can mark that section to be ignored by wrapping it in `@no_prealloc`.
+
 ### If a function resizes any array that it allocates you need to call `reinitialize!`
 The allocated memory is reuses.
 Which means if you call an operation like `push!` or `append!` that resizes it,

src/AutoPreallocation.jl

@@ -1,10 +1,11 @@
 module AutoPreallocation
 using Cassette
 
-export avoid_allocations, record_allocations, freeze, reinitialize!
+export avoid_allocations, record_allocations, freeze, reinitialize!, @no_prealloc
 
 include("record_types.jl")
 include("recording.jl")
 include("replaying.jl")
+include("no_prealloc.jl")
 
 end # module

src/no_prealloc.jl

@@ -0,0 +1,22 @@
+"""
+    @no_prealloc(expr)
+
+Anything inside a `@no_prealloc` will be ignored by AutoPreallocation.jl.
+This applies recursively to functions called by functions called etc within the
+`@no_prealloc`.
+If it allocates then that allocation will always occur -- it won't be avoided by
+auto-preallocation.
+
+This can be used if there is a part of a function that may allocation a unknowable
+number of intermediary variables to get to the final answer -- and which thus is not
+suitable for use with AutoPreallocation.
+"""
+macro no_prealloc(expr)
+    return :(no_prealloc(()->$(esc(expr))))
+end
+
+@inline no_prealloc(f) = f()
+
+# Prevent Cassette from recursing into body of f
+@inline Cassette.overdub(ctx::RecordingCtx, ::typeof(no_prealloc), f) = f()
+@inline Cassette.overdub(ctx::ReplayCtx, ::typeof(no_prealloc), f) = f()

src/recording.jl

@@ -16,7 +16,6 @@ record_alloc!(ctx::RecordingCtx, val) = record_alloc!(ctx.metadata, val)
     return ret
 end
 
-
 function record_allocations(f, args...; kwargs...)
     ctx = new_recording_ctx()
     value = Cassette.overdub(ctx, f, args...; kwargs...)

test/integration_tests.jl

@@ -6,11 +6,12 @@ using Test
 f_ones() = ones(64)
 f_matmul() = ones(32,64) * ones(64, 2)
 
+f_matmul_noprealloc() = ones(32,64) * @no_prealloc(ones(64, 2))
 
 @testset "ones example" begin
     @assert (@ballocated f_ones()) === 624
     val, record = record_allocations(f_ones)
-    @test (@ballocated avoid_allocations($record, f_ones)) == 64
+    @test (@ballocated avoid_allocations($record, f_ones)) <= 64
 end
 
 @testset "matmul example" begin
@@ -19,3 +20,11 @@ end
     # NOTE: (@Roger-luo) not sure why this is 256 on my machine
     @test (@ballocated avoid_allocations($record, f_matmul)) <= 352
 end
+
+@testset "noprealloc example" begin
+    @assert (@ballocated f_matmul_noprealloc()) === 18_304
+    val, record = record_allocations(f_matmul_noprealloc)
+    @test length(record.allocations) == 2
+    @test record.initial_sizes == [(32, 64), (32, 2)]
+    @test (@ballocated avoid_allocations($record, f_matmul_noprealloc)) <= 1520
+end