Skip to content

Commit c0c5ec8

Browse files
committed
Remove assume_effects and use StableTasks.jl in apply[reduce]
StableTasks.jl allows type stable tasks (previously type unstable), so it's significantly easier to multithread. This commit also shows how you could implement another "execution engine" a la what JuliaFolds2 does (maybe work stealing, etc) if you wanted. In benchmarking, StableTasks does decrease the number of allocations which can only be a good thing.
1 parent 458ee35 commit c0c5ec8

File tree

4 files changed

+20
-14
lines changed

4 files changed

+20
-14
lines changed

GeometryOpsCore/Project.toml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@ version = "0.1.2"
66
[deps]
77
DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
88
GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
9+
StableTasks = "91464d47-22a1-43fe-8b7f-2d57ee82463f"
910
Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
1011

1112
[compat]
12-
julia = "1.9"
1313
DataAPI = "1"
1414
GeoInterface = "1.2"
15+
StableTasks = "0.1.5"
1516
Tables = "1"
17+
julia = "1.9"

GeometryOpsCore/src/GeometryOpsCore.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ end
1616

1717
using Tables
1818
using DataAPI
19+
import StableTasks
1920

2021
include("keyword_docs.jl")
2122
include("types.jl")

GeometryOpsCore/src/apply.jl

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,18 @@ end
319319

320320
using Base.Threads: nthreads, @threads, @spawn
321321

322+
#=
323+
Here we used to use the compiler directive `@assume_effects :foldable` to force the compiler
324+
to lookup through the closure. This alone makes e.g. `flip` 2.5x faster!
325+
326+
But it caused inference to fail, so we've removed it. No effect on runtime so far as we can tell,
327+
at least in Julia 1.11.
328+
=#
329+
@inline function _maptasks(f::F, taskrange, threaded::False)::Vector where F
330+
map(f, taskrange)
331+
end
332+
333+
322334
# Threading utility, modified Mason Protters threading PSA
323335
# run `f` over ntasks, where f receives an AbstractArray/range
324336
# of linear indices
@@ -333,22 +345,12 @@ using Base.Threads: nthreads, @threads, @spawn
333345
# Map over the chunks
334346
tasks = map(task_chunks) do chunk
335347
# Spawn a task to process this chunk
336-
@spawn begin
348+
StableTasks.@spawn begin
337349
# Where we map `f` over the chunk indices
338350
map(f, chunk)
339351
end
340352
end
341353

342354
# Finally we join the results into a new vector
343355
return mapreduce(fetch, vcat, tasks)
344-
end
345-
#=
346-
Here we used to use the compiler directive `@assume_effects :foldable` to force the compiler
347-
to lookup through the closure. This alone makes e.g. `flip` 2.5x faster!
348-
349-
But it caused inference to fail, so we've removed it. No effect on runtime so far as we can tell,
350-
at least in Julia 1.11.
351-
=#
352-
@inline function _maptasks(f::F, taskrange, threaded::False)::Vector where F
353-
map(f, taskrange)
354356
end

GeometryOpsCore/src/applyreduce.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ import Base.Threads: nthreads, @threads, @spawn
135135
# Map over the chunks
136136
tasks = map(task_chunks) do chunk
137137
# Spawn a task to process this chunk
138-
@spawn begin
138+
StableTasks.@spawn begin
139139
# Where we map `f` over the chunk indices
140140
mapreduce(f, op, chunk; init)
141141
end
@@ -144,6 +144,7 @@ import Base.Threads: nthreads, @threads, @spawn
144144
# Finally we join the results into a new vector
145145
return mapreduce(fetch, op, tasks; init)
146146
end
147-
Base.@assume_effects :foldable function _mapreducetasks(f::F, op, taskrange, threaded::False; init) where F
147+
148+
function _mapreducetasks(f::F, op, taskrange, threaded::False; init) where F
148149
mapreduce(f, op, taskrange; init)
149150
end

0 commit comments

Comments
 (0)