Metal Extension with fix for tanh_fast (#666)

* metal ext

* tests

* add test

* don't test on nightly

Author: CarloLucibello

.buildkite/pipeline.yml

@@ -19,8 +19,7 @@ steps:
     env:
       JULIA_NUM_THREADS: 4
       NNLIB_TEST_CUDA: "true"
-      NNLIB_TEST_CPU: "true" # Could be useful to uncover multithreading related issues
-                             # Buildkite workers have more threads.
+      NNLIB_TEST_CPU: "false"
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 180
     matrix:
@@ -34,6 +33,7 @@ steps:
             julia: "nightly"
           soft_fail: true
 
+
   - label: ":julia: Julia {{matrix.julia}} - AMD GPU"
     command:
       - echo 'AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"' >> test/Project.toml
@@ -65,10 +65,43 @@ steps:
           # - "1.10"  
           - "1"
           # - "nightly"
-      adjustments:
-        - with:
-            julia: "nightly"
-          soft_fail: true
+      # adjustments:
+      #   - with:
+      #       julia: "nightly"
+      #     soft_fail: true
+
+
+  - label: ":julia: Julia {{matrix.julia}} - Metal GPU"
+    command:
+      - echo 'Metal = "dde4c033-4e86-420c-a63e-0dd931031962"' >> test/Project.toml
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+          dirs:
+            - src
+            - ext
+    agents:
+      os: "macos"
+      arch: "aarch64"
+    timeout_in_minutes: 180
+    env:
+      NNLIB_TEST_METAL: "true"
+      NNLIB_TEST_CPU: "false"
+      JULIA_NUM_THREADS: 4
+    matrix:
+      setup:
+        julia:
+          # - "1.10"
+          - "1"
+          # - "nightly  "
+      # adjustments:
+      #   - with:
+      #       julia: "nightly"
+      #     soft_fail: true
 
 
   - label: "Benchmarks"

Project.toml

@@ -19,6 +19,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
@@ -29,6 +30,7 @@ NNlibCUDAExt = "CUDA"
 NNlibEnzymeCoreExt = "EnzymeCore"
 NNlibFFTWExt = "FFTW"
 NNlibForwardDiffExt = "ForwardDiff"
+NNlibMetalExt = "Metal"
 NNlibSpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -43,9 +45,10 @@ ForwardDiff = "1"
 GPUArraysCore = "0.2"
 KernelAbstractions = "0.9.2"
 LinearAlgebra = "1"
+Metal = "1.6"
 Random = "1"
 ScopedValues = "1.3.0"
 SpecialFunctions = "2"
 Statistics = "1"
 cuDNN = "1"
-julia = "1.10"
+julia = "1.10"

ext/NNlibMetalExt.jl

@@ -0,0 +1,9 @@
+module NNlibMetalExt
+
+
+using Metal: method_table, @device_override
+using NNlib: NNlib
+
+@device_override NNlib.tanh_fast(x) = Base.FastMath.tanh_fast(x)
+
+end

test/Project.toml

@@ -14,6 +14,7 @@ Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/ext_metal/activations.jl

@@ -0,0 +1,17 @@
+@testset "activation broadcast" begin
+    broken_f = (:hardσ, :leakyrelu) 
+    for name in NNlib.ACTIVATIONS
+        # println("Testing forward diff for activation: ", name)
+        f = @eval $name
+        @test gputest(DEVICE, x -> f.(x), rand(5)) broken=name ∈ broken_f
+    end
+end
+
+@testset "forward diff" begin
+    broken_f = (:hardσ, :leakyrelu) 
+    for name in NNlib.ACTIVATIONS
+        # println("Testing forward diff for activation: ", name)
+        f = @eval $name
+        @test gputest(DEVICE, x -> f.(x), Dual.(rand(Float32, 5), 1)) broken=name ∈ broken_f
+    end
+end

test/ext_metal/runtests.jl

@@ -0,0 +1,35 @@
+using NNlib
+using Test
+using Metal
+using Zygote: gradient
+using MLDataDevices: gpu_device
+using ForwardDiff: Dual
+
+Metal.allowscalar(false)
+
+#TODO move this to test/ test_utils.jl and use it with all backends
+function gputest(device, f, xs...; checkgrad=true, atol=1e-6, kws...)
+    cpu_in = xs
+    gpu_in = device(xs)
+
+    cpu_out = f(cpu_in...; kws...)
+    gpu_out = f(gpu_in...; kws...)
+    @test collect(cpu_out) ≈ collect(gpu_out)
+
+    if checkgrad
+        cpu_grad = gradient((x...) -> sum(f(x...; kws...)), cpu_in...)
+        gpu_grad = gradient((x...) -> sum(f(x...; kws...)), gpu_in...)
+        for (cpu_g, gpu_g) in zip(cpu_grad, gpu_grad)
+            if cpu_g === nothing
+                @test gpu_g === nothing
+            else
+                @test collect(cpu_g) ≈ collect(gpu_g) atol=atol
+            end
+        end
+    end
+    return true
+end
+
+DEVICE = gpu_device(force=true)
+
+include("activations.jl")

test/runtests.jl

@@ -24,6 +24,7 @@ DocMeta.setdocmeta!(NNlib, :DocTestSetup, :(using NNlib, UnicodePlots); recursiv
 
 # ENV["NNLIB_TEST_CUDA"] = "true" # uncomment to run CUDA tests
 # ENV["NNLIB_TEST_AMDGPU"] = "true" # uncomment to run AMDGPU tests
+# ENV["NNLIB_TEST_METAL"] = "true" # uncomment to run Metal tests
 # ENV["NNLIB_TEST_CPU"] = "false" # uncomment to skip CPU tests
 
 const rng = StableRNG(123)
@@ -184,4 +185,20 @@ end
     else
         @info "Skipping AMDGPU tests, set NNLIB_TEST_AMDGPU=true to run them."
     end
+
+    if get(ENV, "NNLIB_TEST_METAL", "false") == "true"
+        Pkg.add("Metal")
+
+        using Metal
+        if Metal.functional()
+            @testset "Metal" begin
+                # nnlib_testsuite(MetalBackend)
+                include("ext_metal/runtests.jl")
+            end
+        else
+            @info "Insufficient version or Metal not found; Skipping Metal tests"
+        end
+    else
+        @info "Skipping Metal tests, set NNLIB_TEST_METAL=true to run them"
+    end
 end