Skip to content

Distributed GEMM #302

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.

Sign up for GitHub

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

Jump to bottom
Merged
panditsa merged 27 commits into from
Oct 3, 2025
Merged

Distributed GEMM #302

panditsa merged 27 commits into from
Oct 3, 2025

Conversation

@panditsa
Copy link
Contributor

@panditsa panditsa commented Sep 16, 2025 •
edited
Loading

Enable distributed matrix multiplication operations across multiple GPU devices in wave, where the distribution factor for each dimension is managed through DeviceConstraint

Key Changes

  • In the host_codegen.py, we split the input tensors across devices based on device constraint, dispatch computation to each device, and merge results back into the original tensor shape
  • We now use a HostSignature class to manage full problem-size buffers (e.g., 1024x8192 matrix), while KernelSignature handles per-device tile buffers (e.g., 512x4096 tiles for 2x2 distribution). Non-distributed workloads have no impact, because in absence of DeviceConstraint, HostSignature is the same as KernelSignature
  • In the host_utils.py, added functions that help split input and output tensors per device, manage tensor distribution across devices using device_constraint_map.
  • Updated distributed GEMM template to accept device_m and device_n parameters, with test coverage for distribution factors from 1x1 up to 4x2 across problem sizes up to 4096x20480x2560
  • Use MultiDeviceLaunchable class to orchestrate execution across multiple GPU devices

@panditsa panditsa force-pushed the sanket/dist_gemm branch 7 times, most recently from 985e5cb to 7e97598 Compare September 18, 2025 04:06
@panditsa panditsa marked this pull request as ready for review September 19, 2025 04:06
@panditsa panditsa requested a review from Hardcode84 September 19, 2025 04:06
@panditsa panditsa force-pushed the sanket/dist_gemm branch 9 times, most recently from 1491a50 to 8f64897 Compare September 22, 2025 16:58
Hardcode84
Hardcode84 reviewed Sep 25, 2025
View reviewed changes
Copy link
Contributor

@Hardcode84 Hardcode84 left a comment

Choose a reason for hiding this comment

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

Any lit test for codegen changes? Also, please fix merge conflicts.

@Hardcode84
Copy link
Contributor

Hardcode84 commented Sep 30, 2025

Please fix the DCO and merge conflicts.

@panditsa panditsa force-pushed the sanket/dist_gemm branch 4 times, most recently from f9a9485 to 875041d Compare October 2, 2025 19:24
@panditsa
adding zeroes for device_dimension in kernel emitter
baca7a8 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
device layout in host_codegen
0542735 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
WIP: host_codegen changes
170fb20 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix bad merge on host_codegen
7d5249c 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
host_codegen create a host_signature class
52f8c56 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
using multi_device_launchable in compile
d0332af 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
panditsa and others added 19 commits October 2, 2025 12:27
@panditsa
lint emitter
fd4008c 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
remove num_devices arg from compile.py
187d2f4 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
host_codegen fix single device usecase
8b6fbab 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
working dist_gemm and test
043594a 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
[Dist.] Async support in host_codegen
6cdae19 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
[Dist.] Working async GEMM
fb7c1bb 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add host_utils and cleanup host_codegen
9934c17 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
cleaned up host codegen
e8d3932 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix compile.py
2d3e915 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
use slice_cache to reuse the slices
04f4a3d 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add shapes to enable upto 8 devices
f1d0aca 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix the dispatchop error when using only one device
fa51530 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
dist_gemm kernel gen and test
c5b315c 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add dynamic arguments to dispatchop
c9afcd1 
Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix n-dimensional split
d5b3685 
Side-effect, device constraint on multiple dimensions functional

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
no async support, doesn't work when running in a loop
3e46e06 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix compile_to_mlir flow
a1d67d8 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add lit test
41a7b10 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
Update codegen.py to add newline
875041d 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Hardcode84
Hardcode84 approved these changes Oct 2, 2025
View reviewed changes
@panditsa
add header
fc6db29 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add require_gpus marker
7f14a43 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa panditsa merged commit 6d60f87 into iree-org:main Oct 3, 2025
19 checks passed
Megan0704-1 pushed a commit to Megan0704-1/wave that referenced this pull request Oct 28, 2025
@panditsa @Megan0704-1
Distributed GEMM (iree-org#302)
fe9c510 
Enable distributed matrix multiplication operations across multiple GPU
devices in wave, where the distribution factor for each dimension is
managed through DeviceConstraint

### Key Changes

- In the host_codegen.py, we split the input tensors across devices
based on device constraint, dispatch computation to each device, and
merge results back into the original tensor shape
- We now use a `HostSignature` class to manage full problem-size buffers
(e.g., 1024x8192 matrix), while `KernelSignature` handles per-device
tile buffers (e.g., 512x4096 tiles for 2x2 distribution).
Non-distributed workloads have no impact, because in absence of
`DeviceConstraint`, `HostSignature` is the same as` KernelSignature`
- In the `host_utils.py`, added functions that help split input and
output tensors per device, manage tensor distribution across devices
using `device_constraint_map`.
- Updated distributed GEMM template to accept` device_m` and `device_n`
parameters, with test coverage for distribution factors from 1x1 up to
4x2 across problem sizes up to 4096x20480x2560
- Use MultiDeviceLaunchable class to orchestrate execution across
multiple GPU devices

---------

Signed-off-by: Sanket Pandit <sanketp@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Hardcode84 Hardcode84 Hardcode84 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@panditsa @Hardcode84