Merge pull request #186 from JuliaDiff/core

Change to ArrayInterfaceCore

Author: ChrisRackauckas

Project.toml

@@ -5,7 +5,8 @@ version = "1.21.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
+ArrayInterfaceStaticArrays = "b0d46f97-bff5-4637-a19a-dd75974142cd"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -19,7 +20,8 @@ VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [compat]
 Adapt = "1, 2.0, 3.0"
-ArrayInterface = "2.8, 3.0, 4, 5"
+ArrayInterfaceCore = "0.1.1"
+ArrayInterfaceStaticArrays = "0.1"
 Compat = "2.2, 3"
 DataStructures = "0.17, 0.18"
 FiniteDiff = "2.8.1"
@@ -31,6 +33,8 @@ VertexSafeGraphs = "0.2"
 julia = "1.6"
 
 [extras]
+ArrayInterfaceBlockBandedMatrices = "5331f1e9-51c7-46b0-a9b0-df4434785e0a"
+ArrayInterfaceBandedMatrices = "2e50d22c-5be1-4042-81b1-c572ed69783d"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
@@ -43,4 +47,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "BandedMatrices", "BlockBandedMatrices", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Zygote", "SparsityDetection", "StaticArrays"]
+test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Zygote", "SparsityDetection", "StaticArrays"]

README.md

@@ -112,7 +112,7 @@ gmres!(res,J,v)
 
 ### Matrix Coloring
 
-This library extends the common `ArrayInterface.matrix_colors` function to allow
+This library extends the common `ArrayInterfaceCore.matrix_colors` function to allow
 for coloring sparse matrices using graphical techniques.
 
 Matrix coloring allows you to reduce the number of times finite differencing
@@ -163,7 +163,7 @@ forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
 
 Notice that if a sparsity pattern is not supplied then the built Jacobian will
 be the compressed Jacobian: `sparsity` must be a sparse matrix or a structured matrix
-(`Tridiagonal`, `Banded`, etc. conforming to the ArrayInterface.jl specs) with the
+(`Tridiagonal`, `Banded`, etc. conforming to the ArrayInterfaceCore.jl specs) with the
 appropriate sparsity pattern to allow for decompression.
 
 This call will allocate the cache variables each time. To avoid allocating the
@@ -317,3 +317,14 @@ These all have the same interface, where `J*v` utilizes the out-of-place
 Jacobian-vector or Hessian-vector function, whereas `mul!(res,J,v)` utilizes
 the appropriate in-place versions. To update the location of differentiation
 in the operator, simply mutate the vector `u`: `J.u .= ...`.
+
+# Note about sparse differentiation of GPUArrays, BandedMatrices, and BlockBandedMatrices
+
+These two matrix types need the dependencies ArrayInterfaceBandedMatrices.jl and
+ArrayInterfaceBlockBandedMatrices.jl to basically work with any functionality
+(anywhere). For now, the right thing to do is to add these libraries and do
+`import` on them if you are using BandedMatrices.jl or BlockBandedMatrices.jl
+for sparsity patterns. In the future, those two packages should just depend on
+ArrayInterface.jl and remove this issue entirely from the user space.
+
+Additionally, GPUs need ArrayInterfaceGPUArrays for proper determination of the indexing.

src/SparseDiffTools.jl

@@ -10,14 +10,14 @@ using VertexSafeGraphs
 using Adapt
 
 using LinearAlgebra
-using SparseArrays, ArrayInterface
+using SparseArrays, ArrayInterfaceCore
 
-import StaticArrays
+import StaticArrays, ArrayInterfaceStaticArrays
 
 using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using DataStructures: DisjointSets, find_root!, union!
 
-using ArrayInterface: matrix_colors
+using ArrayInterfaceCore: matrix_colors
 
 export  contract_color,
         greedy_d1,

src/coloring/high_level.jl

@@ -1,4 +1,4 @@
-abstract type SparseDiffToolsColoringAlgorithm <: ArrayInterface.ColoringAlgorithm end
+abstract type SparseDiffToolsColoringAlgorithm <: ArrayInterfaceCore.ColoringAlgorithm end
 struct GreedyD1Color <: SparseDiffToolsColoringAlgorithm end
 struct BacktrackingColor <: SparseDiffToolsColoringAlgorithm end
 struct ContractionColor <: SparseDiffToolsColoringAlgorithm end
