JuliaAstro / FITSIO.jl

Flexible Image Transport System (FITS) file support for Julia
http://juliaastro.org/FITSIO.jl/
MIT License
55 stars 29 forks source link

Use ndims in size to reduce allocations #167

Closed jishnub closed 2 years ago

jishnub commented 2 years ago

This PR relies on https://github.com/JuliaAstro/CFITSIO.jl/pull/8, so tests will only pass with that PR merged and a version tagged.

On master and using the latest tagged CFITSIO:

julia> filename = tempname();

julia> FITSIO.fitswrite(filename, collect(reshape(1:10_000, 100, 100)));

julia> using BenchmarkTools

julia> g(f::FITS) = size(f[1]);  # runtime dispatch

julia> h(f::FITS) = size(f[1]::ImageHDU{Int,2}); # compile-time dispatch

julia> f = FITS(filename, "r");

julia> @btime g($f);
  161.138 ns (2 allocations: 128 bytes)

julia> @btime h($f);
  127.562 ns (1 allocation: 96 bytes)

After this PR:

julia> @btime g($f);
  97.467 ns (1 allocation: 32 bytes)

julia> @btime h($f);
  67.117 ns (0 allocations: 0 bytes)

With this, in-place read becomes allocation-free:

On master and using the latest tagged CFITSIO:

julia> hdu = f[1];

julia> @btime read!($hdu, $m);
  23.605 μs (3 allocations: 288 bytes)

After this PR:

julia> @btime read!($hdu, $m);
  23.086 μs (0 allocations: 0 bytes)

Even with non-inferred code and runtime dispatch, performance seems identical and there are no allocations:

julia> read1!(f, m) = read!(f[1], m);

julia> @btime read1!($f, $m);
  23.320 μs (0 allocations: 0 bytes)

julia> @code_warntype read1!(f, m)
Variables
  #self#::Core.Const(read1!)
  f::FITS
  m::Matrix{Int64}

Body::Array{Int64, N} where N
1 ─ %1 = Base.getindex(f, 1)::HDU
│   %2 = Main.read!(%1, m)::Array{Int64, N} where N
└──      return %2

One may of course help with the inference if the eltype and ndim is known:

julia> read1!(f, m) = read!(f[1]::ImageHDU{Int,2}, m);

julia> @code_warntype read1!(f, m)
Variables
  #self#::Core.Const(read1!)
  f::FITS
  m::Matrix{Int64}

Body::Matrix{Int64}
1 ─ %1 = Base.getindex(f, 1)::HDU
│   %2 = Core.apply_type(Main.ImageHDU, Main.Int, 2)::Core.Const(ImageHDU{Int64, 2})
│   %3 = Core.typeassert(%1, %2)::ImageHDU{Int64, 2}
│   %4 = Main.read!(%3, m)::Matrix{Int64}
└──      return %4

julia> @btime read1!($f, $m);
  23.328 μs (0 allocations: 0 bytes)
jishnub commented 2 years ago

After https://github.com/JuliaAstro/FITSIO.jl/pull/167/commits/73408ee2471443b3a4493b095f44d31a89e4cfd7 there are no allocations in write:

julia> f = FITS(tempname(), "w");

julia> a = Float64[1 3; 2 4];

julia> @btime write($f, $a);
  25.898 μs (0 allocations: 0 bytes)

julia> hdu = f[1];

julia> @btime write($hdu, $a);
  235.786 ns (0 allocations: 0 bytes)
codecov[bot] commented 2 years ago

Codecov Report

Merging #167 (b21e4a8) into master (3ec2708) will not change coverage. The diff coverage is 100.00%.

:exclamation: Current head b21e4a8 differs from pull request most recent head 2a24db0. Consider uploading reports for the commit 2a24db0 to get more accurate results Impacted file tree graph

@@           Coverage Diff           @@
##           master     #167   +/-   ##
=======================================
  Coverage   91.20%   91.20%           
=======================================
  Files           5        5           
  Lines         614      614           
=======================================
  Hits          560      560           
  Misses         54       54           
Impacted Files Coverage Δ
src/image.jl 94.89% <100.00%> (ø)

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 3ec2708...2a24db0. Read the comment docs.

jishnub commented 2 years ago

Test failure in building the documentation is likely due to https://github.com/JuliaRegistries/General/issues/16777. I've seen this happen if the compat version of a dependency is bumped.