Skip to content

Offset matrix multiplication via generic_matmatmul! #270

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ version = "1.10.8"

[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[compat]
Adapt = "2, 3"
Expand All @@ -24,9 +25,8 @@ DistributedArrays = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Aqua", "CatIndices", "DistributedArrays", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation", "StaticArrays", "FillArrays"]
test = ["Aqua", "CatIndices", "DistributedArrays", "DelimitedFiles", "Documenter", "Test", "EllipsisNotation", "StaticArrays", "FillArrays"]
3 changes: 3 additions & 0 deletions src/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -651,6 +651,7 @@
no_offset_view(a::Base.Slice) = Base.Slice(UnitRange(a))
no_offset_view(S::SubArray) = view(parent(S), map(no_offset_view, parentindices(S))...)
end
no_offset_view(A::PermutedDimsArray{T,N,perm,iperm,P}) where {T,N,perm,iperm,P} = PermutedDimsArray(no_offset_view(parent(A)), perm)

Check warning on line 654 in src/OffsetArrays.jl

View check run for this annotation

Codecov / codecov/patch

src/OffsetArrays.jl#L654

Added line #L654 was not covered by tests
no_offset_view(a::Array) = a
no_offset_view(i::Number) = i
no_offset_view(A::AbstractArray) = _no_offset_view(axes(A), A)
Expand Down Expand Up @@ -853,6 +854,8 @@
import Adapt
Adapt.adapt_structure(to, O::OffsetArray) = parent_call(x -> Adapt.adapt(to, x), O)

include("linearalgebra.jl")

if Base.VERSION >= v"1.4.2"
include("precompile.jl")
_precompile_()
Expand Down
56 changes: 56 additions & 0 deletions src/linearalgebra.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
using LinearAlgebra
using LinearAlgebra: MulAddMul, mul!, AdjOrTrans

@inline LinearAlgebra.generic_matvecmul!(C::OffsetVector, fA::Function, A::AbstractVecOrMat, B::AbstractVector,
alpha, beta) = unwrap_matvecmul!(C, fA, A, B, alpha, beta)

Check warning on line 5 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L5

Added line #L5 was not covered by tests

@inline function unwrap_matvecmul!(C::OffsetVector, fA, A::AbstractVecOrMat, B::AbstractVector,

Check warning on line 7 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L7

Added line #L7 was not covered by tests
alpha, beta)

mB_axis = Base.axes1(B)
mA_axis, nA_axis = axes(fA(A))

Check warning on line 11 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L10-L11

Added lines #L10 - L11 were not covered by tests

if mB_axis != nA_axis
throw(DimensionMismatch("mul! can't contract axis $(UnitRange(nA_axis)) from A with axes(B) == ($(UnitRange(mB_axis)),)"))

Check warning on line 14 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L13-L14

Added lines #L13 - L14 were not covered by tests
end
if mA_axis != Base.axes1(C)
throw(DimensionMismatch("mul! got axes(C) == ($(UnitRange(Base.axes1(C))),), expected $(UnitRange(mA_axis))"))

Check warning on line 17 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L16-L17

Added lines #L16 - L17 were not covered by tests
end

mul!(no_offset_view(C), fA(no_offset_view(A)), no_offset_view(B), alpha, beta)
C

Check warning on line 21 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L20-L21

Added lines #L20 - L21 were not covered by tests
end

# The signatures of these differs from LinearAlgebra's *only* on C:
@inline LinearAlgebra.generic_matmatmul!(C::OffsetMatrix, fA::Function, fB::Function, A::AbstractMatrix, B::AbstractMatrix,
alpha, beta) = unwrap_matmatmul!(C, fA, fB, A, B, alpha, beta)

Check warning on line 26 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L26

Added line #L26 was not covered by tests

@inline LinearAlgebra.generic_matmatmul!(C::Union{OffsetMatrix, OffsetVector}, fA::Function, fB::Function, A::AbstractVecOrMat, B::AbstractVecOrMat,
alpha, beta) = unwrap_matmatmul!(C, fA, fB, A, B, alpha, beta)

Check warning on line 29 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L29

Added line #L29 was not covered by tests

@inline LinearAlgebra.generic_matmatmul!(C::AdjOrTrans{<:Any, <:OffsetArray}, fA::Function, fB::Function, A::AbstractMatrix, B::AbstractMatrix,
alpha, beta) = unwrap_matmatmul!(C, fA, fB, A, B, alpha, beta)

Check warning on line 32 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L32

Added line #L32 was not covered by tests

@inline function unwrap_matmatmul!(C::AbstractVecOrMat, fA, fB, A::AbstractVecOrMat, B::AbstractVecOrMat,

Check warning on line 34 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L34

Added line #L34 was not covered by tests
alpha, beta)

mA_axis, nA_axis = axes(fA(A))
mB_axis, nB_axis = axes(fB(B))

Check warning on line 38 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L37-L38

Added lines #L37 - L38 were not covered by tests

if nA_axis != mB_axis
throw(DimensionMismatch("mul! can't contract axis $(UnitRange(nA_axis)) from A with $(UnitRange(mB_axis)) from B"))
elseif mA_axis != axes(C,1)
throw(DimensionMismatch("mul! got axes(C,1) == $(UnitRange(axes(C,1))), expected $(UnitRange(mA_axis)) from A"))
elseif nB_axis != axes(C,2)
throw(DimensionMismatch("mul! got axes(C,2) == $(UnitRange(axes(C,2))), expected $(UnitRange(nB_axis)) from B"))

Check warning on line 45 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L40-L45

Added lines #L40 - L45 were not covered by tests
end

# Must be sure `no_offset_view(C)` won't match signature above!
mul!(no_offset_view(C), fA(no_offset_view(A)), fB(no_offset_view(B)), alpha, beta)
C

Check warning on line 50 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L49-L50

Added lines #L49 - L50 were not covered by tests
end

no_offset_view(A::Adjoint) = adjoint(no_offset_view(parent(A)))
no_offset_view(A::Transpose) = transpose(no_offset_view(parent(A)))

Check warning on line 54 in src/linearalgebra.jl

View check run for this annotation

Codecov / codecov/patch

src/linearalgebra.jl#L53-L54

Added lines #L53 - L54 were not covered by tests
no_offset_view(D::Diagonal) = Diagonal(no_offset_view(parent(D)))