Add C++ helpers to manage Unranked Memref in native code

[mehdi_amini]

Summary

This added an OwningUnrankedMemRef which wraps a UnrankedMemRefType
descriptor. New helpers method to initialize and index the descriptor, as
well as an iterator (UnrankedMemrefIterator) are also provided.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

401

Note to self: refactor this to share the implementation with the ranked iterator.

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

CRunner utils is meant to run on HW that does not have C++ runtime library.
In particular, uses of vector are prohibited and array is used instead.
Does this file fit the bill?

387

An implementation that is easier to follow and reuses common utils is to just keep the offset, add 1 to it and delinearize the indices in the stride basis.

422

That's not ok for the CRunnerUtils.
Maybe put those in RunnerUtils?

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

We should probably refactor all this, it seems all quite not well organized to me right now, including mixing the definition of the descriptors with random set of utilities that look quite "ad-hoc" (printComma(), getTensorFilename(), ...).

387

I didn't quite follow what you describe here?
Why adding 1 to the offset? The offset is physical in the buffer right now, are you thinking about making it a logical offset? When would the logical -> physical take place?

422

I can use a SmallVector :)

[mehdi_amini]

Use LLVM SmallVector instead of std::vector and refactor iterator indexing to share the code between Ranked/Unranked iterators

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

I'm very happy if you have the cycles to devote to this and tidy things up.

Basically the original version of CRunnerUtils was running on ARM micro but it was a half manual process with kludges to connect to a Makefile based system that I didn't have access to (i.e. shipping pre-cross-compiled .o across the wall).
As time passed and this was only used as the runtime library support to write some basic integration tests, other stuff got added.

It would be great if there was a simple way to test the works without a C++ runtime aspect, as this was the original

387

I meant you could make available and reuse these functions from mlir/include/mlir/Dialect/Vector/VectorUtils.h

/// Computes and returns the linearized index of 'offsets' w.r.t. 'basis'.
int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis);

/// Given the strides together with a linear index in the dimension
/// space, returns the vector-space offsets in each dimension for a
/// de-linearized index.
SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> strides,
                                    int64_t linearIndex);

Then your function would just delinearize(strides, linearize(offset, strides) + 1);.
If you were interested in performance (which I would claim you shouldn't be here, or you would just use some prebaked Torch-like library), you could additionally store the linearOffset in to avoid 1/2 the compute.

[mehdi_amini]

Use a templated Range argument for operator []

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

To be clear, by "C++ runtime" do you mean the non-header part of the STL?
Otherwise, without RTTI enabled there shouldn't be anything left I think

387

Ah I see what you mean now, thanks.

It isn't clear to me that the + 1 is enough here if the outermost stride isn't 1? I'm not sure if the behavior of delinearize should even be specific for an offset that is impossible to get from linearize...
But I guess I could replace +1 with + strides[rank-1]?

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

372

This could be: linearize(indices, strides) ?

387

You're right sorry, I took the wrong basis.

Here is the thinking, dropping the word "offset" as it is overloaded by now.

Indexing in the basis of memref "sizes" gives you the virtual index, you can go from linear to multi-D as:
delinearize(sizes, linearize(indices, sizes) + 1).

Indexing in the basis of memref "strides" gives you the physical displacement which skips holes, you can go from linear to multi-D as:
linearize(delinearize(sizes, linearize(indices, sizes) + 1), *strides*).

So assuming you keep a list of indices, to get the physical displacement of your next element compared to the base, you can:
linearize(delinearize(sizes, linearize(indices, sizes) + 1), *strides*)

If you prefer instead to keep the 1-D physical displacement, you can get the displacement of the next entry by:
linearize(delinearize(sizes, linearize(delinearize(displacement, *strides*), sizes) + 1), *strides*)

Does this make sense ?

[bondhugula]

mlir/unittests/ExecutionEngine/Invoke.cpp

421

Is this a typo here?

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

There is another wrinkle here: this file is built in C++11 mode? Why?
It is a bit confusing right now to have this piece of the project compiled differently and in a way that it can't interact with any other data-structure from ADT/Support, what is the purpose here?
If there is a specific needs, I'd like to see isolate in its own component (separate directory, etc.) rather than mixed with the ExecutionEngine folder which is part of the regular C++ infra.

[bondhugula]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

I think this file is built the same way as the rest of MLIR. The comment there perhaps exists in the interest of avoiding link time issues with libmlir_runner_utils.so - for eg. to link objects compiled from MLIR outside of the mlir-cpu-runner JIT?

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

We have this: mlir/lib/ExecutionEngine/CMakeLists.txt:set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11)

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

