Add alternative strategy for batched matrix multiplication #51
Conversation
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
robertknight
force-pushed
the
optimize-batched-matmul
branch
from
b142348
to
0a7344f
Compare
Closed
robertknight
force-pushed
the
optimize-batched-matmul
branch
from
0a7344f
to
74f13d9
Compare
Batched matrix multiplication was handled by prepacking one or neither of the inputs, depending on how often each is re-used, and then performing one `gemm` call per matrix in the output shape. This can be inefficient the LHS input has a small number of rows. For example in [1], the LHS / "A" input is a row vector. In the case where the "A" input is a batch and the "B" input is a single matrix, the "A" input can be reshaped so a single gemm call can be used, with the output reshaped afterwards to restore the batch dimensions. Implement this alternate approach and add a simple benchmark for batched matmul. [1] #50
Refactor MatMul tests into a single table-driven test that has cases for when neither, one or both of the inputs is a batch. Also add tests for various invalid inputs.
robertknight
force-pushed
the
optimize-batched-matmul
branch
from
abeda07
to
e33f6a8
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Previously batched matrix multiplication was handled by prepacking one or neither of the inputs, depending on how often each is re-used, and then performing one
gemmcall per matrix in the output shape. This is inefficient if theAinput has only a small number of rows (as in #50). This PR implements a new strategy for theMatMuloperator whenAis a batch andBis a single matrix, by reshaping the inputs so that instead of many low-arithmetic intensitygemmcalls, a single higher-arithmetic intensity call is performed. The output is then reshaped to restore the batch dimensions.Testing with the benchmark added here and slight variations, the new method is a big improvement when M <= 8, a modest win for M ~ 8-24 and is roughly even, or a very slight win after that. The AVX kernel has MR=6, so this seems as-expected.
See #50