Merge pull request #450 from JuliaGPU/tb/repl_allowscalar

Allow scalar iteration in the REPL on 1.9+.

Author: maleadt

lib/GPUArraysCore/src/GPUArraysCore.jl

@@ -28,6 +28,7 @@ const AnyGPUArray{T,N} = Union{AbstractGPUArray{T,N}, WrappedGPUArray{T,N}}
 const AnyGPUVector{T} = AnyGPUArray{T, 1}
 const AnyGPUMatrix{T} = AnyGPUArray{T, 2}
 
+
 ## broadcasting
 
 """
@@ -38,6 +39,7 @@ this supertype.
 """
 abstract type AbstractGPUArrayStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end
 
+
 ## scalar iteration
 
 export allowscalar, @allowscalar, assertscalar
@@ -46,42 +48,47 @@ export allowscalar, @allowscalar, assertscalar
 
 # 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)
+const requested_scalar_indexing = Ref{Union{Nothing,ScalarIndexing}}(nothing)
 
-"""
-    allowscalar() do
-        # code that can use scalar indexing
+const _repl_frontend_task = Ref{Union{Nothing,Missing,Task}}()
+function repl_frontend_task()
+    if !isassigned(_repl_frontend_task)
+        _repl_frontend_task[] = get_repl_frontend_task()
+    end
+    _repl_frontend_task[]
+end
+function get_repl_frontend_task()
+    @static if VERSION >= v"1.10.0-DEV.444" || v"1.9-beta4" <= VERSION < v"1.10-"
+        if isdefined(Base, :active_repl)
+            Base.active_repl.frontend_task
+        else
+             missing
+        end
+    else
+        nothing
     end
-
-Denote which operations can use scalar indexing.
-
-See also: [`@allowscalar`](@ref).
-"""
-function allowscalar(f::Base.Callable)
-    task_local_storage(f, :ScalarIndexing, ScalarAllowed)
 end
 
-"""
-    allowscalar(::Bool)
-
-Calling this with `false` replaces the default warning about scalar indexing
-(show once per session) with an error.
-
-Instead of calling this with `true`, the preferred style is to allow this locally.
-This can be done with the `allowscalar(::Function)` method (with a `do` block)
-or with the [`@allowscalar`](@ref) macro.
-
-Writes to `task_local_storage` for `:ScalarIndexing`. The default is `:ScalarWarn`,
-and this function sets `:ScalarAllowed` or `:ScalarDisallowed`.
-"""
-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)
+@noinline function default_scalar_indexing()
+    if isinteractive()
+        # try to detect the REPL
+        repl_task = repl_frontend_task()
+        if repl_task isa Task
+            if repl_task === current_task()
+                # we always allow scalar iteration on the REPL's frontend task,
+                # where we often trigger scalar indexing by displaying GPU objects.
+                ScalarAllowed
+            else
+                ScalarDisallowed
+            end
+        else
+            # we couldn't detect a REPL in this interactive session, so default to a warning
+            ScalarWarn
+        end
+    else
+        # non-interactively, we always disallow scalar iteration
+        ScalarDisallowed
     end
-    setting = allow ? ScalarAllowed : ScalarDisallowed
-    task_local_storage(:ScalarIndexing, setting)
-    default_scalar_indexing[] = setting
-    return
 end
 
 """
@@ -90,26 +97,65 @@ end
 Assert that a certain operation `op` performs scalar indexing. If this is not allowed, an
 error will be thrown ([`allowscalar`](@ref)).
 """
-function assertscalar(op = "operation")
-    val = get!(task_local_storage(), :ScalarIndexing) do
-        something(default_scalar_indexing[], isinteractive() ? ScalarWarn : ScalarDisallowed)
+function assertscalar(op::String)
+    behavior = get(task_local_storage(), :ScalarIndexing, nothing)
+    if behavior === nothing
+        behavior = requested_scalar_indexing[]
+        if behavior === nothing
+            behavior = default_scalar_indexing()
+        end
+        task_local_storage(:ScalarIndexing, behavior)
     end
-    desc = """Invocation of $op resulted in scalar indexing of a GPU array.
+
+    behavior = behavior::ScalarIndexing
+    if behavior === ScalarAllowed
+        # fast path
+        return
+    end
+
+    _assertscalar(op, behavior)
+end
+
+@noinline function _assertscalar(op, behavior)
+    desc = """Invocation of '$op' resulted in scalar indexing of a GPU array.
               This is typically caused by calling an iterating implementation of a method.
               Such implementations *do not* execute on the GPU, but very slowly on the CPU,
-              and therefore are only permitted from the REPL for prototyping purposes.
-              If you did intend to index this array, annotate the caller with @allowscalar."""
-    if val == ScalarDisallowed
-        error("""Scalar indexing is disallowed.
-                 $desc""")
-    elseif val == ScalarWarn
-        @warn("""Performing scalar indexing on task $(current_task()).
-                 $desc""")
+              and therefore should be avoided.
+
+              If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
+              to enable scalar iteration globally or for the operations in question."""
+    if behavior == ScalarDisallowed
+        errorscalar(op)
+    elseif behavior == ScalarWarn
+        warnscalar(op)
         task_local_storage(:ScalarIndexing, ScalarWarned)
     end
+
     return
 end
 
+function scalardesc(op)
+    desc = """Invocation of $op resulted in scalar indexing of a GPU array.
+              This is typically caused by calling an iterating implementation of a method.
+              Such implementations *do not* execute on the GPU, but very slowly on the CPU,
+              and therefore should be avoided.
+
+              If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
+              to enable scalar iteration globally or for the operations in question."""
+end
+
+@noinline function warnscalar(op)
+    desc = scalardesc(op)
+    @warn("""Performing scalar indexing on task $(current_task()).
+             $desc""")
+end
+
+@noinline function errorscalar(op)
+    desc = scalardesc(op)
+    error("""Scalar indexing is disallowed.
+             $desc""")
+end
+
 # Like a try-finally block, except without introducing the try scope
 # NOTE: This is deprecated and should not be used from user logic. A proper solution to
 # this problem will be introduced in https://github.com/JuliaLang/julia/pull/39217
@@ -120,6 +166,34 @@ macro __tryfinally(ex, fin)
        )
 end
 
+"""
+    allowscalar([true])
+    allowscalar([true]) do
+        ...
+    end
+
+Use this function to allow or disallow scalar indexing, either globall or for the
+duration of the do block.
+
+See also: [`@allowscalar`](@ref).
+"""
+allowscalar
+
+function allowscalar(f::Base.Callable)
+    task_local_storage(f, :ScalarIndexing, ScalarAllowed)
+end
+
+function allowscalar(allow::Bool=true)
+    if allow
+        @warn """It's not recommended to use allowscalar([true]) to allow scalar indexing.
+                 Instead, use `allowscalar() do end` or `@allowscalar` to denote exactly which operations can use scalar operations.""" maxlog=1
+    end
+    setting = allow ? ScalarAllowed : ScalarDisallowed
+    task_local_storage(:ScalarIndexing, setting)
+    requested_scalar_indexing[] = setting
+    return
+end
+
 """
     @allowscalar() begin
         # code that can use scalar indexing
@@ -139,6 +213,9 @@ macro allowscalar(ex)
     end
 end
 
+
+## other
+
 """
     backend(T::Type)
     backend(x)