Skip to content

Conversation

@cferry-AMD
Copy link
Collaborator

@cferry-AMD cferry-AMD commented May 2, 2025

llvm#118478 changed the definition of what a symbol is: it made it so that an operation that is Pure and whose operands are valid symbols, returns a valid symbol as result. This causes linalg.index, that matches this definition, to become a symbol -- contradicting its lowering to a loop counter by another pass convert-linalg-to-affine-loops.
Example (introduced as a test by this PR):

#set = affine_set<(d0) : (-d0 + 5 >= 0)>

func.func @linalg_index_affine_if(%in: tensor<10xf32>) -> tensor<10xf32> {
  %empty = tensor.empty() : tensor<10xf32>
  %out = linalg.generic {
      indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], 
      iterator_types = ["parallel"]}
    ins(%in : tensor<10xf32>)
    outs(%empty : tensor<10xf32>) {
      ^bb0(%a: f32, %b: f32):
        %c42f = arith.constant 42.0 : f32
        %i = linalg.index 0 : index
        %ret = affine.if #set(%i) -> f32 {
          affine.yield %a : f32
        } else {
          affine.yield %c42f : f32
        }
        linalg.yield %ret : f32
    } -> tensor<10xf32>
  return %out : tensor<10xf32>
}

The upstream behavior canonicalizes the set (attribute of the affine.if) into:

#set = affine_set<()[s0] : (-s0 + 5 >= 0)>

where with the current PR, it stays as

#set = affine_set<(d0) : (-d0 + 5 >= 0)>

This PR allows to directly substitute linalg.index operands to affine ops with loop counters, without causing an error like operand cannot be used as a symbol due to how the map/set gets canonicalized with linalg.index.

Once downstream no longer needs it, this PR should be reverted, and let upstream take their stance between Affine and Linalg compatibility.

This PR also reverts the changes from upstream affecting the tests where linalg.index is used (all the tests touched by this PR except affine-if.mlir).

@cferry-AMD cferry-AMD requested review from TinaAMD and josel-amd May 5, 2025 07:45
@cferry-AMD cferry-AMD marked this pull request as ready for review May 5, 2025 07:45
@cferry-AMD cferry-AMD merged commit 065d0c0 into feature/fused-ops May 5, 2025
4 checks passed
@cferry-AMD cferry-AMD deleted the corentin.fix_affine_if branch May 5, 2025 11:32
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