There is another wrinkle here: this file is built in C++11 mode? Why?
It is a bit confusing right now to have this piece of the project compiled differently and in a way that it can't interact with any other data-structure from ADT/Support, what is the purpose here?
If there is a specific needs, I'd like to see isolate in its own component (separate directory, etc.) rather than mixed with the ExecutionEngine folder which is part of the regular C++ infra.

Yes there are specific needs for this file: when it was created it was used to run ModelBuilder with some ARM micro (both arm32 and arm64 I believe) HW. I don't know offhand the exact processor but we can find out if relevant. Anyway, what's relevant is that the SDK, compiler etc for those required what is documented at the top of the file:

1. use C++11
2. don't require linking in libstdc++.

If you want to iterate on a better set of requirements and think there is a better place/way to do this, by all means let's start the discussion.
Maybe we should start an integration_test/xxx directory for some e2e test on a restricted embedded HW.

In the meantime, let's not break it.

35

I don't see value in introducing a dependency to LLVM at runtime here.
Things were absolutely fine with std::array.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

In the meantime, let's not break it.

I don't think this is reasonable to have this kind of requirement here right now.

Wanting to support an embedded environment which does not support the LLVM C++ standard is interesting, but I don't think we have an agreed goal to have this in-tree at the moment. I'm more inclined to align everything right now for what the project supports, rather that having a "project within the project" for out-of-tree reasons.

35

std::array is fine if you have a fixed rank at compile time, it does not scale with unranked...

[mehdi_amini]

Rebase and fix API changes + typo found by Uday

[bondhugula]

Looks good. Could you add an additional line to the commit summary for better context? The commit title is a bit cryptic and the connection to runtime utils and runners is missing.

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

144

Nit: * -> .

322–323

Doc comments here please.

mlir/unittests/ExecutionEngine/Invoke.cpp

316–320

Can these five lines be put in a helper registerDialectsAndParseSourceString? You have four repetitions.

[mehdi_amini]

I haven't worked around the specific requirement in the CRunnerUtils file and I suspect it requires some refactoring. But I don't know also the exact use-case and how to test it in order to be able to propose a path forward here, @nicolasvasilache can you help figure out who's using this and how?

[nicolasvasilache]

I haven't worked around the specific requirement in the CRunnerUtils file and I suspect it requires some refactoring. But I don't know also the exact use-case and how to test it in order to be able to propose a path forward here, @nicolasvasilache can you help figure out who's using this and how?

I don't either which is why the best proxy I found is mlir/lib/ExecutionEngine/CMakeLists.txt:set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11).
I don't think there are active users though so this could be your window of opportunity to move everything into RunnerUtils as I suggested above.

I still question the decision to not use the lowest common denominator when the logic you need is barely more than ptr + offset.
But I also see more value in progress and iteration than YAGNI discussions: it is perfectly valid to say that interested parties should grow/bring their own tensor library and this is merely for providing the batteries than MLIR is lacking.
If you want to grow this into something more than simple 2xAAA, an RFC would be useful.

Please address the comments on the indexing logic.

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

34

As usual, with MLIR everything is on a per-need basis, the need for setting up such goals has not come up before.
This probably deserves an RFC.

In the meantime, to unblock progress, I recommend moving all the code from this file into RunnerUtils.cpp and revisit later.
CRunnerUtils.* should not stay around if you're breaking the requirements.

35

pointer + offset ?

372

I don't believe this has been addressed?

387

I don't believe this has been addressed?

[mehdi_amini]

In D96397#2678641, @nicolasvasilache wrote:

I haven't worked around the specific requirement in the CRunnerUtils file and I suspect it requires some refactoring. But I don't know also the exact use-case and how to test it in order to be able to propose a path forward here, @nicolasvasilache can you help figure out who's using this and how?

