add frozen function (#8)

* add frozen function

* fix method

* add test & pretty print

* fix test

* fix type stablility problem

* oops

* this might be better?

* Update src/replaying.jl

Co-Authored-By: Lyndon White <oxinabox@ucc.asn.au>

* rename -> preallocate

* bump version

* fix test

* export preallocation

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Project.toml

@@ -1,7 +1,7 @@
 name = "AutoPreallocation"
 uuid = "e7028de2-df94-4053-9fdc-99272086b8d4"
 authors = ["Lyndon White"]
-version = "0.1.1"
+version = "0.2.0"
 
 [deps]
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"

src/AutoPreallocation.jl

@@ -2,7 +2,7 @@ module AutoPreallocation
 using Cassette
 using LinearAlgebra: LinearAlgebra
 
-export avoid_allocations, record_allocations, freeze, reinitialize!, @no_prealloc
+export avoid_allocations, record_allocations, freeze, preallocate, reinitialize!, @no_prealloc
 
 include("record_types.jl")
 include("recording.jl")

src/record_types.jl

@@ -21,7 +21,7 @@ Freezes an allocation `record`, to not allow new allocations to be added.
 This converts the memory of the what happenned to be stored as a `Tuple`.
 This could give better performance during replay; or it could give worse.
 """
-freeze(record) = FrozenAllocationRecord(record)
+freeze(record::AbstractAllocationRecord) = FrozenAllocationRecord(record)
 
 """
     reinitialize!(record)

src/replaying.jl

@@ -44,3 +44,36 @@ function avoid_allocations(record, f, args...; kwargs...)
     ctx = new_replay_ctx(record)
     return Cassette.overdub(ctx, f, args...; kwargs...)
 end
+
+struct PreallocatedFunction{F}
+    f::F
+    ctx::Dict{Tuple, ReplayCtx}  # maps from argument types to the ReplayCtx
+    PreallocatedFunction(f) = new{typeof(f)}(f, Dict{Tuple, ReplayCtx}())
+end
+
+@generated function (f::PreallocatedFunction)(xs...)
+    return quote
+        if haskey(f.ctx, $xs)
+            ctx = f.ctx[$xs]
+            # step = ctx.metadata.step::Ref{Int}
+            ctx.metadata.step[] = 1
+            return Cassette.overdub(ctx, f.f, xs...)
+        else
+            x, record = record_allocations(f.f, xs...)
+            ctx = AutoPreallocation.new_replay_ctx(record)
+            f.ctx[$xs] = ctx
+            return x
+        end
+    end
+end
+
+"""
+    preallocate(f)
+
+Preallocate a function. This will preallocate the allocation behaviour of the function by creating
+a [`PreallocatedFunction`](@ref). This function will record all the allocations at the first run,
+then in the following run, it will not allocate anymore.
+"""
+preallocate(f) = PreallocatedFunction(f)
+
+Base.show(io::IO, f::PreallocatedFunction) = print(io, "preallocate(", f.f, ")")

test/integration_tests.jl

@@ -19,6 +19,9 @@ end
     val, record = record_allocations(f_matmul)
     # NOTE: (@Roger-luo) not sure why this is 256 on my machine
     @test (@ballocated avoid_allocations($record, f_matmul)) <= 352
+
+    f = preallocate(f_matmul)
+    @test (@ballocated $f()) <= 352
 end
 
 @testset "noprealloc example" begin