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

Add reduction_size and broadcast_size attributes to XeTile.reduction and XeTile.broadcast #996

Merged
merged 7 commits into from
Jan 13, 2025

Conversation

Jianhui-Li
Copy link
Contributor

@Jianhui-Li Jianhui-Li commented Jan 9, 2025

Add broadcast_size to XeTile.reduction and XeTile.broadcast to support partial reduction.

We found two use case for partial reduction:


1. Support the staged work group reduction - first reduce within workgroup and then reduce across workgroup.
 %18 = math.exp %17#0 {map = #xetile.wg_map<sg_layout = [8, 4], sg_data = [32, 32]>} : vector<256x128xf32>
%19= vector.shape_cast %18 {map = #xetile.wg_map<sg_layout = [8, 4], sg_data = [32, 32]>} : vector<256x128xf32> to vector<8x32x128xf32>
%20 = vector.multi_reduction <add>, %19, %cst_0 {map = #xetile.wg_map<sg_layout = [8, 4], sg_data = [1, 32]>} [1] : vector<8x32x128xf32> to vector<8x128xf32>

        => 
 
 %18 = math.exp %17#0 {map = #xetile.wg_map<sg_layout = [8, 4], sg_data = [32, 32]>} : vector<256x128xf32>
 %20 = xetile.reduction <add>, %18, %cst_0 {map = #xetile.wg_map<sg_layout = [8, 4], sg_data = [1, 32]>} [1] {$reduction_size = [32]} : vector<256x128xf32> to vector<8x128xf32>

Support MXFP reduction - the reduction only reduce 32 elements to one element, so not reduce the whole dimension to one element.

This PR also fixes a few name inconsistency: tile_broadcast to broadcast, tile_reduce to reduction, tile_transpose to transpose, atomic_rmw_tile to atomic_rmw, tile_conv_layout to convert_layout.

Please review these guidelines to help with the review process:

  • Have you provided a meaningful PR description?
  • Have you added a test, a reproducer, or a reference to an issue with a reproducer?
  • Have you tested your changes locally for CPU and GPU devices?
  • Have you made sure that new changes do not introduce compiler warnings?
  • If this PR is a work in progress, are you filing the PR as a draft?
  • Have you organized your commits logically and ensured each can be built by itself?

Add broadcast_size to xeTile.broadcast operation
to make name consistent: tile_broadcast to broadcast, tile_reduce to reduction, tile_transpose to transpose, atomic_rmw_tile to atomic_rmw.
Copy link
Contributor

@nbpatel nbpatel left a comment

Choose a reason for hiding this comment

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

LGTM. Can we also change xetile.tile_conv_layout to xetile.conv_layout

tile_conv_layout to convert_layout
```mlir
%vector_a = xetile.tile_reduce <add> %vector_b [1]: vector<64x32xfloat> into vector<64x1xfloat>
%vector_a = xetile.reduction <add> %vector_b [0] {$reduction_size=32}: vector<64x64xfloat>, vector<2x64xfloat> into vector<2x64xfloat>
Copy link
Contributor

Choose a reason for hiding this comment

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

do we need to mention vector<2x64xfloat> twice in this operation?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

fixed. thanks

@silee2 silee2 merged commit 5dbeec7 into intel:main Jan 13, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants