From 836e373377034c326513faeac04f34788ef63b75 Mon Sep 17 00:00:00 2001 From: Rogerluo Date: Mon, 9 Mar 2020 14:36:10 -0400 Subject: [PATCH] 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 * rename -> preallocate * bump version * fix test * export preallocation Co-authored-by: Lyndon White --- Project.toml | 2 +- src/AutoPreallocation.jl | 2 +- src/record_types.jl | 2 +- src/replaying.jl | 33 +++++++++++++++++++++++++++++++++ test/integration_tests.jl | 3 +++ 5 files changed, 39 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index f9ae2a0..14c4a92 100644 --- a/Project.toml +++ b/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" diff --git a/src/AutoPreallocation.jl b/src/AutoPreallocation.jl index 49b431a..6de00d5 100644 --- a/src/AutoPreallocation.jl +++ b/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") diff --git a/src/record_types.jl b/src/record_types.jl index 82d2f2a..bc0c483 100644 --- a/src/record_types.jl +++ b/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) diff --git a/src/replaying.jl b/src/replaying.jl index 8b96f43..8a114c4 100644 --- a/src/replaying.jl +++ b/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, ")") diff --git a/test/integration_tests.jl b/test/integration_tests.jl index 54f7122..bb235ba 100644 --- a/test/integration_tests.jl +++ b/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