Skip to content

Commit

Permalink
Inherit default allowscalar settings in child tasks.
Browse files Browse the repository at this point in the history
  • Loading branch information
maleadt committed Jun 17, 2021
1 parent efb14c0 commit 52ecb4d
Showing 1 changed file with 9 additions and 7 deletions.
16 changes: 9 additions & 7 deletions src/host/indexing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ export allowscalar, @allowscalar, assertscalar

@enum ScalarIndexing ScalarAllowed ScalarWarn ScalarWarned ScalarDisallowed

# if the user explicitly calls allowscalar, use that setting for all new tasks
# XXX: use context variables to inherit the parent task's setting, once available.
const default_scalar_indexing = Ref{Union{Nothing,ScalarIndexing}}(nothing)

"""
allowscalar() do
# code that can use scalar indexing
Expand All @@ -24,7 +28,9 @@ function allowscalar(allow::Bool=true)
if allow
Base.depwarn("allowscalar([true]) is deprecated, use `allowscalar() do end` or `@allowscalar` to denote exactly which operations can use scalar operations.", :allowscalar)
end
task_local_storage(:ScalarIndexing, allow ? ScalarAllowed : ScalarDisallowed)
setting = allow ? ScalarAllowed : ScalarDisallowed
task_local_storage(:ScalarIndexing, setting)
default_scalar_indexing[] = setting
return
end

Expand All @@ -36,11 +42,7 @@ error will be thrown ([`allowscalar`](@ref)).
"""
function assertscalar(op = "operation")
val = get!(task_local_storage(), :ScalarIndexing) do
if isinteractive()
ScalarWarn
else
ScalarDisallowed
end
something(default_scalar_indexing[], isinteractive() ? ScalarWarn : ScalarDisallowed)
end
desc = """Invocation of $op resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Expand All @@ -51,7 +53,7 @@ function assertscalar(op = "operation")
error("""Scalar indexing is disallowed.
$desc""")
elseif val == ScalarWarn
@warn("""Performing scalar indexing.
@warn("""Performing scalar indexing on task $(current_task()).
$desc""")
task_local_storage(:ScalarIndexing, ScalarWarned)
end
Expand Down

0 comments on commit 52ecb4d

Please sign in to comment.