Skip to content

Commit 048f947

Browse files
evelyne-ringootevelyne-ringoot
committed
Changing support for triu, tril and mul from AbstractArrays to AnyArrays
1 parent 03fe360 commit 048f947

File tree

1 file changed

+42
-42
lines changed

1 file changed

+42
-42
lines changed

‎src/host/linalg.jl

Lines changed: 42 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriang
117117
@eval Base.copyto!(A::$T{T, <:AbstractGPUArray{T,N}}, B::$T{T, <:AbstractGPUArray{T,N}}) where {T,N} = $T(copyto!(parent(A), parent(B)))
118118
end
119119

120-
function LinearAlgebra.tril!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
120+
function LinearAlgebra.tril!(A::AnyGPUMatrix{T}, d::Integer = 0) where T
121121
gpu_call(A, d; name="tril!") do ctx, _A, _d
122122
I = @cartesianidx _A
123123
i, j = Tuple(I)
@@ -129,7 +129,7 @@ function LinearAlgebra.tril!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
129129
return A
130130
end
131131

132-
function LinearAlgebra.triu!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
132+
function LinearAlgebra.triu!(A::AnyGPUMatrix{T}, d::Integer = 0) where T
133133
gpu_call(A, d; name="triu!") do ctx, _A, _d
134134
I = @cartesianidx _A
135135
i, j = Tuple(I)
@@ -206,9 +206,9 @@ if VERSION < v"1.8-"
206206
return B
207207
end
208208
else
209-
function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
210-
D::Diagonal{<:Any, <:AbstractGPUArray},
211-
A::AbstractGPUVecOrMat)
209+
function LinearAlgebra.mul!(B::AnyGPUVecOrMat,
210+
D::Diagonal{<:Any, <:AnyGPUArray},
211+
A::AnyGPUVecOrMat)
212212
dd = D.diag
213213
d = length(dd)
214214
m, n = size(A, 1), size(A, 2)
@@ -220,9 +220,9 @@ else
220220
B
221221
end
222222

223-
function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
224-
D::Diagonal{<:Any, <:AbstractGPUArray},
225-
A::AbstractGPUVecOrMat,
223+
function LinearAlgebra.mul!(B::AnyGPUVecOrMat,
224+
D::Diagonal{<:Any, <:AnyGPUArray},
225+
A::AnyGPUVecOrMat,
226226
α::Number,
227227
β::Number)
228228
dd = D.diag
@@ -236,9 +236,9 @@ else
236236
B
237237
end
238238

239-
function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
240-
A::AbstractGPUVecOrMat,
241-
D::Diagonal{<:Any, <:AbstractGPUArray})
239+
function LinearAlgebra.mul!(B::AnyGPUVecOrMat,
240+
A::AnyGPUVecOrMat,
241+
D::Diagonal{<:Any, <:AnyGPUArray})
242242
dd = D.diag
243243
d = length(dd)
244244
m, n = size(A, 1), size(A, 2)
@@ -250,9 +250,9 @@ else
250250
B
251251
end
252252

253-
function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
254-
A::AbstractGPUVecOrMat,
255-
D::Diagonal{<:Any, <:AbstractGPUArray},
253+
function LinearAlgebra.mul!(B::AnyGPUVecOrMat,
254+
A::AnyGPUVecOrMat,
255+
D::Diagonal{<:Any, <:AnyGPUArray},
256256
α::Number,
257257
β::Number)
258258
dd = D.diag
@@ -266,9 +266,9 @@ else
266266
B
267267
end
268268

269-
function LinearAlgebra.ldiv!(B::AbstractGPUVecOrMat,
270-
D::Diagonal{<:Any, <:AbstractGPUArray},
271-
A::AbstractGPUVecOrMat)
269+
function LinearAlgebra.ldiv!(B::AnyGPUVecOrMat,
270+
D::Diagonal{<:Any, <:AnyGPUArray},
271+
A::AnyGPUVecOrMat)
272272
dd = D.diag
273273
d = length(dd)
274274
m, n = size(A, 1), size(A, 2)
@@ -289,7 +289,7 @@ end
289289

290290
## matrix multiplication
291291

292-
function generic_matmatmul!(C::AbstractArray{R}, A::AbstractArray{T}, B::AbstractArray{S}, a::Number, b::Number) where {T,S,R}
292+
function generic_matmatmul!(C::AnyArray{R}, A::AnyArray{T}, B::AnyArray{S}, a::Number, b::Number) where {T,S,R}
293293
if size(A,2) != size(B,1)
294294
throw(DimensionMismatch("matrix A has dimensions $(size(A)), matrix B has dimensions $(size(B))"))
295295
end
@@ -319,29 +319,29 @@ function generic_matmatmul!(C::AbstractArray{R}, A::AbstractArray{T}, B::Abstrac
319319
C
320320
end
321321

322-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::AbstractGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
323-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
324-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
325-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::AbstractGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
326-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::AbstractGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
327-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
328-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
329-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
330-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
322+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::AnyGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
323+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
324+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
325+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::AnyGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
326+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::AnyGPUVecOrMat, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
327+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
328+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
329+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
330+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Number, b::Number) = generic_matmatmul!(C, A, B, a, b)
331331

332332
# specificity hacks
333-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::AbstractGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
334-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
335-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::AbstractGPUVecOrMat, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
336-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::AbstractGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
337-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::AbstractGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
338-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
339-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
340-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
341-
LinearAlgebra.mul!(C::AbstractGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AbstractGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
342-
343-
344-
function generic_rmul!(X::AbstractArray, s::Number)
333+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::AnyGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
334+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
335+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::AnyGPUVecOrMat, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
336+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::AnyGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
337+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::AnyGPUVecOrMat, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
338+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
339+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
340+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Adjoint{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
341+
LinearAlgebra.mul!(C::AnyGPUVecOrMat, A::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, B::LinearAlgebra.Transpose{<:Any, <:AnyGPUVecOrMat}, a::Real, b::Real) = generic_matmatmul!(C, A, B, a, b)
342+
343+
344+
function generic_rmul!(X::AnyArray, s::Number)
345345
gpu_call(X, s; name="rmul!") do ctx, X, s
346346
i = @linearidx X
347347
@inbounds X[i] *= s
@@ -350,9 +350,9 @@ function generic_rmul!(X::AbstractArray, s::Number)
350350
return X
351351
end
352352

353-
LinearAlgebra.rmul!(A::AbstractGPUArray, b::Number) = generic_rmul!(A, b)
353+
LinearAlgebra.rmul!(A::AnyGPUArray, b::Number) = generic_rmul!(A, b)
354354

355-
function generic_lmul!(s::Number, X::AbstractArray)
355+
function generic_lmul!(s::Number, X::AnyArray)
356356
gpu_call(X, s; name="lmul!") do ctx, X, s
357357
i = @linearidx X
358358
@inbounds X[i] = s*X[i]
@@ -361,7 +361,7 @@ function generic_lmul!(s::Number, X::AbstractArray)
361361
return X
362362
end
363363

364-
LinearAlgebra.lmul!(a::Number, B::AbstractGPUArray) = generic_lmul!(a, B)
364+
LinearAlgebra.lmul!(a::Number, B::AnyGPUArray) = generic_lmul!(a, B)
365365

366366

367367
## permutedims

0 commit comments

Comments
 (0)