Skip to content

Commit d1c18ee

Browse files
authored
Merge pull request #351 from JuliaGPU/tb/global_assertscalar
Add an env var for setting allowscalar for all tasks.
2 parents 4de9fbb + 9018626 commit d1c18ee

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

src/host/indexing.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,13 @@ Assert that a certain operation `op` performs scalar indexing. If this is not al
5151
error will be thrown ([`allowscalar`](@ref)).
5252
"""
5353
function assertscalar(op = "operation")
54-
val = get(task_local_storage(), :ScalarIndexing, ScalarWarn)
54+
val = get!(task_local_storage(), :ScalarIndexing) do
55+
if haskey(ENV, "JULIA_GPU_ALLOWSCALAR")
56+
parse(Bool, ENV["JULIA_GPU_ALLOWSCALAR"]) ? ScalarAllowed : ScalarDisallowed
57+
else
58+
ScalarWarn
59+
end
60+
end
5561
if val == ScalarDisallowed
5662
error("$op is disallowed")
5763
elseif val == ScalarWarn

0 commit comments

Comments
 (0)