-
Notifications
You must be signed in to change notification settings - Fork 18
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Add HyperDualNumbersExt * add test and dependency * minor fix on test * tests fix with proper testsets * fix hyperduals ext * bump version * some format fix
- Loading branch information
1 parent
00d50b3
commit a37ae26
Showing
5 changed files
with
357 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,249 @@ | ||
| module HyperDualNumbersExt | ||
|
|
||
| using HyperDualNumbers: Hyper | ||
| using Octavian: ArrayInterface, | ||
| @turbo, @tturbo, | ||
| One, Zero, | ||
| indices, static | ||
| import Octavian: real_rep, _matmul!, _matmul_serial! | ||
|
|
||
| real_rep(a::AbstractArray{DualT}) where {T,DualT<:Hyper{T}} = | ||
| reinterpret(reshape, T, a) | ||
| _view1(B::AbstractMatrix) = @view(B[1, :]) | ||
| _view1(B::AbstractArray{<:Any,3}) = @view(B[1, :, :]) | ||
|
|
||
| for AbstractVectorOrMatrix in (:AbstractVector, :AbstractMatrix) | ||
| # multiplication of dual vector/matrix by standard matrix from the left | ||
| @eval function _matmul!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| A::AbstractMatrix, | ||
| _B::$(AbstractVectorOrMatrix){DualT}, | ||
| α, | ||
| β = Zero(), | ||
| nthread::Nothing = nothing, | ||
| MKN = nothing, | ||
| contig_axis = nothing | ||
| ) where {T, DualT<:Hyper{T}} | ||
| B = real_rep(_B) | ||
| C = real_rep(_C) | ||
|
|
||
| @tturbo for n ∈ indices((C, B), 3), | ||
| m ∈ indices((C, A), (2, 1)), | ||
| l in indices((C, B), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), 2) | ||
| Cₗₘₙ += A[m, k] * B[l, k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
|
|
||
| _C | ||
| end | ||
|
|
||
| # multiplication of dual matrix by standard vector/matrix from the right | ||
| @eval @inline function _matmul!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| _A::AbstractMatrix{DualT}, | ||
| B::$(AbstractVectorOrMatrix), | ||
| α = One(), | ||
| β = Zero(), | ||
| nthread::Nothing = nothing, | ||
| MKN = nothing | ||
| ) where {T,DualT<:Hyper{T}} | ||
| if Bool(ArrayInterface.is_dense(_C)) && | ||
| Bool(ArrayInterface.is_column_major(_C)) && | ||
| Bool(ArrayInterface.is_dense(_A)) && | ||
| Bool(ArrayInterface.is_column_major(_A)) | ||
| # we can avoid the reshape and call the standard method | ||
| A = reinterpret(T, _A) | ||
| C = reinterpret(T, _C) | ||
| _matmul!(C, A, B, α, β, nthread, nothing) | ||
| else | ||
| # we cannot use the standard method directly | ||
| A = real_rep(_A) | ||
| C = real_rep(_C) | ||
|
|
||
| @tturbo for n ∈ indices((C, B), (3, 2)), | ||
| m ∈ indices((C, A), 2), | ||
| l in indices((C, A), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 1)) | ||
| Cₗₘₙ += A[l, m, k] * B[k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
| end | ||
|
|
||
| _C | ||
| end | ||
|
|
||
| @eval @inline function _matmul!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| _A::AbstractMatrix{DualT}, | ||
| _B::$(AbstractVectorOrMatrix){DualT}, | ||
| α = One(), | ||
| β = Zero(), | ||
| nthread::Nothing = nothing, | ||
| MKN = nothing, | ||
| contig = nothing | ||
| ) where {T,DualT<:Hyper{T}} | ||
| A = real_rep(_A) | ||
| C = real_rep(_C) | ||
| B = real_rep(_B) | ||
| if Bool(ArrayInterface.is_dense(_C)) && | ||
| Bool(ArrayInterface.is_column_major(_C)) && | ||
| Bool(ArrayInterface.is_dense(_A)) && | ||
| Bool(ArrayInterface.is_column_major(_A)) | ||
| # we can avoid the reshape and call the standard method | ||
| Ar = reinterpret(T, _A) | ||
| Cr = reinterpret(T, _C) | ||
| _matmul!(Cr, Ar, _view1(B), α, β, nthread, nothing) | ||
| else | ||
| # we cannot use the standard method directly | ||
| @tturbo for n ∈ indices((C, B), 3), | ||
| m ∈ indices((C, A), 2), | ||
| l in indices((C, A), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₗₘₙ += A[l, m, k] * B[1, k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
| end | ||
| @tturbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:3 | ||
| Cₚₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₚₘₙ += A[1, m, k] * B[p+1, k, n] | ||
| end | ||
| C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ | ||
| end | ||
|
|
||
| @tturbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2) | ||
| Cₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₘₙ += A[2, m, k] * B[3, k, n] + A[3, m, k] * B[2, k, n] | ||
| end | ||
| C[4, m, n] = C[4, m, n] + α * Cₘₙ | ||
| end | ||
| _C | ||
| end | ||
|
|
||
| # multiplication of dual vector/matrix by standard matrix from the left | ||
| @eval function _matmul_serial!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| A::AbstractMatrix, | ||
| _B::$(AbstractVectorOrMatrix){DualT}, | ||
| α, | ||
| β, | ||
| MKN | ||
| ) where {T, DualT<:Hyper{T}} | ||
| B = real_rep(_B) | ||
| C = real_rep(_C) | ||
|
|
||
| @turbo for n ∈ indices((C, B), 3), | ||
| m ∈ indices((C, A), (2, 1)), | ||
| l in indices((C, B), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), 2) | ||
| Cₗₘₙ += A[m, k] * B[l, k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
|
|
||
| _C | ||
| end | ||
|
|
||
| # multiplication of dual matrix by standard vector/matrix from the right | ||
| @eval @inline function _matmul_serial!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| _A::AbstractMatrix{DualT}, | ||
| B::$(AbstractVectorOrMatrix), | ||
| α, | ||
| β, | ||
| MKN | ||
| ) where {T,DualT<:Hyper{T}} | ||
| if Bool(ArrayInterface.is_dense(_C)) && | ||
| Bool(ArrayInterface.is_column_major(_C)) && | ||
| Bool(ArrayInterface.is_dense(_A)) && | ||
| Bool(ArrayInterface.is_column_major(_A)) | ||
| # we can avoid the reshape and call the standard method | ||
| A = reinterpret(T, _A) | ||
| C = reinterpret(T, _C) | ||
| _matmul_serial!(C, A, B, α, β, nothing) | ||
| else | ||
| # we cannot use the standard method directly | ||
| A = real_rep(_A) | ||
| C = real_rep(_C) | ||
|
|
||
| @turbo for n ∈ indices((C, B), (3, 2)), | ||
| m ∈ indices((C, A), 2), | ||
| l in indices((C, A), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 1)) | ||
| Cₗₘₙ += A[l, m, k] * B[k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
| end | ||
|
|
||
| _C | ||
| end | ||
|
|
||
| @eval @inline function _matmul_serial!( | ||
| _C::$(AbstractVectorOrMatrix){DualT}, | ||
| _A::AbstractMatrix{DualT}, | ||
| _B::$(AbstractVectorOrMatrix){DualT}, | ||
| α, | ||
| β, | ||
| MKN | ||
| ) where {T, DualT<:Hyper{T}} | ||
| A = real_rep(_A) | ||
| C = real_rep(_C) | ||
| B = real_rep(_B) | ||
| if Bool(ArrayInterface.is_dense(_C)) && | ||
| Bool(ArrayInterface.is_column_major(_C)) && | ||
| Bool(ArrayInterface.is_dense(_A)) && | ||
| Bool(ArrayInterface.is_column_major(_A)) | ||
| # we can avoid the reshape and call the standard method | ||
| Ar = reinterpret(T, _A) | ||
| Cr = reinterpret(T, _C) | ||
| _matmul_serial!(Cr, Ar, _view1(B), α, β, nothing) | ||
| else | ||
| # we cannot use the standard method directly | ||
| @turbo for n ∈ indices((C, B), 3), | ||
| m ∈ indices((C, A), 2), | ||
| l in indices((C, A), 1) | ||
|
|
||
| Cₗₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₗₘₙ += A[l, m, k] * B[1, k, n] | ||
| end | ||
| C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n] | ||
| end | ||
| end | ||
|
|
||
| @turbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:3 | ||
| Cₚₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₚₘₙ += A[1, m, k] * B[p+1, k, n] | ||
| end | ||
| C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ | ||
| end | ||
|
|
||
| @tturbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2) | ||
| Cₘₙ = zero(eltype(C)) | ||
| for k ∈ indices((A, B), (3, 2)) | ||
| Cₘₙ += A[2, m, k] * B[3, k, n] + A[3, m, k] * B[2, k, n] | ||
| end | ||
| C[4, m, n] = C[4, m, n] + α * Cₘₙ | ||
| end | ||
| _C | ||
| end | ||
| end # for | ||
|
|
||
| end # module |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
a37ae26There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@JuliaRegistrator register
a37ae26There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Registration pull request created: JuliaRegistries/General/85800
After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.
This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via: