Skip to content

Commit 3c09d19

Browse files
charleskawczynskimaleadt
charleskawczynski
and
maleadt
authored
Add Aqua.jl tests (#1965)
Co-authored-by: Tim Besard <tim.besard@gmail.com>
1 parent 6574bd2 commit 3c09d19

File tree

7 files changed

+189
-148
lines changed

7 files changed

+189
-148
lines changed

‎Project.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
1010
CUDA_Driver_jll = "4ee394cb-3365-5eb0-8335-949819d2adfc"
1111
CUDA_Runtime_Discovery = "1af6417a-86b4-443c-805f-a4643ffb695f"
1212
CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
13-
CompilerSupportLibraries_jll = "e66e0078-7015-5450-92f7-15fbd957f2ae"
1413
ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
1514
GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
1615
GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"

‎lib/cusparse/linalg.jl

Lines changed: 34 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ function sum_dim1(A::CuSparseMatrixCSR{T}) where {T}
1616
m, n = size(A)
1717
rowsum = CuVector{Float64}(undef, m)
1818
kernel_f = @cuda launch=false kernel(T, rowsum, A)
19-
19+
2020
config = launch_configuration(kernel_f.fun)
2121
threads = min(n, config.threads)
2222
blocks = cld(n, threads)
@@ -40,7 +40,7 @@ function sum_dim2(A::CuSparseMatrixCSR{T}) where {T}
4040
m, n = size(A)
4141
colsum = CuVector{Float64}(undef, n)
4242
kernel_f = @cuda launch=false kernel(T, colsum, A)
43-
43+
4444
config = launch_configuration(kernel_f.fun)
4545
threads = min(m, config.threads)
4646
blocks = cld(m, threads)
@@ -97,7 +97,7 @@ function SparseArrays.droptol!(A::CuSparseMatrixCOO, tol::Real)
9797
copyto!(A, B)
9898
end
9999

100-
function Base.reshape(A::CuSparseMatrixCOO, dims::NTuple{N,Int}) where {N}
100+
function Base.reshape(A::CuSparseMatrixCOO, dims::Dims)
101101
nrows, ncols = size(A)
102102
flat_indices = nrows .* (A.colInd .- 1) .+ A.rowInd .- 1
103103
new_col, new_row = div.(flat_indices, dims[1]) .+ 1, rem.(flat_indices, dims[1]) .+ 1
@@ -125,7 +125,7 @@ function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::CuSparseMatrixCOO{T}) wh
125125
col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1
126126

127127
data .*= repeat(B.nzVal, outer = Annz)
128-
128+
129129
sparse(row, col, data, out_shape..., fmt = :coo)
130130
end
131131

@@ -150,7 +150,7 @@ function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::Diagonal) where {T}
150150
col .+= CuVector(repeat(0:nB-1, outer = Annz)) .+ 1
151151

152152
data .*= repeat(CUDA.ones(T, nB), outer = Annz)
153-
153+
154154
sparse(row, col, data, out_shape..., fmt = :coo)
155155
end
156156

@@ -175,62 +175,62 @@ function LinearAlgebra.kron(A::Diagonal, B::CuSparseMatrixCOO{T}) where {T}
175175
col .+= repeat(B.colInd .- 1, outer = Annz) .+ 1
176176

177177
data .*= repeat(B.nzVal, outer = Annz)
178-
178+
179179
sparse(row, col, data, out_shape..., fmt = :coo)
180180
end
181181