I don't either which is why the best proxy I found is mlir/lib/ExecutionEngine/CMakeLists.txt:set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11).
I don't think there are active users though so this could be your window of opportunity to move everything into RunnerUtils as I suggested above.

I'm not sure how I could "move" things without knowing who depends on the things I'd move :)

That said my best shot at a plan for a future layering would be:

• A low-level layer to manipulate memref descriptor written purely in C. That would make it a C runtime very suitable for embedding on micro environment as well I think.
• A higher-level layer written in modern C++ with all the possible niceties one want. This layer could target the low-level C API to manipulate the descriptors, that said it is possible that template+inline would give better perf for C++ users.

WDYT?
@bondhugula as well?

[nicolasvasilache]

That said my best shot at a plan for a future layering would be:

This sounds great to me!

[bondhugula]

In D96397#2680121, @mehdi_amini wrote:

In D96397#2678641, @nicolasvasilache wrote:

I haven't worked around the specific requirement in the CRunnerUtils file and I suspect it requires some refactoring. But I don't know also the exact use-case and how to test it in order to be able to propose a path forward here, @nicolasvasilache can you help figure out who's using this and how?

I don't either which is why the best proxy I found is mlir/lib/ExecutionEngine/CMakeLists.txt:set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11).
I don't think there are active users though so this could be your window of opportunity to move everything into RunnerUtils as I suggested above.

I'm not sure how I could "move" things without knowing who depends on the things I'd move :)

That said my best shot at a plan for a future layering would be:

• A low-level layer to manipulate memref descriptor written purely in C. That would make it a C runtime very suitable for embedding on micro environment as well I think.
• A higher-level layer written in modern C++ with all the possible niceties one want. This layer could target the low-level C API to manipulate the descriptors, that said it is possible that template+inline would give better perf for C++ users.

WDYT?
@bondhugula as well?

This sounds great to me too.

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

 #define MLIR_CRUNNERUTILS_EXPORT __declspec(dllimport)
 #endif // mlir_c_runner_utils_EXPORTS
 #endif // MLIR_CRUNNERUTILS_EXPORT
 #else  // _WIN32
 #define MLIR_CRUNNERUTILS_EXPORT
 #define MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
 #endif // _WIN32
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
 #include <array>
 #include <cassert>
 #include <cstdint>
 #include <initializer_list>
 
 //===----------------------------------------------------------------------===//
 // Codegen-compatible structures for Vector type.
 //===----------------------------------------------------------------------===//
struct Vector1D<T, Dim, /*IsPowerOf2=*/false> {
   }
   inline T &operator[](unsigned i) { return vector[i]; }
   inline const T &operator[](unsigned i) const { return vector[i]; }
 
 private:
   T vector[Dim];
   char padding[nextPowerOf2(sizeof(T[Dim])) - sizeof(T[Dim])];
 };
+
+/// Update the `indices` array representing the subscript for a memref to point
+/// to the next element. The `sizes` is the current shape of the memref, and
+/// `strides` the physical size for each dimension. `offset` is the physical
+/// offset in the buffer for the current element. Returns the physical offset in
+/// the buffer for the next element, or -1 if the end of the memref has been
+/// reached.
+inline int64_t offsetForNextElement(llvm::MutableArrayRef<int64_t> indices,
+                                    llvm::ArrayRef<int64_t> sizes,
+                                    llvm::ArrayRef<int64_t> strides,
+                                    int64_t offset) {
+  int dim = indices.size() - 1;
+  // Start from the end of the indices and try to find the outer most dimension
+  // where the index can be incremented. For each dimension where the index is
+  // at the end, roll back the offset to the beginning for this dimension (the
+  // next outer dimension will increment by the stride size).
+  while (dim >= 0 && indices[dim] == (sizes[dim] - 1)) {
+    offset -= indices[dim] * strides[dim];
+    indices[dim] = 0;
+    --dim;
+  }
+  if (dim < 0) {
+    // All indices are at the end for every dimensions: end of the memref.
+    offset = -1;
+  } else {
+    // Increment the index for the current dimension, and move the offset
+    // forward by the current stride.
+    ++indices[dim];
+    offset += strides[dim];
+  }
+  return offset;
+}
+
 } // end namespace detail
 } // end namespace mlir
 
 // N-D vectors recurse down to 1-D.
 template <typename T, int Dim, int... Dims>
 struct Vector {
   inline Vector<T, Dims...> &operator[](unsigned i) { return vector[i]; }
   inline const Vector<T, Dims...> &operator[](unsigned i) const {

 /// Iterate over all elements in a strided memref.
 template <typename T, int Rank>
 class StridedMemrefIterator {
 public:
   StridedMemrefIterator(StridedMemRefType<T, Rank> &descriptor,
                         int64_t offset = 0)
       : offset(offset), descriptor(descriptor) {}
   StridedMemrefIterator<T, Rank> &operator++() {
-    int dim = Rank - 1;
-    while (dim >= 0 && indices[dim] == (descriptor.sizes[dim] - 1)) {
-      offset -= indices[dim] * descriptor.strides[dim];
-      indices[dim] = 0;
-      --dim;
-    }
-    if (dim < 0) {
-      offset = -1;
-      return *this;
-    }
-    ++indices[dim];
-    offset += descriptor.strides[dim];
+    offset = mlir::detail::offsetForNextElement(indices, descriptor.sizes,
+                                                descriptor.strides, offset);
     return *this;
   }
 
   T &operator*() { return descriptor.data[offset]; }
   T *operator->() { return &descriptor.data[offset]; }
 
   const std::array<int64_t, Rank> &getIndices() { return indices; }
 

 private:
   /// Pointer to the single element in the zero-ranked memref.
   T *elt;
 };
 
 //===----------------------------------------------------------------------===//
 // Codegen-compatible structure for UnrankedMemRef type.
 //===----------------------------------------------------------------------===//
+
+template <typename T>
+class UnrankedMemrefIterator;
+
 // Unranked MemRef
 template <typename T>
 struct UnrankedMemRefType {
   int64_t rank;
+  /// This is a pointer to an instantiation of StridedMemref<T, rank>.
   void *descriptor;
+
+  // Return a pointer to the `basePtr` member of the Strided descriptor.
+  T *basePtr() { return *reinterpret_cast<T **>(descriptor); }
+
+  // Return a pointer to the `data` member of the Strided descriptor.
+  T *data() {
+    return *reinterpret_cast<T **>(reinterpret_cast<int8_t *>(descriptor) +
+                                   sizeof(T *));
+  }
+
+  // Return the `offset` member of the Strided descriptor.
+  int64_t offset() {
+    return *reinterpret_cast<int64_t *>(reinterpret_cast<int8_t *>(descriptor) +
+                                        sizeof(T *) + sizeof(T *));
+  }
+
+  // Return the `sizes` member of the Strided descriptor.
+  llvm::ArrayRef<int64_t> sizes() {
+    return llvm::makeArrayRef(reinterpret_cast<int64_t *>(
+                                  reinterpret_cast<int8_t *>(descriptor) +
+                                  sizeof(T *) + sizeof(T *) + sizeof(int64_t)),
+                              rank);
+  }
+
+  // Return the `strides` member of the Strided descriptor.
+  llvm::ArrayRef<int64_t> strides() {
+    return llvm::makeArrayRef(
+        reinterpret_cast<int64_t *>(reinterpret_cast<int8_t *>(descriptor) +
+                                    sizeof(T *) + sizeof(T *) +
+                                    sizeof(int64_t) + rank * sizeof(int64_t)),
+        rank);
+  }
+
+  template <typename Range,
+            typename sfinae = decltype(std::declval<Range>().begin())>
+  T &operator[](Range indices) {
+    assert(static_cast<int64_t>(indices.size()) == rank);
+    int64_t offset = 0;
+    // This is a very ugly hard-coded way to access the data and strides
+    // inside a StridedMemrefType<T, rank>!!
+    llvm::ArrayRef<int64_t> strides = this->strides();
+
+    for (int dim = rank - 1; dim >= 0; --dim)
+      offset += *(indices.begin() + dim) * strides[dim];
+    return data()[offset];
+  }
+
+  UnrankedMemrefIterator<T> begin() { return {*this}; }
+  UnrankedMemrefIterator<T> end() { return {*this, -1}; }
+};
+
+/// Iterate over all elements in an unranked strided memref.
+template <typename T>
+class UnrankedMemrefIterator {
+public:
+  UnrankedMemrefIterator(UnrankedMemRefType<T> &descriptor, int64_t offset = 0)
+      : offset(offset), descriptor(descriptor) {
+    indices.resize(descriptor.rank);
+  }
+  UnrankedMemrefIterator &operator++() {
+    offset = mlir::detail::offsetForNextElement(indices, descriptor.sizes(),
+                                                descriptor.strides(), offset);
+    return *this;
+  }
+
+  T &operator*() { return descriptor.data()[offset]; }
+  T *operator->() { return descriptor.data()[offset]; }
+
+  bool operator==(const UnrankedMemrefIterator &other) const {
+    return &other.descriptor == &descriptor && other.offset == offset;
+  }
+
+  bool operator!=(const UnrankedMemrefIterator &other) const {
+    return !(*this == other);
+  }
+
+  llvm::ArrayRef<int64_t> getIndices() { return indices; }
+
+private:
+  /// Offset in the buffer. This can be derived from the indices and the
+  /// descriptor.
+  int64_t offset = 0;
+  /// Array of indices in the multi-dimensional memref.
+  llvm::SmallVector<int64_t, 4> indices = {};
+  /// Descriptor for the unranked memref.
+  UnrankedMemRefType<T> &descriptor;
 };
 
 //===----------------------------------------------------------------------===//
 // DynamicMemRefType type.
 //===----------------------------------------------------------------------===//
 // A reference to one of the StridedMemRef types.
 template <typename T>
 class DynamicMemRefType {
 public:
   explicit DynamicMemRefType(const StridedMemRefType<T, 0> &mem_ref)
       : rank(0), basePtr(mem_ref.basePtr), data(mem_ref.data),
         offset(mem_ref.offset), sizes(nullptr), strides(nullptr) {}
   template <int N>
   explicit DynamicMemRefType(const StridedMemRefType<T, N> &mem_ref)
       : rank(N), basePtr(mem_ref.basePtr), data(mem_ref.data),
         offset(mem_ref.offset), sizes(mem_ref.sizes), strides(mem_ref.strides) {

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

 #ifndef MLIR_EXECUTIONENGINE_MEMREFUTILS_H_
 #define MLIR_EXECUTIONENGINE_MEMREFUTILS_H_
 
 namespace mlir {
 using AllocFunType = llvm::function_ref<void *(size_t)>;
 
 namespace detail {
 
+/// Given a shape with sizes greater than 0 along all dimensions, populate the
+/// provided `strides` array with the distance, in number of elements, between a
+/// slice in a dimension and the next slice in the same dimension.
+///    e.g. shape[3, 4, 5] -> strides[20, 5, 1]
+inline void makeStrides(ArrayRef<int64_t> shape,
+                        MutableArrayRef<int64_t> strides) {
+  assert(shape.size() == strides.size() &&
+         "expect shapes and strides size to match");
+  int64_t running = 1;
+  for (int64_t idx = shape.size() - 1; idx >= 0; --idx) {
+    assert(shape[idx] && "size must be non-negative for all shape dimensions");
+    strides[idx] = running;
+    running *= shape[idx];
+  }
+}
+
 /// Given a shape with sizes greater than 0 along all dimensions, returns the
 /// distance, in number of elements, between a slice in a dimension and the next
 /// slice in the same dimension.
 ///    e.g. shape[3, 4, 5] -> strides[20, 5, 1]
 template <size_t N>
 inline std::array<int64_t, N> makeStrides(ArrayRef<int64_t> shape) {
   assert(shape.size() == N && "expect shape specification to match rank");
-  std::array<int64_t, N> res;
-  int64_t running = 1;
-  for (int64_t idx = N - 1; idx >= 0; --idx) {
-    assert(shape[idx] && "size must be non-negative for all shape dimensions");
-    res[idx] = running;
-    running *= shape[idx];
-  }
-  return res;
+  std::array<int64_t, N> strides;
+  makeStrides(shape, strides);
+  return strides;
 }
 
 /// Build a `StridedMemRefDescriptor<T, N>` that matches the MLIR ABI.
 /// This is an implementation detail that is kept in sync with MLIR codegen
 /// conventions.  Additionally takes a `shapeAlloc` array which
 /// is used instead of `shape` to allocate "more aligned" data and compute the
 /// corresponding strides.
 template <int N, typename T>
makeStridedMemRefDescriptor(T *ptr, T *alignedPtr, ArrayRef<int64_t> shape = {},
   assert(shapeAlloc.size() == N);
   StridedMemRefType<T, 0> descriptor;
   descriptor.basePtr = static_cast<T *>(ptr);
   descriptor.data = static_cast<T *>(alignedPtr);
   descriptor.offset = 0;
   return descriptor;
 }
 
+// Mallocs an UnrankedMemRefType<T>* that contains a ranked
+// StridedMemRefDescriptor<T, Rank>* and matches the MLIR ABI. This is an
+// implementation detail that is kept in sync with MLIR codegen conventions.
+template <typename T>
+::UnrankedMemRefType<T>
+makeUnrankedDescriptor(T *data, T *alignedData, ArrayRef<int64_t> shape,
+                       AllocFunType allocFun = &::malloc) {
+  ::UnrankedMemRefType<T> res{};
+  res.rank = shape.size();
+  if (res.rank == 0) {
+    res.descriptor = allocFun(sizeof(StridedMemRefType<T, 0>));
+    *static_cast<StridedMemRefType<T, 0> *>(res.descriptor) =
+        makeStridedMemRefDescriptor<0>(data, alignedData, shape, shape);
+  } else {
+    // Allocate and build a StridedMemRefType descriptor for Rank >= 1.
+    // The allocated size is the size of the descriptor for rank 1 plus two
+    // int64_t for each additional rank for one shape and one stride. The
+    // reinterpret_cast computations are computing the offset to the individual
+    // fields based on the rank.
+    res.descriptor = allocFun(sizeof(StridedMemRefType<T, 1>) +
+                              (res.rank - 1) * 2 * sizeof(int64_t));
+    *reinterpret_cast<T **>(res.descriptor) = data;
+    *reinterpret_cast<T **>(reinterpret_cast<int8_t *>(res.descriptor) +
+                            sizeof(T *)) = alignedData;
+    *reinterpret_cast<int64_t **>(reinterpret_cast<int8_t *>(res.descriptor) +
+                                  2 * sizeof(T *)) = 0;
+    std::copy(shape.begin(), shape.end(),
+              const_cast<int64_t *>(res.sizes().data()));
+    llvm::SmallVector<int64_t> strides;
+    strides.resize(res.rank);
+    makeStrides(shape, strides);
+    std::copy(strides.begin(), strides.end(),
+              const_cast<int64_t *>(res.strides().data()));
+  }
+  return res;
+}
+
+// Frees an UnrankedMemRefType<T>*
+template <typename T>
+void freeUnrankedDescriptor(::UnrankedMemRefType<T> *desc) {
+  free(desc->descriptor);
+  free(desc);
+}
+
 /// Align `nElements` of type T with an optional `alignment`.
 /// This replaces a portable `posix_memalign`.
 /// `alignment` must be a power of 2 and greater than the size of T. By default
 /// the alignment is sizeof(T).
 template <typename T>
 std::pair<T *, T *>
 allocAligned(size_t nElements, AllocFunType allocFun = &::malloc,
              llvm::Optional<uint64_t> alignment = llvm::Optional<uint64_t>()) {

 private:
   /// Custom deleter used to release the data buffer manager with the descriptor
   /// below.
   FreeFunType freeFunc;
   /// The descriptor is an instance of StridedMemRefType<T, rank>.
   DescriptorType descriptor;
 };
 
+/// Owning Unranked MemRef type that abstracts over the runtime type for memref.
+template <typename T>
+class OwningUnrankedMemRef {
+public:
+  using DescriptorType = ::UnrankedMemRefType<T>;
+  using FreeFunType = std::function<void(DescriptorType)>;
+
+  // Allocate a new UnrankedMemref with a given `shape` and initializer
+  // of type ElementWiseVisitor. Can optionally take specific `alloc` and `free`
+  // functions.
+  OwningUnrankedMemRef(
+      ArrayRef<int64_t> shape, ElementWiseVisitor<T> init = {},
+      llvm::Optional<uint64_t> alignment = llvm::Optional<uint64_t>(),
+      AllocFunType alloc = &::malloc,
+      FreeFunType freeFun =
+          [](DescriptorType descriptor) {
+            auto *strided_descriptor =
+                reinterpret_cast<StridedMemRefType<T, 0> *>(
+                    descriptor.descriptor);
+            ::free(strided_descriptor->data);
+          })
+      : freeFunc(freeFun) {
+    int64_t nElements = 1;
+    for (int64_t s : shape)
+      nElements *= s;
+    T *data, *alignedData;
+    std::tie(data, alignedData) =
+        detail::allocAligned<T>(nElements, alloc, alignment);
+    descriptor =
+        detail::makeUnrankedDescriptor(data, alignedData, shape, alloc);
+    if (init) {
+      for (UnrankedMemrefIterator<T> it = descriptor.begin(),
+                                     end = descriptor.end();
+           it != end; ++it)
+        init(*it, it.getIndices());
+    } else {
+      memset(descriptor.data(), 0,
+             nElements * sizeof(T) +
+                 alignment.getValueOr(detail::nextPowerOf2(sizeof(T))));
+    }
+  }
+
+  /// Take ownership of an existing UnrankMemRef descriptor with a custom
+  /// deleter;
+  OwningUnrankedMemRef(FreeFunType freeFunc, DescriptorType descriptor)
+      : freeFunc(freeFunc), descriptor(descriptor) {}
+  ~OwningUnrankedMemRef() { freeFunc(descriptor); }
+  T &operator[](std::initializer_list<int64_t> indices) {
+    return descriptor[std::move(indices)];
+  }
+
+  DescriptorType &operator*() { return descriptor; }
+  DescriptorType *operator->() { return &descriptor; }
+
+private:
+  FreeFunType freeFunc;
+  DescriptorType descriptor;
+};
+
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_MEMREFUTILS_H_

mlir/lib/ExecutionEngine/CMakeLists.txt

 
 add_mlir_library(mlir_c_runner_utils
   SHARED
   CRunnerUtils.cpp
   SparseUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
   )
-set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11)
 target_compile_definitions(mlir_c_runner_utils PRIVATE mlir_c_runner_utils_EXPORTS)
 
 add_mlir_library(mlir_runner_utils
   SHARED
   RunnerUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
 )

mlir/unittests/ExecutionEngine/Invoke.cpp

TEST(NativeMemRefJit, JITCallback) {
   int32_t coefficient = 3.;
   llvm::Error error = jit->invoke("caller_for_callback", &*A, coefficient);
   ASSERT_TRUE(!error);
   count = 1;
   for (float elt : *A)
     ASSERT_EQ(elt, coefficient * count++);
 }
 
+TEST(NativeMemRefJit, UnrankedMemref_Rank0) {
+  OwningUnrankedMemRef<float> A({});
+  A[{}] = 42.;
+  ASSERT_EQ(*A->data(), 42);
+  A[{}] = 0;
+  std::string moduleStr = R"mlir(
+  func @unranked_zero(%arg0 : memref<*xf32>) attributes { llvm.emit_c_interface } {
+    %cst42 = constant 42.0 : f32
+    %casted = memref_cast %arg0 : memref<*xf32> to memref<f32>
+    store %cst42, %casted[] : memref<f32>
+    return
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  auto module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  auto jit = std::move(jitOrError.get());
+
+  llvm::Error error = jit->invoke("unranked_zero", &*A);
+  ASSERT_TRUE(!error);
+  EXPECT_EQ((A[{}]), 42.);
+  for (float &elt : *A)
+    EXPECT_EQ(&elt, &(A[{}]));
+}
+
+TEST(NativeMemRefJit, UnrankedMemref_Rank1) {
+  int64_t shape[] = {9};
+  OwningUnrankedMemRef<float> A(shape);
+  int count = 1;
+  for (float &elt : *A) {
+    EXPECT_EQ(&elt, &(A[{count - 1}]));
+    elt = count++;
+  }
+
+  std::string moduleStr = R"mlir(
+  func @unranked_one(%arg0 : memref<*xf32>) attributes { llvm.emit_c_interface } {
+    %cst42 = constant 42.0 : f32
+    %cst5 = constant 5 : index
+    %casted = memref_cast %arg0 : memref<*xf32> to memref<?xf32>
+    store %cst42, %casted[%cst5] : memref<?xf32>
+    return
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  auto module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  auto jit = std::move(jitOrError.get());
+
+  llvm::Error error = jit->invoke("unranked_one", &*A);
+  ASSERT_TRUE(!error);
+  count = 1;
+  for (float elt : *A) {
+    if (count == 6)
+      EXPECT_EQ(elt, 42.);
+    else
+      EXPECT_EQ(elt, count);
+    count++;
+  }
+}
+
+TEST(NativeMemRefJit, UnrankedMemref_Rank2) {
+  constexpr int K = 3;
+  constexpr int M = 7;
+  // Prepare arguments beforehand.
+  auto init = [=](float &elt, ArrayRef<int64_t> indices) {
+    assert(indices.size() == 2);
+    elt = M * indices[0] + indices[1];
+  };
+  int64_t shape[] = {K, M};
+  OwningUnrankedMemRef<float> A(shape, init);
+
+  for (int i = 0; i < K; ++i)
+    for (int j = 0; j < M; ++j)
+      EXPECT_EQ((A[{i, j}]), i * M + j);
+
+  std::string moduleStr = R"mlir(
+  func @unranked_memref(%arg0 : memref<*xf32>, %arg1 : memref<*xf32>) attributes { llvm.emit_c_interface } {
+    %x = constant 2 : index
+    %y = constant 1 : index
+    %cst42 = constant 42.0 : f32
+    %arg0_cast = memref_cast %arg0 : memref<*xf32> to memref<?x?xf32>
+    %arg1_cast = memref_cast %arg1 : memref<*xf32> to memref<?x?xf32>
+    store %cst42, %arg0_cast[%y, %x] : memref<?x?xf32>
+    store %cst42, %arg1_cast[%x, %y] : memref<?x?xf32>
+    return
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  auto module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  auto jit = std::move(jitOrError.get());
+  llvm::Error error = jit->invoke("unranked_memref", &*A, &*A);
+  ASSERT_TRUE(!error);
+  EXPECT_EQ((A[{1, 2}]), 42.);
+  EXPECT_EQ((A[{2, 1}]), 42.);
+}
+
+// A helper function operating on unranked memref that will be called from the
+// JIT.
+static void unranked_memref_multiply(::UnrankedMemRefType<float> *memref,
+                                     int32_t coefficient) {
+  for (float &elt : *memref)
+    elt *= coefficient;
+}
+
+TEST(NativeMemRefJit, UnrankedMemrefCallback) {
+  constexpr int K = 3;
+  constexpr int M = 7;
+  // Prepare arguments beforehand.
+  int64_t shape[] = {K, M};
+  OwningUnrankedMemRef<float> A(shape);
+  int count = 0;
+  for (float &elt : *A)
+    elt = count++;
+
+  for (int i = 0; i < K; ++i)
+    for (int j = 0; j < M; ++j)
+      EXPECT_EQ((A[{i, j}]), i * M + j);
+
+  std::string moduleStr = R"mlir(
+  func private @callback(%arg0: memref<*xf32>, %coefficient: i32)  attributes { llvm.emit_c_interface } 
+  func @unrankedcaller_for_callback(%arg0: memref<*xf32>, %coefficient: i32) attributes { llvm.emit_c_interface } {
+    call @callback(%arg0, %coefficient) : (memref<*xf32>, i32) -> ()
+    return
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  auto module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  auto jit = std::move(jitOrError.get());
+  // Define any extra symbols so they're available at runtime.
+  jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) {
+    llvm::orc::SymbolMap symbolMap;
+    symbolMap[interner("_mlir_ciface_callback")] =
+        llvm::JITEvaluatedSymbol::fromPointer(unranked_memref_multiply);
+    return symbolMap;
+  });
+
+  int32_t coefficient = 7;
+  llvm::Error error =
+      jit->invoke("unrankedcaller_for_callback", &*A, coefficient);
+  ASSERT_TRUE(!error);
+  count = 0;
+  for (float elt : *A)
+    ASSERT_EQ(elt, coefficient * count++);
+}
+
 #endif // _WIN32