Canonicalization of affine maps in affine loads and stores

There is canonicalization of affine maps in affine loads and stores already, like in the example below. That it is reflected by addition of “symbol(%arg1)” wrapper in affine store index in compact textual representation, or a change in the affine map (d0) -> (d0) to ()[s0] -> (s0) in verbose representation (-mlir-print-op-generic) see canonicalizePromotedSymbols()@AffineOps.cpp.

func @affine_symbol_add(%arg0 : memref<4xf32>, %arg1: index) {
  %cst = constant 1.0 : f32
  // CHECK: affine.store %{{.*}}, %{{.*}}[symbol(%{{.*}})]
  affine.store %cst, %arg0[%arg1] : memref<4xf32>
  return
}

However there is no inverse transformation. E.g. if parent affine.for is added, the affine map is not changed from ()[s0] -> (s0) to (d0) -> (d0) by canonicalization pass.

func @affine_symbol_not_removed(%arg0 : memref<4xf32>) {
  %cst = constant 1.0 : f32
  affine.for %arg1 = 0 to 4 {
     affine.store %cst, %arg0[symbol(%arg1)] : memref<4xf32>
  }
  return
}

Does inverse canonicalization make sense, would it be reasonable to implement it in MLIR? (I am pretty sure I miss something fundamental). Thank you!

An affine.for IV can’t be used as a symbol. So, the IR would be invalid in the first place with the IV as a symbol. If you are specifically referring to making symbolic identifiers dimensional identifiers, that’s always valid, but there hasn’t been a use case to do that I think. Symbols are unknown constants and it’s always more information for everything if an SSA value is known to be a symbol as opposed to being a dimensional identifier.

1 Like