182182
for SparseMatrixType in [:CuSparseMatrixCSC, :CuSparseMatrixCSR]
183183
@eval begin
184-
LinearAlgebra.triu(A::$SparseMatrixType{T,M}, k::Integer) where {T,M} =
184+
LinearAlgebra.triu(A::$SparseMatrixType{T}, k::Integer) where {T} =
185185
$SparseMatrixType( triu(CuSparseMatrixCOO(A), k) )
186-
LinearAlgebra.triu(A::Transpose{T,<:$SparseMatrixType}, k::Integer) where {T} =
186+
LinearAlgebra.triu(A::Transpose{T,<:$SparseMatrixType}, k::Integer) where {T} =
187187
$SparseMatrixType( triu(CuSparseMatrixCOO(_sptranspose(parent(A))), k) )
188-
LinearAlgebra.triu(A::Adjoint{T,<:$SparseMatrixType}, k::Integer) where {T} =
188+
LinearAlgebra.triu(A::Adjoint{T,<:$SparseMatrixType}, k::Integer) where {T} =
189189
$SparseMatrixType( triu(CuSparseMatrixCOO(_spadjoint(parent(A))), k) )
190-
191-
LinearAlgebra.tril(A::$SparseMatrixType{T,M}, k::Integer) where {T,M} =
190+
191+
LinearAlgebra.tril(A::$SparseMatrixType{T}, k::Integer) where {T} =
192192
$SparseMatrixType( tril(CuSparseMatrixCOO(A), k) )
193-
LinearAlgebra.tril(A::Transpose{T,<:$SparseMatrixType}, k::Integer) where {T} =
193+
LinearAlgebra.tril(A::Transpose{T,<:$SparseMatrixType}, k::Integer) where {T} =
194194
$SparseMatrixType( tril(CuSparseMatrixCOO(_sptranspose(parent(A))), k) )
195-
LinearAlgebra.tril(A::Adjoint{T,<:$SparseMatrixType}, k::Integer) where {T} =
195+
LinearAlgebra.tril(A::Adjoint{T,<:$SparseMatrixType}, k::Integer) where {T} =
196196
$SparseMatrixType( tril(CuSparseMatrixCOO(_spadjoint(parent(A))), k) )
197-
198-
LinearAlgebra.triu(A::Union{$SparseMatrixType{T,M}, Transpose{T,<:$SparseMatrixType}, Adjoint{T,<:$SparseMatrixType}}) where {T,M} =
197+
198+
LinearAlgebra.triu(A::Union{$SparseMatrixType{T}, Transpose{T,<:$SparseMatrixType}, Adjoint{T,<:$SparseMatrixType}}) where {T} =
199199
$SparseMatrixType( triu(CuSparseMatrixCOO(A), 0) )
200-
LinearAlgebra.tril(A::Union{$SparseMatrixType{T,M}, Transpose{T,<:$SparseMatrixType}, Adjoint{T,<:$SparseMatrixType}}) where {T,M} =
200+
LinearAlgebra.tril(A::Union{$SparseMatrixType{T},Transpose{T,<:$SparseMatrixType}, Adjoint{T,<:$SparseMatrixType}}) where {T} =
201201
$SparseMatrixType( tril(CuSparseMatrixCOO(A), 0) )
202202

203-
LinearAlgebra.kron(A::$SparseMatrixType{T,M}, B::$SparseMatrixType{T,M}) where {T,M} =
203+
LinearAlgebra.kron(A::$SparseMatrixType{T}, B::$SparseMatrixType{T}) where {T} =
204204
$SparseMatrixType( kron(CuSparseMatrixCOO(A), CuSparseMatrixCOO(B)) )
205-
LinearAlgebra.kron(A::$SparseMatrixType{T,M}, B::Diagonal) where {T,M} =
205+
LinearAlgebra.kron(A::$SparseMatrixType{T}, B::Diagonal) where {T} =
206206
$SparseMatrixType( kron(CuSparseMatrixCOO(A), B) )
207-
LinearAlgebra.kron(A::Diagonal, B::$SparseMatrixType{T,M}) where {T,M} =
207+
LinearAlgebra.kron(A::Diagonal, B::$SparseMatrixType{T}) where {T} =
208208
$SparseMatrixType( kron(A, CuSparseMatrixCOO(B)) )
209-
210-
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::$SparseMatrixType{T,M}) where {T,M} =
209+
210+
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::$SparseMatrixType{T}) where {T} =
211211
$SparseMatrixType( kron(CuSparseMatrixCOO(_sptranspose(parent(A))), CuSparseMatrixCOO(B)) )
212-
LinearAlgebra.kron(A::$SparseMatrixType{T,M}, B::Transpose{T,<:$SparseMatrixType}) where {T,M} =
212+
LinearAlgebra.kron(A::$SparseMatrixType{T}, B::Transpose{T,<:$SparseMatrixType}) where {T} =
213213
$SparseMatrixType( kron(CuSparseMatrixCOO(A), CuSparseMatrixCOO(_sptranspose(parent(B)))) )
214-
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::Transpose{T,<:$SparseMatrixType}) where {T} =
214+
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::Transpose{T,<:$SparseMatrixType}) where {T} =
215215
$SparseMatrixType( kron(CuSparseMatrixCOO(_sptranspose(parent(A))), CuSparseMatrixCOO(_sptranspose(parent(B)))) )
216-
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::Diagonal) where {T} =
216+
LinearAlgebra.kron(A::Transpose{T,<:$SparseMatrixType}, B::Diagonal) where {T} =
217217
$SparseMatrixType( kron(CuSparseMatrixCOO(_sptranspose(parent(A))), B) )
218-
LinearAlgebra.kron(A::Diagonal, B::Transpose{T,<:$SparseMatrixType}) where {T} =
218+
LinearAlgebra.kron(A::Diagonal, B::Transpose{T,<:$SparseMatrixType}) where {T} =
219219
$SparseMatrixType( kron(A, CuSparseMatrixCOO(_sptranspose(parent(B)))) )
220220

221-
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::$SparseMatrixType{T,M}) where {T,M} =
221+
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::$SparseMatrixType{T}) where {T} =
222222
$SparseMatrixType( kron(CuSparseMatrixCOO(_spadjoint(parent(A))), CuSparseMatrixCOO(B)) )
223-
LinearAlgebra.kron(A::$SparseMatrixType{T,M}, B::Adjoint{T,<:$SparseMatrixType}) where {T,M} =
223+
LinearAlgebra.kron(A::$SparseMatrixType{T}, B::Adjoint{T,<:$SparseMatrixType}) where {T} =
224224
$SparseMatrixType( kron(CuSparseMatrixCOO(A), CuSparseMatrixCOO(_spadjoint(parent(B)))) )
225-
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::Adjoint{T,<:$SparseMatrixType}) where {T} =
225+
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::Adjoint{T,<:$SparseMatrixType}) where {T} =
226226
$SparseMatrixType( kron(CuSparseMatrixCOO(_spadjoint(parent(A))), CuSparseMatrixCOO(_spadjoint(parent(B)))) )
227-
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::Diagonal) where {T} =
227+
LinearAlgebra.kron(A::Adjoint{T,<:$SparseMatrixType}, B::Diagonal) where {T} =
228228
$SparseMatrixType( kron(CuSparseMatrixCOO(_spadjoint(parent(A))), B) )
229-
LinearAlgebra.kron(A::Diagonal, B::Adjoint{T,<:$SparseMatrixType}) where {T} =
229+
LinearAlgebra.kron(A::Diagonal, B::Adjoint{T,<:$SparseMatrixType}) where {T} =
230230
$SparseMatrixType( kron(A, CuSparseMatrixCOO(_spadjoint(parent(B)))) )
231231

232232

