MoveLoadToAffine in affine-cfg doesn't work for affine maps of indices?

[jumerckx]

edit: I added a smaller MWE in the comment below.

Hi,

I'm banging my head against the wall trying to figure out why the second memref.load and memref.store in this IR aren't rewritten to their affine counterpart when running Polygeist's affine-cfg pass. AFAIK, this should be possible?

// polygeist-opt --allow-unregistered-dialect --affine-cfg
#map = affine_map<()[s0] -> (s0 + 2)>
#map1 = affine_map<()[s0] -> (s0 + 1)>
module {
  func.func @example1(%arg0: memref<?xf64, strided<[1], offset: ?>>, %arg1: memref<?x?xf64, strided<[1, ?], offset: ?>>, %arg2: memref<?xf64, strided<[1], offset: ?>>) attributes {llvm.emit_c_interface} {
    %cst = arith.constant 0.000000e+00 : f64
    %c1 = arith.constant 1 : index
    %c0 = arith.constant 0 : index
    %dim = memref.dim %arg0, %c0 : memref<?xf64, strided<[1], offset: ?>>
    %dim_0 = memref.dim %arg1, %c1 : memref<?x?xf64, strided<[1, ?], offset: ?>>
    %dim_1 = memref.dim %arg2, %c0 : memref<?xf64, strided<[1], offset: ?>>
    affine.for %arg3 = 1 to #map()[%dim] {
      %0 = affine.for %arg4 = 1 to #map1()[%dim_0] iter_args(%arg5 = %cst) -> (f64) {
        %2 = memref.load %arg1[%arg3, %arg4] : memref<?x?xf64, strided<[1, ?], offset: ?>> // DOES get rewritten
        %3 = affine.delinearize_index %arg4 into (%dim_1) : index
        %4 = memref.load %arg2[%3] : memref<?xf64, strided<[1], offset: ?>> // doesn't get rewritten
        %5 = arith.mulf %2, %4 : f64
        %6 = arith.addf %arg5, %5 : f64
        affine.yield %6 : f64
      }
      %1 = affine.delinearize_index %arg3 into (%dim) : index
      memref.store %0, %arg0[%1] : memref<?xf64, strided<[1], offset: ?>> // doesn't get rewritten
    }
    return
  }
}

In the code, the pass fails at isValidIndex which is called from here. Could this perhaps be fixed by adding another recursive check for delinearize_index in that function?

When I leave the delinearize_index out (it's a no-op anyway in this case) the memory operations do get raised.
I tried further reducing this example but below IR does work as expected, so it's not only the presence of delinearize_index that's problematic.

module {
  func.func @example1(%arg0: memref<?xf64, strided<[1], offset: ?>>, %arg1: i64) -> f64 attributes {llvm.emit_c_interface} {
    %c1 = arith.constant 1 : index
    %c0 = arith.constant 0 : index
    %dim = memref.dim %arg0, %c0 : memref<?xf64, strided<[1], offset: ?>>
    %0 = affine.delinearize_index %c1 into (%dim) : index
    %1 = memref.load %arg0[%0] : memref<?xf64, strided<[1], offset: ?>>
    return %1 : f64
  }
}

I hope I'm not missing anything too obvious, looking forward to hear if this can be fixed. If you point me to where the change needs to happen, I'll gladly create a PR :).

Jules

[jumerckx]

I was able to reduce the example quite a bit, but I'm still unable to see why this doesn't simply work.

// polygeist-opt --allow-unregistered-dialect --affine-cfg
#map = affine_map<(d0) -> (d0 floordiv 10)>
#map1 = affine_map<(d0) -> (d0 mod 10)>
module {
  func.func @sum(%arg0: memref<10x2xf64>) -> f64 attributes {llvm.emit_c_interface} {
    %cst = arith.constant 0.000000e+00 : f64
    %0 = affine.for %arg1 = 0 to 10 iter_args(%arg2 = %cst) -> (f64) {
      %2 = affine.apply #map(%arg1)
      %3 = affine.apply #map1(%arg1)
      %4 = memref.load %arg0[%3, %2] : memref<10x2xf64>
      %5 = arith.addf %arg2, %4 : f64
      affine.yield %5 : f64
    }
    return %0 : f64
  }
}