Introduce a macro for marking multiple functions as device-only. (#1117)

Use it to mark the indexing functions device-only.

Author: maleadt

src/device/intrinsics/indexing.jl

@@ -61,6 +61,8 @@ for dim in (:x, :y, :z)
     @eval @inline $fn() = Int(_index($(Val(intr)), $(Val(1:max_grid_size[dim]))))
 end
 
+@device_functions begin
+
 """
     gridDim()::CuDim3
 
@@ -111,4 +113,6 @@ executing thread.
 """
 @inline active_mask() = @asmcall("activemask.b32 \$0;", "=r", false, UInt32, Tuple{})
 
+end
+
 const FULL_MASK = 0xffffffff

src/device/utils.jl

@@ -36,3 +36,27 @@ macro device_function(ex)
         @device_override $ex
     end)
 end
+
+macro device_functions(ex)
+    ex = macroexpand(__module__, ex)
+
+    # recursively prepend `@device_function` to all function definitions
+    function rewrite(block)
+        out = Expr(:block)
+        for arg in block.args
+            if Meta.isexpr(arg, :block)
+                # descend in blocks
+                push!(out.args, rewrite(arg))
+            elseif Meta.isexpr(arg, [:function, :(=)])
+                # rewrite function definitions
+                push!(out.args, :(@device_function $arg))
+            else
+                # preserve all the rest
+                push!(out.args, arg)
+            end
+        end
+        out
+    end
+
+    esc(rewrite(ex))
+end