[MLIR][memref] The Normalizing Memref on removing extra input arguments on`memref.load`

[tavakkoliamirmohammad]

Hi. If you want to normalize the memref layout that reduces the dimensionality, the -normalize-memrefs pass will work perfectly if there is affine.load, but in the case that the operation is memref.load, the code fails to reduce the inputs to memref.load.

Complier explorer link: https://godbolt.org/z/zYT7jxs9n

#map0 = affine_map<(i,k) -> (2*(i mod 2) + (k mod 2) + 4 * (i floordiv 2) + 8*(k floordiv 2))>
#map1 = affine_map<(k,j) -> ((k mod 2) + 2 * (j mod 2) + 8 * (k floordiv 2) + 4*(j floordiv 2))>
#map2 = affine_map<(i,j) -> (4*i+j)>

func.func @test_norm(%arg0 :  memref<4x4xf32,#map2>) -> () {
    %0 = memref.alloc() : memref<4x8xf32,#map0>
    %1 = memref.alloc() : memref<8x4xf32,#map1>
    %2 = memref.alloc() : memref<4x4xf32,#map2>

    %cst = arith.constant 3.0 : f32
    %cst0 = arith.constant 0 : index

    affine.for %i = 0 to 4 {
        affine.for %j = 0 to 8 {
                affine.for %k = 0 to 8 {
                    %a = memref.load %0[%i, %k] : memref<4x8xf32,#map0>

                    %b = memref.load %1[%k, %j] :memref<8x4xf32,#map1>
                    %c = memref.load %2[%i, %j] : memref<4x4xf32,#map2>

                    %3 = arith.mulf %a, %b : f32 
                    %4 = arith.addf %3, %c : f32
                    affine.store %4, %arg0[%i, %j] : memref<4x4xf32,#map2>
            }
        }
    }
    return
}

Reproducing the error

mlir-opt -normalize-memrefs test.mlir

The error

# For more information see the output window
MLIR opt (trunk) - 908ms
<source>:16:26: error: 'memref.load' op incorrect number of indices for load, expected 1 but got 2
                    %a = memref.load %0[%i, %k] : memref<4x8xf32,#map0>
                         ^
<source>:16:26: note: see current operation: %5 = "memref.load"(%0, %arg1, %arg3) : (memref<32xf32>, index, index) -> f32
Compiler returned: 1

[DarshanRamakant]

This is happening because the memref memory operations are not handled when there is a diminsionality ( airity) change.
I have created a patch for the issue :
#107771

[llvmbot]

@llvm/issue-subscribers-mlir-affine

Author: Amir Mohammad Tavakkoli (tavakkoliamirmohammad)

Hi. If you want to normalize the `memref` layout that reduces the dimensionality, the `-normalize-memrefs` pass will work perfectly if there is `affine.load`, but in the case that the operation is `memref.load`, the code fails to reduce the inputs to `memref.load`.

Complier explorer link: https://godbolt.org/z/zYT7jxs9n

#map0 = affine_map<(i,k) -> (2*(i mod 2) + (k mod 2) + 4 * (i floordiv 2) + 8*(k floordiv 2))>
#map1 = affine_map<(k,j) -> ((k mod 2) + 2 * (j mod 2) + 8 * (k floordiv 2) + 4*(j floordiv 2))>
#map2 = affine_map<(i,j) -> (4*i+j)>

func.func @<!-- -->test_norm(%arg0 :  memref&lt;4x4xf32,#map2&gt;) -&gt; () {
    %0 = memref.alloc() : memref&lt;4x8xf32,#map0&gt;
    %1 = memref.alloc() : memref&lt;8x4xf32,#map1&gt;
    %2 = memref.alloc() : memref&lt;4x4xf32,#map2&gt;

    %cst = arith.constant 3.0 : f32
    %cst0 = arith.constant 0 : index

    affine.for %i = 0 to 4 {
        affine.for %j = 0 to 8 {
                affine.for %k = 0 to 8 {
                    %a = memref.load %0[%i, %k] : memref&lt;4x8xf32,#map0&gt;

                    %b = memref.load %1[%k, %j] :memref&lt;8x4xf32,#map1&gt;
                    %c = memref.load %2[%i, %j] : memref&lt;4x4xf32,#map2&gt;

                    %3 = arith.mulf %a, %b : f32 
                    %4 = arith.addf %3, %c : f32
                    affine.store %4, %arg0[%i, %j] : memref&lt;4x4xf32,#map2&gt;
            }
        }
    }
    return
}

Reproducing the error

mlir-opt -normalize-memrefs test.mlir

The error

# For more information see the output window
MLIR opt (trunk) - 908ms
<source>:16:26: error: 'memref.load' op incorrect number of indices for load, expected 1 but got 2
                    %a = memref.load %0[%i, %k] : memref<4x8xf32,#map0>
                         ^
<source>:16:26: note: see current operation: %5 = "memref.load"(%0, %arg1, %arg3) : (memref<32xf32>, index, index) -> f32
Compiler returned: 1