Repository: JuliaFolds2/OhMyThreads.jl (julia)

Issue:
Rename `@init` to `@tasklocal` or `@local`, make `@set init` set the `init` kwarg
I just realized that in `@tasks`, we 

1) have no way to set the `init` keyword argument
2) the obvious way to do that would be with `@init` or `@set init = ...`, but that seems unnececarily confusing given that we also have `@init` for task local values.

Given that, I think it'd be good to do a renaming here before too many people start relying on `@tasks`. Any thoughts?

Pull request:
Closes https://github.com/JuliaFolds2/OhMyThreads.jl/issues/64

* [x] `@init` → `@local`
* [x] `@set init = ...`
* [x] drop deprecated `SpawnAllScheduler` (since we're preparing a breaking release here anyways)

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 48e65fc1..99cc506d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,13 @@
 OhMyThreads.jl Changelog
 =========================
 
+Version 0.5.0
+-------------
+
+- ![Feature][badge-feature] `@set init = ...` may now be used to specify an initial value for a reduction (only has an effect in conjuction with `@set reducer=...` and triggers a warning otherwise).
+- ![BREAKING][badge-breaking] Within a `@tasks` block, task-local values must from now on be defined via `@local` instead of `@init` (renamed).
+- ![BREAKING][badge-breaking] The (already deprecated) `SpawnAllScheduler` has been dropped.
+
 Version 0.4.6
 -------------
 
diff --git a/docs/src/literate/tls/tls.jl b/docs/src/literate/tls/tls.jl
index 58eb6344..dc12f0d7 100644
--- a/docs/src/literate/tls/tls.jl
+++ b/docs/src/literate/tls/tls.jl
@@ -165,7 +165,7 @@ res ≈ res_tlv
 # (instead of $O(\textrm{iterations})$) allocations.
 #
 # Note that if you use our `@tasks` macro API, there is built-in support for task-local
-# values via `@init`.
+# values via `@local`.
 #
 
 using OhMyThreads: @tasks
@@ -174,7 +174,7 @@ function matmulsums_tlv_macro(As, Bs; kwargs...)
     N = size(first(As), 1)
     @tasks for i in eachindex(As,Bs)
         @set collect=true
-        @init C::Matrix{Float64} = Matrix{Float64}(undef, N, N)
+        @local C::Matrix{Float64} = Matrix{Float64}(undef, N, N)
         mul!(C, As[i], Bs[i])
         sum(C)
     end
@@ -183,7 +183,7 @@ end
 res_tlv_macro = matmulsums_tlv_macro(As, Bs)
 res ≈ res_tlv_macro
 
-# Here, `@init` simply expands to the explicit pattern around `TaskLocalValue` above.
+# Here, `@local` simply expands to the explicit pattern around `TaskLocalValue` above.
 #
 #
 # ### Benchmark
diff --git a/docs/src/literate/tls/tls.md b/docs/src/literate/tls/tls.md
index 0392e993..6ea93a4c 100644
--- a/docs/src/literate/tls/tls.md
+++ b/docs/src/literate/tls/tls.md
@@ -215,7 +215,7 @@ This solves our issues above and leads to $O(\textrm{parallel tasks})$
 (instead of $O(\textrm{iterations})$) allocations.
 
 Note that if you use our `@tasks` macro API, there is built-in support for task-local
-values via `@init`.
+values via `@local`.
 
 ````julia
 using OhMyThreads: @tasks
@@ -224,7 +224,7 @@ function matmulsums_tlv_macro(As, Bs; kwargs...)
     N = size(first(As), 1)
     @tasks for i in eachindex(As,Bs)
         @set collect=true
-        @init C::Matrix{Float64} = Matrix{Float64}(undef, N, N)
+        @local C::Matrix{Float64} = Matrix{Float64}(undef, N, N)
         mul!(C, As[i], Bs[i])
         sum(C)
     end
@@ -238,7 +238,7 @@ res ≈ res_tlv_macro
 true
 ````
 
-Here, `@init` simply expands to the explicit pattern around `TaskLocalValue` above.
+Here, `@local` simply expands to the explicit pattern around `TaskLocalValue` above.
 
 
 ### Benchmark
diff --git a/docs/src/refs/api.md b/docs/src/refs/api.md
index f2393efe..a7805eab 100644
--- a/docs/src/refs/api.md
+++ b/docs/src/refs/api.md
@@ -13,7 +13,7 @@ Order   = [:function, :macro, :type]
 ```@docs
 @tasks
 @set
-@init
+@local
 ```
 
 ### Functions
diff --git a/src/OhMyThreads.jl b/src/OhMyThreads.jl
index a00d3335..50fe693c 100644
--- a/src/OhMyThreads.jl
+++ b/src/OhMyThreads.jl
@@ -1,7 +1,7 @@
 module OhMyThreads
 
 using StableTasks: StableTasks
-for mac ∈ Symbol.(["@spawn", "@spawnat", "@fetch", "@fetchfrom"])
+for mac in Symbol.(["@spawn", "@spawnat", "@fetch", "@fetchfrom"])
     @eval const $mac = getproperty(StableTasks, $(QuoteNode(mac)))
 end
 
@@ -16,12 +16,11 @@ include("macros.jl")
 
 include("tools.jl")
 include("schedulers.jl")
-using .Schedulers: Scheduler,
-    DynamicScheduler, StaticScheduler, GreedyScheduler, SpawnAllScheduler
+using .Schedulers: Scheduler, DynamicScheduler, StaticScheduler, GreedyScheduler
 include("implementation.jl")
 
-export @tasks, @set, @init
+export @tasks, @set, @local
 export treduce, tmapreduce, treducemap, tmap, tmap!, tforeach, tcollect
-export Scheduler, DynamicScheduler, StaticScheduler, GreedyScheduler, SpawnAllScheduler
+export Scheduler, DynamicScheduler, StaticScheduler, GreedyScheduler
 
 end # module OhMyThreads
diff --git a/src/macro_impl.jl b/src/macro_impl.jl
index 995d24be..246000d7 100644
--- a/src/macro_impl.jl
+++ b/src/macro_impl.jl
@@ -17,8 +17,8 @@ function tasks_macro(forex)
 
     settings = Settings()
 
-    inits_before, init_inner = _maybe_handle_init_block!(forbody.args)
-    _maybe_handle_set_block!(settings, forbody.args)
+    locals_before, local_inner = _maybe_handle_atlocal_block!(forbody.args)
+    _maybe_handle_atset_block!(settings, forbody.args)
 
     forbody = esc(forbody)
     itrng = esc(itrng)
@@ -26,24 +26,37 @@ function tasks_macro(forex)
 
     # @show settings
     q = if !isnothing(settings.reducer)
-        quote
-            tmapreduce(
-                $(settings.reducer), $(itrng); scheduler = $(settings.scheduler)) do $(itvar)
-                $(init_inner)
-                $(forbody)
+        if !isnothing(settings.init)
+            quote
+                tmapreduce(
+                    $(settings.reducer), $(itrng); scheduler = $(settings.scheduler),
+                    init = $(settings.init)) do $(itvar)
+                    $(local_inner)
+                    $(forbody)
+                end
+            end
+        else
+            quote
+                tmapreduce(
+                    $(settings.reducer), $(itrng); scheduler = $(settings.scheduler)) do $(itvar)
+                    $(local_inner)
+                    $(forbody)
+                end
             end
         end
     elseif settings.collect
+        maybe_warn_useless_init(settings)
         quote
             tmap($(itrng); scheduler = $(settings.scheduler)) do $(itvar)
-                $(init_inner)
+                $(local_inner)
                 $(forbody)
             end
         end
     else
+        maybe_warn_useless_init(settings)
         quote
             tforeach($(itrng); scheduler = $(settings.scheduler)) do $(itvar)
-                $(init_inner)
+                $(local_inner)
                 $(forbody)
             end
         end
@@ -54,8 +67,8 @@ function tasks_macro(forex)
     result = :(let
     end)
     push!(result.args[2].args, q)
-    if !isnothing(inits_before)
-        for x in inits_before
+    if !isnothing(locals_before)
+        for x in locals_before
             push!(result.args[1].args, x)
         end
     end
@@ -63,10 +76,16 @@ function tasks_macro(forex)
     result
 end
 
+function maybe_warn_useless_init(settings)
+    !isnothing(settings.init) &&
+        @warn("The @set init = ... settings won't have any effect because no reduction is performed.")
+end
+
 Base.@kwdef mutable struct Settings
     scheduler::Expr = :(DynamicScheduler())
     reducer::Union{Expr, Symbol, Nothing} = nothing
     collect::Bool = false
+    init::Union{Expr, Symbol, Nothing} = nothing
 end
 
 function _sym2scheduler(s)
@@ -81,64 +100,64 @@ function _sym2scheduler(s)
     end
 end
 
-function _maybe_handle_init_block!(args)
-    inits_before = nothing
-    init_inner = nothing
+function _maybe_handle_atlocal_block!(args)
+    locals_before = nothing
+    local_inner = nothing
     tlsidx = findfirst(args) do arg
-        arg isa Expr && arg.head == :macrocall && arg.args[1] == Symbol("@init")
+        arg isa Expr && arg.head == :macrocall && arg.args[1] == Symbol("@local")
     end
     if !isnothing(tlsidx)
-        inits_before, init_inner = _unfold_init_block(args[tlsidx].args[3])
+        locals_before, local_inner = _unfold_atlocal_block(args[tlsidx].args[3])
         deleteat!(args, tlsidx)
     end
-    return inits_before, init_inner
+    return locals_before, local_inner
 end
 
-function _unfold_init_block(ex)
-    inits_before = Expr[]
+function _unfold_atlocal_block(ex)
+    locals_before = Expr[]
     if ex.head == :(=)
-        initb, init_inner = _init_assign_to_exprs(ex)
-        push!(inits_before, initb)
+        localb, local_inner = _atlocal_assign_to_exprs(ex)
+        push!(locals_before, localb)
     elseif ex.head == :block
         tlsexprs = filter(x -> x isa Expr, ex.args) # skip LineNumberNode
-        init_inner = quote end
+        local_inner = quote end
         for x in tlsexprs
-            initb, initi = _init_assign_to_exprs(x)
-            push!(inits_before, initb)
-            push!(init_inner.args, initi)
+            localb, locali = _atlocal_assign_to_exprs(x)
+            push!(locals_before, localb)
+            push!(local_inner.args, locali)
         end
     else
-        throw(ErrorException("Wrong usage of @init. You must either provide a typed assignment or multiple typed assignments in a `begin ... end` block."))
+        throw(ErrorException("Wrong usage of @local. You must either provide a typed assignment or multiple typed assignments in a `begin ... end` block."))
     end
-    return inits_before, init_inner
+    return locals_before, local_inner
 end
 
-function _init_assign_to_exprs(ex)
+function _atlocal_assign_to_exprs(ex)
     left_ex = ex.args[1]
     if left_ex isa Symbol || left_ex.head != :(::)
-        throw(ErrorException("Wrong usage of @init. Expected typed assignment, e.g. `A::Matrix{Float} = rand(2,2)`."))
+        throw(ErrorException("Wrong usage of @local. Expected typed assignment, e.g. `A::Matrix{Float} = rand(2,2)`."))
     end
     tls_sym = esc(left_ex.args[1])
     tls_type = esc(left_ex.args[2])
     tls_def = esc(ex.args[2])
     @gensym tls_storage
-    init_before = :($(tls_storage) = TaskLocalValue{$tls_type}(() -> $(tls_def)))
-    init_inner = :($(tls_sym) = $(tls_storage)[])
-    return init_before, init_inner
+    local_before = :($(tls_storage) = TaskLocalValue{$tls_type}(() -> $(tls_def)))
+    local_inner = :($(tls_sym) = $(tls_storage)[])
+    return local_before, local_inner
 end
 
-function _maybe_handle_set_block!(settings, args)
+function _maybe_handle_atset_block!(settings, args)
     idcs = findall(args) do arg
         arg isa Expr && arg.head == :macrocall && arg.args[1] == Symbol("@set")
     end
-    isnothing(idcs) && return # no set block found
+    isnothing(idcs) && return # no @set block found
     for i in idcs
         ex = args[i].args[3]
         if ex.head == :(=)
-            _handle_set_single_assign!(settings, ex)
+            _handle_atset_single_assign!(settings, ex)
         elseif ex.head == :block
             exprs = filter(x -> x isa Expr, ex.args) # skip LineNumberNode
-            _handle_set_single_assign!.(Ref(settings), exprs)
+            _handle_atset_single_assign!.(Ref(settings), exprs)
         else
             throw(ErrorException("Wrong usage of @set. You must either provide an assignment or multiple assignments in a `begin ... end` block."))
         end
@@ -150,7 +169,7 @@ function _maybe_handle_set_block!(settings, args)
     end
 end
 
-function _handle_set_single_assign!(settings, ex)
+function _handle_atset_single_assign!(settings, ex)
     if ex.head != :(=)
         throw(ErrorException("Wrong usage of @set. Expected assignment, e.g. `scheduler = StaticScheduler()`."))
     end
diff --git a/src/macros.jl b/src/macros.jl
index ed56b416..fd3669c1 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -11,7 +11,7 @@ Supports reductions (`@set reducer=<reducer function>`) and collecting the resul
 Under the hood, the `for` loop is translated into corresponding parallel
 [`tforeach`](@ref), [`tmapreduce`](@ref), or [`tmap`](@ref) calls.
 
-See also: [`@set`](@ref), [`@init`](@ref)
+See also: [`@set`](@ref), [`@local`](@ref)
 
 ## Examples
 
@@ -67,46 +67,49 @@ Multiple settings are supported, either as separate `@set` statements or via
 * `scheduler` (e.g. `scheduler=:static`): Can be either a [`Scheduler`](@ref) or a `Symbol` (e.g. `:dynamic` or `:static`)
 * `reducer` (e.g. `reducer=+`): Indicates that a reduction should be performed with the provided binary function. See [`tmapreduce`](@ref) for more information.
 * `collect` (e.g. `collect=true`): Indicates that results should be collected (similar to `map`).
+* `init` (e.g. `init=0.0`): Initial value to be used in a reduction (requires `reducer=...`).
 """
 macro set(args...)
     error("The @set macro may only be used inside of a @tasks block.")
 end
 
-"""
-    @init name::T = value
-
-Can be used inside a `@tasks for ... end` block to specify
-[task-local values](@ref TLS) (TLV) via explicitly typed assignments.
-These values will be allocated once per task
-(rather than once per iteration) and can be re-used between different task-local iterations.
-
-There can only be a single `@init` block in a `@tasks for ... end` block. To specify
-multiple TLVs, use `@init begin ... end`. Compared to regular assignments, there are some
-limitations though, e.g. TLVs can't reference each other.
-
-## Examples
-
-```julia
-using OhMyThreads.Tools: taskid
-@tasks for i in 1:10
-    @set scheduler=DynamicScheduler(; nchunks=2)
-    @init x::Vector{Float64} = zeros(3) # TLV
-
-    x .+= 1
-    println(taskid(), " -> ", x)
-end
-```
-
-```julia
-@tasks for i in 1:10
-    @init begin
-        x::Vector{Int64} = rand(Int, 3)
-        M::Matrix{Float64} = rand(3, 3)
+@eval begin
+    """
+        @local name::T = value
+    
+    Can be used inside a `@tasks for ... end` block to specify
+    [task-local values](@ref TLS) (TLV) via explicitly typed assignments.
+    These values will be allocated once per task
+    (rather than once per iteration) and can be re-used between different task-local iterations.
+    
+    There can only be a single `@local` block in a `@tasks for ... end` block. To specify
+    multiple TLVs, use `@local begin ... end`. Compared to regular assignments, there are some
+    limitations though, e.g. TLVs can't reference each other.
+    
+    ## Examples
+    
+    ```julia
+    using OhMyThreads.Tools: taskid
+    @tasks for i in 1:10
+        @set scheduler=DynamicScheduler(; nchunks=2)
+        @local x::Vector{Float64} = zeros(3) # TLV
+    
+        x .+= 1
+        println(taskid(), " -> ", x)
+    end
+    ```
+    
+    ```julia
+    @tasks for i in 1:10
+        @local begin
+            x::Vector{Int64} = rand(Int, 3)
+            M::Matrix{Float64} = rand(3, 3)
+        end
+        # ...
+    end
+    ```
+    """
+    macro $(Symbol("local"))(args...)
+        error("The @local macro may only be used inside of a @tasks block.")
     end
-    # ...
-end
-```
-"""
-macro init(args...)
-    error("The @init macro may only be used inside of a @tasks block.")
 end
diff --git a/src/schedulers.jl b/src/schedulers.jl
index b1776a92..51badf3a 100644
--- a/src/schedulers.jl
+++ b/src/schedulers.jl
@@ -108,10 +108,6 @@ Base.@kwdef struct GreedyScheduler <: Scheduler
     end
 end
 
-@deprecate SpawnAllScheduler(args...; kwargs...) DynamicScheduler(args...;
-    nchunks = 0,
-    kwargs...)
-
 chunking_enabled(s::Scheduler) = chunking_enabled(typeof(s))
 chunking_enabled(::Type{DynamicScheduler{C}}) where {C} = C
 chunking_enabled(::Type{StaticScheduler{C}}) where {C} = C


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index d22725e7..2bc13c2b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -120,6 +120,22 @@ end
         i
         @set reducer=(+)
     end) == 55
+    # @set init
+    @test @tasks(for i in 1:10
+        @set begin
+            reducer=(+)
+            init=0.0
+        end
+        i
+    end) == 55.0
+    @test @tasks(for i in 1:10
+        @set begin
+            reducer=(+)
+            init=0.0*im
+        end
+        i
+    end) == (55.0 + 0.0im)
+
 
     # TaskLocalValue
     ntd = 2*Threads.nthreads()
@@ -127,7 +143,7 @@ end
     tids = Vector{UInt64}(undef, ntd)
     tid() = OhMyThreads.Tools.taskid()
     @test @tasks(for i in 1:ntd
-        @init C::Vector{Float64} = rand(3)
+        @local C::Vector{Float64} = rand(3)
         @set scheduler=:static
         ptrs[i] = pointer_from_objref(C)
         tids[i] = tid()
@@ -139,7 +155,7 @@ end
     end
     # TaskLocalValue (another fundamental check)
     @test @tasks(for i in 1:ntd
-        @init x::Ref{Int64} = Ref(0)
+        @local x::Ref{Int64} = Ref(0)
         @set reducer = (+)
         @set scheduler = :static
         x[] += 1
@@ -147,7 +163,7 @@ end
     end) == 1.5 * ntd # if a new x would be allocated per iteration, we'd get ntd here.
     # TaskLocalValue (begin ... end block)
     @test @tasks(for i in 1:10
-        @init begin
+        @local begin
             C::Matrix{Int64} = fill(4, 3, 3)
             x::Vector{Float64} = fill(5.0, 3)
         end
@@ -170,7 +186,7 @@ end
         x::Int
     end
     @test @tasks(for _ in 1:10
-        @init C::SingleInt = SingleInt(var)
+        @local C::SingleInt = SingleInt(var)
         @set reducer=+
         C.x
     end) == 10*var