@@ -16,7 +16,7 @@ The coloring defaults to a greedy distance-1 coloring.
 Note that if A isa SparseMatrixCSC, the sparsity pattern is defined by structural nonzeroes,
 ie includes explicitly stored zeros.
 """
-function ArrayInterface.matrix_colors(A::AbstractMatrix, alg::SparseDiffToolsColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false)
+function ArrayInterfaceCore.matrix_colors(A::AbstractMatrix, alg::SparseDiffToolsColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false)
     _A = A isa SparseMatrixCSC ? A : sparse(A) # Avoid the copy
     A_graph = matrix2graph(_A, partition_by_rows)
     return color_graph(A_graph, alg)

src/differentiation/compute_jacobian_ad.jl

@@ -43,15 +43,15 @@ function ForwardColorJacCache(f::F,x,_chunksize = nothing;
     else
         p = adapt.(parameterless_type(x),generate_chunked_partials(x,colorvec,chunksize))
         _t = Dual{T,eltype(x),getsize(chunksize)}.(vec(x),ForwardDiff.Partials.(first(p)))
-        t = ArrayInterface.restructure(x,_t)
+        t = ArrayInterfaceCore.restructure(x,_t)
     end
 
 
     if dx isa Nothing
         fx = similar(t)
         _dx = similar(x)
     else
-        tup = ArrayInterface.allowed_getindex(ArrayInterface.allowed_getindex(p,1),1) .* false
+        tup = ArrayInterfaceCore.allowed_getindex(ArrayInterfaceCore.allowed_getindex(p,1),1) .* false
         _pi = adapt(parameterless_type(dx),[tup for i in 1:length(dx)])
         fx = reshape(Dual{T,eltype(dx),length(tup)}.(vec(dx),ForwardDiff.Partials.(_pi)),size(dx)...)
         _dx = dx
@@ -134,7 +134,7 @@ end
 
 function forwarddiff_color_jacobian(f::F,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing) where F
 
-    if jac_prototype isa Nothing ? ArrayInterface.ismutable(x) : ArrayInterface.ismutable(jac_prototype)
+    if jac_prototype isa Nothing ? ArrayInterfaceCore.ismutable(x) : ArrayInterfaceCore.ismutable(jac_prototype)
         # Whenever J is mutable, we mutate it to avoid allocations
         dx = jac_cache.dx
         vecx = vec(x)
@@ -164,7 +164,7 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f::F,x::Abstract
     nrows,ncols = size(J)
 
     if !(sparsity isa Nothing)
-        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        rows_index, cols_index = ArrayInterfaceCore.findstructralnz(sparsity)
         rows_index = [rows_index[i] for i in 1:length(rows_index)]
         cols_index = [cols_index[i] for i in 1:length(cols_index)]
     end
@@ -232,7 +232,7 @@ function forwarddiff_color_jacobian_immutable(f,x::AbstractArray{<:Number},jac_c
     nrows,ncols = size(J)
 
     if !(sparsity isa Nothing)
-        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        rows_index, cols_index = ArrayInterfaceCore.findstructralnz(sparsity)
         rows_index = [rows_index[i] for i in 1:length(rows_index)]
         cols_index = [cols_index[i] for i in 1:length(cols_index)]
     end
@@ -277,7 +277,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 x::AbstractArray{<:Number};
                 dx = similar(x,size(J,1)),
                 colorvec = 1:length(x),
-                sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
+                sparsity = ArrayInterfaceCore.has_sparsestruct(J) ? J : nothing)
     forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity))
 end
 
@@ -305,7 +305,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     end
 
     if FiniteDiff._use_findstructralnz(sparsity)
-        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        rows_index, cols_index = ArrayInterfaceCore.findstructralnz(sparsity)
     else
         rows_index = 1:size(J,1)
         cols_index = 1:size(J,2)
@@ -344,7 +344,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                     dx .= partials.(fx, j)
                 end
 
-                if ArrayInterface.fast_scalar_indexing(dx)
+                if ArrayInterfaceCore.fast_scalar_indexing(dx)
                     #dx is implicitly used in vecdx
                     if sparseCSC_common_sparsity
                         FiniteDiff._colorediteration!(J,vecdx,colorvec,color_i,ncols)

test/gpu/Project.toml

@@ -1,2 +1,3 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ArrayInterfaceGPUArrays = "6ba088a2-8465-4c0a-af30-387133b534db"