Make allowscalar three-state and use for JLArray tests.

Author: maleadt

src/host/indexing.jl

@@ -5,39 +5,72 @@ export allowscalar, @allowscalar, assertscalar
 
 # mechanism to disallow scalar operations
 
-const scalar_allowed = Ref(true)
+@enum ScalarIndexing ScalarAllowed ScalarWarned ScalarDisallowed
+
+const scalar_allowed = Ref(ScalarWarned)
 const scalar_warned = Ref(false)
 
-function allowscalar(flag = true)
-    scalar_allowed[] = flag
+"""
+    allowscalar(allow=true, warn=true)
+
+Configure whether scalar indexing is allowed depending on the value of `allow`.
+
+If allowed, `warn` can be set to throw a single warning instead. Calling this function will
+reset the state of the warning, and throw a new warning on subsequent scalar iteration.
+"""
+function allowscalar(allow::Bool=true, warn::Bool=true)
     scalar_warned[] = false
+    scalar_allowed[] = if allow && !warn
+        ScalarAllowed
+    elseif allow
+        ScalarWarned
+    else
+        ScalarDisallowed
+    end
     return
 end
 
+"""
+    assertscalar(op::String)
+
+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")
-    if !scalar_allowed[]
+    if scalar_allowed[] == ScalarDisallowed
         error("$op is disallowed")
-    elseif !scalar_warned[]
+    elseif scalar_allowed[] == ScalarWarned && !scalar_warned[]
         @warn "Performing scalar operations on GPU arrays: This is very slow, consider disallowing these operations with `allowscalar(false)`"
         scalar_warned[] = true
     end
     return
 end
 
+"""
+    @allowscalar ex...
+    @disallowscalar ex...
+
+Temporarily allow or disallow scalar iteration.
+
+Note that this functionality is intended for functionality that is known and allowed to use
+scalar iteration (or not), i.e., there is no option to throw a warning. Only use this on
+fine-grained expressions.
+"""
 macro allowscalar(ex)
     quote
         local prev = scalar_allowed[]
-        scalar_allowed[] = true
+        scalar_allowed[] = ScalarAllowed
         local ret = $(esc(ex))
         scalar_allowed[] = prev
         ret
     end
 end
 
+@doc (@doc @allowscalar) ->
 macro disallowscalar(ex)
     quote
         local prev = scalar_allowed[]
-        scalar_allowed[] = false
+        scalar_allowed[] = ScalarDisallowed
         local ret = $(esc(ex))
         scalar_allowed[] = prev
         ret

test/runtests.jl

@@ -4,5 +4,6 @@ include("testsuite.jl")
 
 @testset "JLArray" begin
     using GPUArrays.JLArrays
+    JLArrays.allowscalar(false)
     TestSuite.test(JLArray)
 end