Extend multi-dimensional batched matmul support #2180
Merged
Conversation
This file contains hidden or 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
phschaad
approved these changes
Oct 18, 2025
Collaborator
phschaad
left a comment
phschaad
left a comment
There was a problem hiding this comment.
Nice addition, looks good to me
sophieblock
pushed a commit
to sophieblock/dace
that referenced
this pull request
Oct 20, 2025
# Extended Batched Matrix Multiplication Support ## Summary Extend batched matrix multiplication to support N-dimensional tensors with NumPy-style broadcasting across all implementations (Pure, MKL, OpenBLAS, cuBLAS). ## Changes - **N-D tensor support**: Handles tensors with arbitrary batch dimensions (e.g., `[12, 2, 64, 64] @ [12, 2, 64, 128]`) - **Broadcasting**: Supports broadcasting patterns like `[b, m, k] @ [k, n]` and `[m, k] @ [b, k, n]` - **Batch flattening**: Multi-dimensional batches are flattened internally for efficient BLAS operations ## New Capabilities ```python # 3D broadcasting [b, m, k] @ [k, n] → [b, m, n] [m, k] @ [b, k, n] → [b, m, n] # 4D batched matmul [12, 2, 64, 64] @ [12, 2, 64, 128] → [12, 2, 64, 128] [12, 2, m, k] @ [k, n] → [12, 2, m, n] # with broadcasting ``` ## Files Modified - `dace/libraries/blas/nodes/matmul.py`: Extended `_get_batchmm_opts()` for N-D tensors and broadcasting - `dace/libraries/blas/nodes/batched_matmul.py`: Updated validation and Pure expansion for dynamic dimensions - `dace/frontend/python/replacements/linalg.py`: Removed 3D tensor check - `tests/library/batched_matmul_test.py`: Added tests for newly supported 3D/4D batched matmuls with broadcasting
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.
Extended Batched Matrix Multiplication Support
Summary
Extend batched matrix multiplication to support N-dimensional tensors with NumPy-style broadcasting across all implementations (Pure, MKL, OpenBLAS, cuBLAS).
Changes
[12, 2, 64, 64] @ [12, 2, 64, 128])[b, m, k] @ [k, n]and[m, k] @ [b, k, n]New Capabilities
Files Modified
dace/libraries/blas/nodes/matmul.py: Extended_get_batchmm_opts()for N-D tensors and broadcastingdace/libraries/blas/nodes/batched_matmul.py: Updated validation and Pure expansion for dynamic dimensionsdace/frontend/python/replacements/linalg.py: Removed 3D tensor checktests/library/batched_matmul_test.py: Added tests for newly supported 3D/4D batched matmuls with broadcasting