Merge pull request #486 from pxl-th/ka-kernels

Use KernelAbstractions.jl for upsample kernels
FluxML · Apr 10, 2023 · ee909e6 · ee909e6
2 parents 335ed1c + bf4d6cf
commit ee909e6
Show file tree

Hide file tree

Showing 11 changed files with 823 additions and 1,048 deletions.
diff --git a/Project.toml b/Project.toml
@@ -5,6 +5,7 @@ version = "0.8.19"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -21,6 +22,7 @@ NNlibAMDGPUExt = "AMDGPU"
 AMDGPU = "0.4.8"
 Adapt = "2, 3.2"
 ChainRulesCore = "1.13"
+KernelAbstractions = "0.9"
 Requires = "0.5, 1.0"
 julia = "1.6"
 

diff --git a/ext/NNlibCUDA/src/NNlibCUDA.jl b/ext/NNlibCUDA/src/NNlibCUDA.jl
@@ -6,7 +6,6 @@ using Random, Statistics
 
 const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}
 
-include("upsample.jl")
 include("sampling.jl")
 include("activations.jl")
 include("batchedadjtrans.jl")

diff --git a/ext/NNlibCUDA/src/upsample.jl b/ext/NNlibCUDA/src/upsample.jl
diff --git a/ext/NNlibCUDA/test/runtests.jl b/ext/NNlibCUDA/test/runtests.jl
@@ -13,7 +13,6 @@ include("activations.jl")
 include("dropout.jl")
 include("batchedadjtrans.jl")
 include("batchedmul.jl")
-include("upsample.jl")
 include("conv.jl")
 include("ctc.jl")
 include("fold.jl")