233-
function Base.reshape(A::$SparseMatrixType, dims::NTuple{N,Int}) where {N}
233+
function Base.reshape(A::$SparseMatrixType, dims::Dims)
234234
B = CuSparseMatrixCOO(A)
235235
$SparseMatrixType(reshape(B, dims))
236236
end
@@ -244,16 +244,16 @@ for SparseMatrixType in [:CuSparseMatrixCSC, :CuSparseMatrixCSR]
244244
function LinearAlgebra.exp(A::$SparseMatrixType; threshold = 1e-7, nonzero_tol = 1e-14)
245245
rows = LinearAlgebra.checksquare(A) # Throws exception if not square
246246
typeA = eltype(A)
247-
247+
248248
mat_norm = norm(A, Inf)
249249
scaling_factor = nextpow(2, mat_norm) # Native routine, faster
250250
A = A ./ scaling_factor
251251
delta = 1
252-
252+
253253
P = $SparseMatrixType(spdiagm(0 => ones(eltype(A), rows)))
254254
next_term = P
255255
n = 1
256-
256+
257257
while delta > threshold
258258
next_term = typeA(1 / n) * A * next_term
259259
droptol!(next_term, nonzero_tol)

‎src/array.jl

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -237,9 +237,26 @@ array) or a tuple of the array dimensions. `own` optionally specified whether Ju
237237
take ownership of the memory, calling `cudaFree` when the array is no longer referenced. The
238238
`ctx` argument determines the CUDA context where the data is allocated in.
239239
"""
240-
function Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,N}},Type{CuArray{T,N,B}}},
240+
function Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,N}}},
241+
ptr::CuPtr{T}, dims::NTuple{N,Int};
242+
own::Bool=false, ctx::CuContext=context()) where {T,N}
243+
buf = _unsafe_wrap(T, ptr, dims; own, ctx)
244+
storage = ArrayStorage(buf, own ? 1 : -1)
245+
CuArray{T, length(dims)}(storage, dims)
246+
end
247+
function Base.unsafe_wrap(::Type{CuArray{T,N,B}},
241248
ptr::CuPtr{T}, dims::NTuple{N,Int};
242249
own::Bool=false, ctx::CuContext=context()) where {T,N,B}
250+
buf = _unsafe_wrap(T, ptr, dims; own, ctx)
251+
if typeof(buf) !== B
252+
error("Declared buffer type does not match inferred buffer type.")
253+
end
254+
storage = ArrayStorage(buf, own ? 1 : -1)
255+
CuArray{T, length(dims)}(storage, dims)
256+
end
257+
258+
function _unsafe_wrap(::Type{T}, ptr::CuPtr{T}, dims::NTuple{N,Int};
259+
own::Bool=false, ctx::CuContext=context()) where {T,N}
243260
isbitstype(T) || error("Can only unsafe_wrap a pointer to a bits type")
244261
sz = prod(dims) * sizeof(T)
245262

@@ -259,16 +276,15 @@ function Base.unsafe_wrap(::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,
259276
catch err
260277
error("Could not identify the buffer type; are you passing a valid CUDA pointer to unsafe_wrap?")
261278
end
262-
263-
if @isdefined(B) && typeof(buf) !== B
264-
error("Declared buffer type does not match inferred buffer type.")
265-
end
266-
267-
storage = ArrayStorage(buf, own ? 1 : -1)
268-
CuArray{T, length(dims)}(storage, dims)
279+
return buf
269280
end
270281

271-
function Base.unsafe_wrap(Atype::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,1}},Type{CuArray{T,1,B}}},
282+
function Base.unsafe_wrap(Atype::Union{Type{CuArray},Type{CuArray{T}},Type{CuArray{T,1}}},
283+
p::CuPtr{T}, dim::Int;
284+
own::Bool=false, ctx::CuContext=context()) where {T}
285+
unsafe_wrap(Atype, p, (dim,); own, ctx)
286+
end
287+
function Base.unsafe_wrap(Atype::Type{CuArray{T,1,B}},
272288
p::CuPtr{T}, dim::Int;
273289
own::Bool=false, ctx::CuContext=context()) where {T,B}
274290
unsafe_wrap(Atype, p, (dim,); own, ctx)

0 commit comments

Comments
 (0)