From 23cf805ef74bd12063e9e39b6ca1c9d556673b4c Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Tue, 28 May 2024 09:33:13 -0400
Subject: [PATCH] Add EnzymeCore extension

---
 lib/GPUArraysCore/Project.toml         |  7 +++++++
 lib/GPUArraysCore/ext/EnzymeCoreExt.jl | 27 ++++++++++++++++++++++++++
 test/Project.toml                      |  1 +
 test/gpuarrayscore.jl                  |  9 +++++++++
 test/runtests.jl                       |  4 ++++
 5 files changed, 48 insertions(+)
 create mode 100644 lib/GPUArraysCore/ext/EnzymeCoreExt.jl
 create mode 100644 test/gpuarrayscore.jl

diff --git a/lib/GPUArraysCore/Project.toml b/lib/GPUArraysCore/Project.toml
index c842d718..39c3624d 100644
--- a/lib/GPUArraysCore/Project.toml
+++ b/lib/GPUArraysCore/Project.toml
@@ -6,6 +6,13 @@ version = "0.1.6"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 
+[weakdeps]
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[extensions]
+EnzymeCoreExt = "EnzymeCore"
+
 [compat]
 Adapt = "4.0"
 julia = "1.6"
+EnzymeCore = "0.6, 0.7"
diff --git a/lib/GPUArraysCore/ext/EnzymeCoreExt.jl b/lib/GPUArraysCore/ext/EnzymeCoreExt.jl
new file mode 100644
index 00000000..ff02ddb4
--- /dev/null
+++ b/lib/GPUArraysCore/ext/EnzymeCoreExt.jl
@@ -0,0 +1,27 @@
+# compatibility with EnzymeCore
+
+module EnzymeCoreExt
+
+using GPUArraysCore
+
+if isdefined(Base, :get_extension)
+    using EnzymeCore
+    using EnzymeCore.EnzymeRules
+else
+    using ..EnzymeCore
+    using ..EnzymeCore.EnzymeRules
+end
+
+function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(GPUArraysCore.default_scalar_indexing), args...)
+    return nothing
+end
+
+function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(GPUArraysCore.assertscalar), args...)
+    return nothing
+end
+
+function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(GPUArraysCore.allowscalar), args...)
+    return nothing
+end
+
+end # module
diff --git a/test/Project.toml b/test/Project.toml
index 76e1e22a..fe9429b9 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
diff --git a/test/gpuarrayscore.jl b/test/gpuarrayscore.jl
new file mode 100644
index 00000000..88f1d62b
--- /dev/null
+++ b/test/gpuarrayscore.jl
@@ -0,0 +1,9 @@
+using Test, GPUArraysCore, EnzymeCore
+
+@testset "EnzymeCore" begin
+    @test nothing == EnzymeCore.EnzymeRules.inactive_noinl(GPUArraysCore.assertscalar)
+
+    @test nothing == EnzymeCore.EnzymeRules.inactive_noinl(GPUArraysCore.default_scalar_indexing)
+
+    @test nothing == EnzymeCore.EnzymeRules.inactive_noinl(GPUArraysCore.allowscalar, identity)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 4df72b2b..d5ef5f62 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,9 @@
 using GPUArrays, Test, Pkg
 
+@tests "GPUArraysCore" begin
+    include("gpuarrayscore.jl")
+end
+
 include("testsuite.jl")
 
 @testset "JLArray" begin