LLVM Discussion Forums

Memrefs and maps for tiling

This topic may have been covered before, but I can’t find answers. I am trying to generate code for a device where memory needs padding in the lowest dimensions. Memrefs and maps seems uniquely suited for this:

#paddedMap = affine-map<(d0, d1) -> (d0 floordiv 32, d1 floordiv 32, d0 mod 32, d1 mod 32)
%1 = alloc() : memref<62x61xf32, #paddedMap>

Using the maps is perfect as it let us preserve the original data dimension (i.e. 62x61, data that matters), while also carrying with the memrefs the desired data layout and memory size needed for data allocation (i.e. 2x2x32x32).

All is well until lowering to LLVM where I hit this assert:

Assertion failed: (isStrided(type) && "Non-strided layout maps must have been normalized away"), function convertMemRefSignature, file /Users/alexe/MLIR/llvm-project/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp, line 237.

The example is simple (and actually has no padding as dimensions match):

#pad = affine_map<(i) -> (i floordiv 4, i mod 4)>
func @matmul(%A: memref<16xf64, #pad>) {
  affine.for %arg3 = 0 to 16 {
        %a = affine.load %A[%arg3] : memref<16xf64, #pad>
        %p = mulf %a, %a : f64
        affine.store %p, %A[%arg3] : memref<16xf64, #pad>
  }
  return
}

when compiled with various different options (did not find a combination that works):

mlir-opt simple.mlir --simplify-affine-structures --lower-affine --convert-std-to-llvm

I have looked at some of the test codes, they often seems to flatten N-dimensional tensors to a 1D in maps, not sure why that is needed.

Is my error due to missing functionality, erroneous input, erroneous optimization sequence?

I think the expressivity of N dimensional memrefs with maps to deal with layout is a very strong abstraction. I hope we can make it work for this case, which hardly seems to be a corner case.

Thanks

Just before you lower to LLVM, please call normlizeMemRef (there is no pass for it). It will rewrite a memref’s access functions such that its layout map becomes identity. It can then be handled by the std to LLVM lowering.

Update: I notice that -simplify-affine-structures already calls normalizeMemref for all AllocOps. However, here, your memref is as a function argument. Had your memref been the result of an alloc, that would have worked. If you need to keep your function in that form (say it’s not inlined etc.), normalizeMemRef will have to be reused to create a new interprocedural / module pass to perform the necessary rewrites for the function argument case.

Thanks @bondhugula, using an alloc made it work. Will look into adding a call to normalizememRef if needed. Much appreciated

Great. The case where you have such memrefs being passed to functions should be straightforward to extend to as well if needed - the replacement mechanics are the same. It’s just that the function argument will have to be replaced and non-dereferencing uses (in call ops) could simply be replaced (normalizeMemref would bail out currently on such uses). And finally, it could be put into a separate -normalize-memref pass, which would have to be a pass on ModuleOp (because -simplify-affine-structures is a pass on FuncOps).