Affine.for iteration arguments

MLIR
1

Hello,

I am a Master student in Computer Science and Engineering and I am investigating affine dialect and more closely affine.for operation and the loop optimization possibilities. I have encountered 2 issues related to the iteration arguments that I have not been able to find the explanations for. However, if there is already a similar topic opened that I wasn’t able to find, sorry for the duplication.

  1. Why is -affine-loop-invariant-code-motion pass moving the operations that don’t depend on the induction variable, but they do depend on other iteration variables? Here is an example that is failing when pass is applied as such operation is moved out of the loop although it depends on the iteration argument:
  func @example1(%arg0: memref<1000xi32>, %arg1: memref<1000xi32>) {
    %c2 = constant 2 : index
    %0 = affine.apply affine_map<(d0) -> (d0 - 2)>(%c2)
    %1 = affine.for %arg2 = 2 to 1000 iter_args(%arg3 = %0) -> index {
      %3 = affine.load %arg0[%arg3] : memref<1000xi32>
      affine.store %3, %arg1[%arg2] : memref<1000xi32>
      %4 = affine.apply affine_map<(d0) -> (d0 - 2)>(%arg2)
      affine.yield %4 : index
    }
    %c999 = constant 999 : index
    %2 = affine.load %arg0[%1#0] : memref<1000xi32>
    affine.store %2, %arg1[%c999] : memref<1000xi32>
    return
  }
  • Here is the error I am getting:
 note: see current operation: %1 = "affine.load"(%arg0, %arg3) {map = affine_map<(d0) -> (d0)>} : (memref<1000xi32>, index) -> i32
 note: operand defined as a block argument (block #0 in a child region)
  1. Why are top-level instructions labeled as symbols and thus can be used as operands to the affine.load or affine.store operations? I had issues with the following example (example1). The error I get is that affine.load cannot use the affine.for loop result as it cannot be determined that it is a valid dimension or symbol. However, if I try with the previous example (example2), there are no issues. The only difference is that in the example, affine.for operation whose result we are using is a top-level operation. Why is there a difference between these two situations?
  func @example2(%arg0: memref<1000xi32>, %arg1: memref<1000xi32>) {
    %c2 = constant 2 : index
    affine.for %arg4 = 1 to 10 {
      %0 = affine.apply affine_map<(d0) -> (d0 - 2)>(%c2)
      %1 = affine.for %arg2 = 2 to 1000 iter_args(%arg3 = %0) -> index {
        %3 = affine.load %arg0[%arg3] : memref<1000xi32>
        affine.store %3, %arg1[%arg2] : memref<1000xi32>
        %4 = affine.apply affine_map<(d0) -> (d0 - 2)>(%arg2)
        affine.yield %4 : index
      }
      %c999 = constant 999 : index
      %2 = affine.load %arg0[%1#0] : memref<1000xi32>
      affine.store %2, %arg1[%c999] : memref<1000xi32>
    }
    return
  }

Additionally, there is no problem when using the iteration arguments that are not the induction variable as dimension (This check is passed). Is this the desired behavior?

I tried to make the examples as simple as possible which is why they don’t really do anything meaningful, but I hope they help clarify the questions. Thank you in advance!

2

A symbol is an unknown constant for the region of interest - in this case, the 2-dimensional loop nest you have. In your second example, it has a potentially different value for each iteration of %arg4 and is thus not a symbol. (It may turn out to be a constant after simplification like 997 in this case, but until then, it’s not a symbol.) In your first example, it’s a constant (albeit unknown). Note that if you change the second affine.load to load, the IR becomes valid.

3

Thank you for the answer! What about the first question concerning -affine-loop-invariant-code-motion? Should this method also check for the iteration arguments other than the induction variable or is this desired behavior?

4

The pass wasn’t extended to account for iter args when the latter was introduced. Can you please file a bug report?

1 Like
5

I have another question regarding this topic.

    %2:2 = affine.for %arg2 = 2 to 1000 iter_args(%arg3 = %0, %arg4 = %1) -> index, index {
     ....
    %4 = affine.apply affine_map<(d0) -> (d0 - 2)>(%arg2)
    affine.yield %4, %arg3 : index, index
  }

Is there a way to yield one of the iteration arguments to the next iteration? Like in this example, if we want to pass %0 to the second iteration as %arg4. This code currently would basically pass the value of %arg3 in iteration 3 to %arg4 in iteration 3. How would I pass the value of %arg3 in iteration 2 to %arg4 in iteration 3? Thank you!

6

The operands of affine.yield become the values of iter_args on the next iteration, or the results of the loop after the last iteration. So your code seems to be doing exactly what you want. The values “assigned” to iter_args are only there for the first iteration.

7

Thank you for the answer!