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

computeSliceParameters: Fix offset when m(0) != 0 #442

Merged
merged 1 commit into from
Jan 15, 2025

Conversation

mgehre-amd
Copy link
Collaborator

For affine maps where m(0) != 0,
like affine_map<(d0) -> (d0 + 3) in

  %generic = linalg.generic
    {indexing_maps = [affine_map<(d0) -> (d0 + 3)>,
                      affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]} ins(%arg0: tensor<9xf32>) outs(%empty : tensor<6xf32>) {
    ^bb0(%in : f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<6xf32>

tiling currently computes the wrong slice offsets. When tiling above example with a size of 3, it would compute

scf.for %i = ...
  %slice = tensor.extract_slice %arg0[%i + 3] [6] [1]
  linalg.generic
    {indexing_maps = [affine_map<(d0) -> (d0 + 3)>,
                      affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]} ins(%slice: tensor<6xf32>)

and thus apply the +3 twice (once in the extract slice and a second time in the linalg.generic).

This PR fixes this to yield an offset of
tensor.extract_slice %arg0[%i] [6] [1] instead.

For affine maps where `m(0) != 0`,
like
```
  %generic = linalg.generic
    {indexing_maps = [affine_map<(d0) -> (d0 + 3)>,
                      affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]} ins(%arg0: tensor<9xf32>) outs(%empty : tensor<6xf32>) {
    ^bb0(%in : f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<6xf32>
```
tiling currently computes the wrong slice offsets.
When tiling above example with a size of 3, it would compute
```
scf.for %i = ...
  %slice = tensor.extract_slice %arg0[%i + 3] [6] [1]
  linalg.generic
    {indexing_maps = [affine_map<(d0) -> (d0 + 3)>,
                      affine_map<(d0) -> (d0)>],
     iterator_types = ["parallel"]} ins(%slice: tensor<6xf32>)
```
and thus apply the `+3` twice (once in the extract slice and a second time in the linalg.generic).

This PR fixes this to yield an offset of
`tensor.extract_slice %arg0[%i] [6] [1]` instead.
@mgehre-amd mgehre-amd force-pushed the matthias.fix_tiling_slice_offset branch from ffe0308 to dcbacfd Compare January 10, 2025 16:14
Copy link
Collaborator

@cferry-AMD cferry-AMD left a comment

Choose a reason for hiding this comment

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

The way I understood this is that it results in a larger slice than necessary. I don't know if we have another choice: since the expressions aren't always affine, but still monotonically increasing, so we can't really afford to change the linalg op's access expression by subtracting the constant there (which would work if the expressions were truly affine). Do we share the same point of view?

@mgehre-amd
Copy link
Collaborator Author

The way I understood this is that it results in a larger slice than necessary. I don't know if we have another choice: since the expressions aren't always affine, but still monotonically increasing, so we can't really afford to change the linalg op's access expression by subtracting the constant there (which would work if the expressions were truly affine). Do we share the same point of view?

Yes, one could have "better" solutions if we were allowed to change the linalg affine_maps.

@mgehre-amd mgehre-amd merged commit 2f9fa50 into feature/fused-ops Jan 15, 2025
4 checks passed
@mgehre-amd mgehre-amd deleted the matthias.fix_tiling_slice_offset branch January 15, 2025 08:14
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.

3 participants