How to pass a set of MemRefs to an MLIR function from the C++ side

MLIR
1

I’m trying to pass a set of MemRefs from the C++ side to the MLIR function. The MemRef I want to pass is something like the following code. I cast MemRefs to Unranks and composed them in one dim MemRef, similar to an array of UnrankMemRefs

func.func @foo(..., %arg1: memref<16xmemref<*xf32>>) -> ... {
    ...
    %0 = memref.load %arg1[%c1] : memref<2xmemref<*xf32>>
    %cast = memref.cast %0 : memref<*xf32> to memref<3xf32>
    %1 = memref.load %arg1[%c0] : memref<2xmemref<*xf32>>
...

The MemRef and UnrankMemRef definitions I used:

template<typename T, size_t N>
struct MemRefDescriptor {
  T* allocated;
  T* aligned;
  intptr_t offset;
  intptr_t sizes[N];
  intptr_t strides[N];
};

template<typename T>
struct UnrankedMemRefType {
  int64_t rank;
  void* descriptor;
};

from the C++ side, I use the code:

_mlir_ciface_foo(..., MemRefDescriptor<UnrankedMemRefType<float>, 1>* params);

However, I get memory errors when running the program, so I’m asking if it’s the right way to nest UnrankMemRef in MemRef? Or is there any other way to pass a set of MemRefs from C++ Side to MLIR?

2

Have you looked at the unit tests? llvm-project/mlir/unittests/ExecutionEngine at main · llvm/llvm-project · GitHub

3

The example inside the unit test doesn’t seem to include UnrankMemRefType?I’m not really sure if doing something like memref<16xmemref<*xf32>> is allowed by MLIR.

4

Mmmm I remembered writing a unit-tests a while ago, turns out I didn’t land it. Maybe the patch is salvageable though (contribution welcome here!):
https://reviews.llvm.org/D96397#change-2RDEmSq5iVxi

5

Thanks, I’m trying to follow this patch to modify the code. For passing a set of MemRefs to an MLIR function, it might be faster to flatten all the MemRefs into a 1D array, and then use subview and reinterpret_cast to expand them into the original shape within the MLIR function.