MemRef type in interface

[kumasento]

It is a follow-up to #44, and maybe I've asked about this earlier: could we generate function interface that has MemRef types with explicit dimensionality? For example, take the code from #44:

module attributes {llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu"}  {
  func @matmul(%arg0: memref<?x400xf32>, %arg1: memref<?x300xf32>, %arg2: memref<?x300xf32>) {
    affine.for %arg3 = 0 to 200 {
      affine.for %arg4 = 0 to 300 {
        affine.for %arg5 = 0 to 400 {
          %0 = affine.load %arg0[%arg3, %arg5] : memref<?x400xf32>
          %1 = affine.load %arg1[%arg5, %arg4] : memref<?x300xf32>
          %2 = mulf %0, %1 : f32
          %3 = affine.load %arg2[%arg3, %arg4] : memref<?x300xf32>
          %4 = addf %3, %2 : f32
          affine.store %4, %arg2[%arg3, %arg4] : memref<?x300xf32>
        }
      }
    }
    return
  }
}

How could we make the interface of @matmul to be:

func @matmul(%arg0: memref<200x400xf32>, %arg1: memref<400x300xf32>, %arg2: memref<200x300xf32>)

Thanks!

[wsmoses]

Yes indeed we can. The reason for this is that the Clang AST actually represents this as *[300 x double], stripping the last dimension. However, we don't need to adhere to Clang's precise definitions so long as we remain consistent

[kumasento]

Yes indeed we can. The reason for this is that the Clang AST actually represents this as *[300 x double], stripping the last dimension. However, we don't need to adhere to Clang's precise definitions so long as we remain consistent

Thanks! Would you mind pointing me how I can lift this limitation? I could create a patch for that, maybe add a switchable option.

[wsmoses]

Yeah for sure. Was thinking about this a little bit and I think the easiest thing to do would be:

1. When querying the type of an argument where this applies, manually enforce the correct rule by say expanding here: https://github.com/wsmoses/MLIR-GPU/blob/d18dce36513bebef5060adc786115a85593f461d/mlir/tools/mlir-clang/Lib/clang-mlir.cc#L2690
2. Whenever calling a function do a check for the type to do a memref cast from the default clang convention to our new convention (allowing us to not have to change the definition of how the rest of the AST is lowered)
3. When starting to emit a function, immediately cast to the clang convention pointer.

In theory this should be sufficient to be correct when combined with everything else and maintain the calling convention in the function declaration. Hopefully it should also allow existing optimizations to propagate versions of the type with the extra information whenever its used.