WIP: kernels

Author: William Moses

Project.toml

@@ -21,13 +21,15 @@ Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 [weakdeps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
+ReactantCUDAExt = "CUDA"
 ReactantNNlibExt = "NNlib"
 ReactantStatisticsExt = "Statistics"
 ReactantYaoBlocksExt = "YaoBlocks"
@@ -54,7 +56,8 @@ julia = "1.10"
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
-[sources]
-ReactantCore = { path = "lib/ReactantCore" }
+[sources.ReactantCore]
+path = "lib/ReactantCore"

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"

test/cuda.jl

@@ -0,0 +1,23 @@
+using Reactant
+using Test
+using CUDA
+
+function square_kernel!(x)
+    i = threadIdx().x
+    x[i] *= x[i]
+    sync_threads()
+    return nothing
+end
+
+# basic squaring on GPU
+function square!(x)
+    @cuda blocks = 1 threads = length(x) square_kernel!(x)
+    return nothing
+end
+
+@testset "Square Kernel" begin
+    oA = collect(1:1:64)
+    A = Reactant.to_rarray(oA)
+    func = @compile square!(A)
+    @test all(A .≈ (oA .* oA))
+end

test/runtests.jl

@@ -60,17 +60,18 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"
         @safetestset "Linear Algebra" include("integration/linear_algebra.jl")
+        @safetestset "CUDA" include("cuda.jl")
         @safetestset "AbstractFFTs" include("integration/fft.jl")
     end
 
-    if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "neural_networks"
-        @testset "Neural Networks" begin
-            @safetestset "NNlib Primitives" include("nn/nnlib.jl")
-            @safetestset "Flux.jl Integration" include("nn/flux.jl")
-            if Sys.islinux()
-                @safetestset "LuxLib Primitives" include("nn/luxlib.jl")
-                @safetestset "Lux Integration" include("nn/lux.jl")
-            end
-        end
-    end
+    # if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "neural_networks"
+    #     @testset "Neural Networks" begin
+    #         @safetestset "NNlib Primitives" include("nn/nnlib.jl")
+    #         @safetestset "Flux.jl Integration" include("nn/flux.jl")
+    #         if Sys.islinux()
+    #             @safetestset "LuxLib Primitives" include("nn/luxlib.jl")
+    #             @safetestset "Lux Integration" include("nn/lux.jl")
+    #         end
+    #     end
+    # end
 end