Skip to content
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

Efficient Inplace Accumulating Strided MatMul #273

Merged
merged 11 commits into from
Oct 20, 2020
Merged

Efficient Inplace Accumulating Strided MatMul #273

merged 11 commits into from
Oct 20, 2020

Conversation

oxinabox
Copy link
Member

@oxinabox oxinabox commented Oct 2, 2020

These come out from Nabla,
but i think are wrong as they don't currently work, and were not actually being hit.
See invenia/Nabla.jl#192

I think we need to unwrap the Transpose and Adjoint objects before passing to GEMM
and/or maybe consider the double transpose case carefully.

I may split this PR in two and just get the simple case that has no Transposes or Adjoints in first, since Nabla actually uses that.

@oxinabox oxinabox marked this pull request as draft October 2, 2020 19:13
@oxinabox oxinabox added inplace for things relating to inplace accumulation missing rule labels Oct 2, 2020
@oxinabox
Copy link
Member Author

oxinabox commented Oct 5, 2020

I think this can all be cone much easier now with 5 arg mul!,
though that requires Julia 1.3.
and is not in Compat.jl yet JuliaLang/Compat.jl#681

I think we might even be able to just replace the generic code for the rrule/frule of matmul with that,
and not need this special casing at all

test/runtests.jl Outdated Show resolved Hide resolved
@oxinabox
Copy link
Member Author

oxinabox commented Oct 8, 2020

mul! its all the cases except 1,
and I have a PR open for that one.
JuliaLang/julia#37930
and its not one we are too worries about really

@oxinabox
Copy link
Member Author

oxinabox commented Oct 8, 2020

This is basically done, but can't be merged until we are testing accumulation properly, as it could just be wrong
JuliaDiff/ChainRulesTestUtils.jl#59

@oxinabox oxinabox changed the title WIP: Efficient Inplace Strided MatMul Efficient Inplace Strided MatMul Oct 8, 2020
@oxinabox
Copy link
Member Author

oxinabox commented Oct 8, 2020

Another advantage of this over the Nabla code is that we don't have to worry about the tests not hitting the right one, because of the fact that this is the most general one.

@oxinabox oxinabox marked this pull request as ready for review October 13, 2020 10:46
@@ -4,6 +4,7 @@ version = "0.7.23"

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Version number needs to be bumped

@oxinabox oxinabox changed the title Efficient Inplace Strided MatMul Efficient Inplace Accumulating Strided MatMul Oct 16, 2020
Copy link
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good. Just a couple of comments.

@@ -19,23 +19,47 @@ end
##### `*`
#####

function rrule(::typeof(*), A::AbstractMatrix{<:Real}, B::AbstractMatrix{<:Real})
function rrule(::typeof(*), A::AbstractMatrix{<:Number}, B::AbstractMatrix{<:Number})
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sethaxen pointed out that we probably want to restrict these to Real and Complex numbers.

test/rulesets/Base/arraymath.jl Show resolved Hide resolved
test/rulesets/Base/arraymath.jl Outdated Show resolved Hide resolved
@oxinabox oxinabox merged commit 1844537 into master Oct 20, 2020
@oxinabox oxinabox deleted the ox/strided branch October 20, 2020 15:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
inplace for things relating to inplace accumulation missing rule
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants