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

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

Run the provided MLIR reproducer with polygeist-opt --allow-unregistered-dialect --affine-cfg and inspect lib/polygeist/Passes/AffineCFG.cpp, especially isValidIndex and its call near line 1238. Compare the failing case with the reduced working example and determine how affine.delinearize_index affects validation. Done means the second load and store are rewritten to affine operations without regressing the reduced case.