Speeding up my logsumexp function

General Usage Performance
, , ,
16

I’m trying to use Tullio with sparse arrays.

using Tullio, LoopVectorization, SparseArrays

function logsumexp(mat; dims=1)
           max_ = maximum(mat, dims=1)
           @tullio exp_mat[i,j] := exp(mat[i,j] - max_[1,j]) - (mat[i,j] == max_[1,j])
           sum_exp_ = sum(exp_mat, dims=dims)
           @tullio out[i,j] := log1p(sum_exp_[i,j]) + max_[i,j]
       end

A = sprand(1500, 10000, 0.02);
logsumexp(A)

When I run the above code I get signal (6): Aborted

This is the full stack trace.

double free or corruption (!prev)

signal (6): Aborted
in expression starting at REPL[4]:1
gsignal at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
abort at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x7fe8d09d8906)
unknown function (ip: 0x7fe8d09df979)
cfree at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
jl_realloc_aligned at /buildworker/worker/package_linux64/build/src/gc.c:249 [inlined]
gc_managed_realloc_ at /buildworker/worker/package_linux64/build/src/gc.c:3369 [inlined]
jl_gc_managed_realloc at /buildworker/worker/package_linux64/build/src/gc.c:3386
array_resize_buffer at /buildworker/worker/package_linux64/build/src/array.c:660
jl_array_grow_at_beg at /buildworker/worker/package_linux64/build/src/array.c:785 [inlined]
jl_array_grow_at at /buildworker/worker/package_linux64/build/src/array.c:929
_growat! at ./array.jl:873 [inlined]
insert! at ./array.jl:1215 [inlined]
_insert! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.4/SparseArrays/src/sparsematrix.jl:2485
_setindex_scalar! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.4/SparseArrays/src/sparsematrix.jl:2473
setindex! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.4/SparseArrays/src/sparsematrix.jl:2447 [inlined]
π’œπ’Έπ“‰! at /home/s/.julia/packages/Tullio/HGzih/src/macro.jl:797
unknown function (ip: 0x7fe87f26475e)
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2145 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2323
block_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:104
block_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:108
block_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:107
block_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:107
block_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:101 [inlined]
thread_halves at /home/s/.julia/packages/Tullio/HGzih/src/threads.jl:91
#184 at ./threadingconstructs.jl:126
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2145 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2323
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1700 [inlined]
start_task at /buildworker/worker/package_linux64/build/src/task.c:687
unknown function (ip: (nil))
Allocations: 55469269 (Pool: 55453222; Big: 16047); GC: 56
Aborted (core dumped)

The code runs fine for 100x100 sparse arrays and 1500x10000 normal arrays. My machine has 45 GB free RAM. What is the problem here?