Generic support for ghost types in launch (#93)

Author: AntonOresten

ext/CUDAExt.jl

@@ -2,7 +2,7 @@ module CUDAExt
 
 using cuTile
 using cuTile: TileArray, Constant, CGOpts, CuTileResults, emit_code, sanitize_name,
-              constant_eltype, constant_value
+              constant_eltype, constant_value, is_ghost_type
 
 using CompilerCaching: CacheView, method_instance, results
 
@@ -187,9 +187,8 @@ return their fields in order.
 
 This is used by the launch helper to splat arguments to cudacall.
 """
-flatten(x) = (x,)
+flatten(x) = is_ghost_type(typeof(x)) ? () : (x,)
 flatten(arr::TileArray{T, N}) where {T, N} = (arr.ptr, arr.sizes..., arr.strides...)
-flatten(::Constant) = ()  # Ghost types are not passed to cudacall
 
 """
     to_tile_arg(x)

test/codegen/integration.jl

@@ -1123,3 +1123,39 @@ end
         end
     end
 end
+
+#=============================================================================
+ Ghost Type Arguments
+=============================================================================#
+
+@testset "Non-Constant Ghost Type Arguments" begin
+    spec = ct.ArraySpec{1}(16, true)
+
+    @testset "Nothing argument" begin
+        @test @filecheck begin
+            @check_label "entry"
+            @check "load_view_tko"
+            @check "store_view_tko"
+            code_tiled(Tuple{ct.TileArray{Float32,1,spec}, ct.TileArray{Float32,1,spec}, Nothing}) do a, b, _
+                pid = ct.bid(1)
+                tile = ct.load(a, pid, (16,))
+                ct.store(b, pid, tile)
+                return
+            end
+        end
+    end
+
+    @testset "Val argument" begin
+        @test @filecheck begin
+            @check_label "entry"
+            @check "load_view_tko"
+            @check "store_view_tko"
+            code_tiled(Tuple{ct.TileArray{Float32,1,spec}, ct.TileArray{Float32,1,spec}, Val{16}}) do a, b, _
+                pid = ct.bid(1)
+                tile = ct.load(a, pid, (16,))
+                ct.store(b, pid, tile)
+                return
+            end
+        end
+    end
+end

test/execution/basic.jl

@@ -1065,3 +1065,36 @@ end
     ct.launch(kernel!, 1)
 end
 
+@testset "non-Constant ghost type argument (nothing)" begin
+    function ghost_nothing_kernel(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, ::Nothing)
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        ct.store(b, pid, tile)
+        return
+    end
+
+    n = 256
+    a = CUDA.rand(Float32, n)
+    b = CUDA.zeros(Float32, n)
+
+    ct.launch(ghost_nothing_kernel, cld(n, 16), a, b, nothing)
+
+    @test Array(b) ≈ Array(a)
+end
+
+@testset "non-Constant ghost type argument (Val)" begin
+    function ghost_val_kernel(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, ::Val{n}) where n
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (n,))
+        ct.store(b, pid, tile)
+        return
+    end
+
+    n = 256
+    a = CUDA.rand(Float32, n)
+    b = CUDA.zeros(Float32, n)
+
+    ct.launch(ghost_val_kernel, cld(n, 16), a, b, Val(16))
+
+    @test Array(b) ≈ Array(a)
+end