Add convenience C++ helper to manipulate ranked strided memref

[bondhugula]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

81

Missing doc comment (for the field above as well btw). All fields should be documented. On another note, the doc comment on Vector1D should have been triple ///.

214

If Rank is 0, offset will become -1. Is this intended? A comment is missing at least here and above where you have int dim = Rank - 1. Otherwise, you need a static_assert on Rank not being zero. The offset field is also missing a doc comment.

236–238

Missing doc comments however obvious it may appear now.

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

10

Nit: Missing period at the end.

41–44

Triple /// comments please.

47

Assert message please.

51

negatice -> negative

58

Triple /// here and everywhere below.

143

-> for ranked strided memrefs.

207

Missing doc comment.

214

Missing new line at end.

mlir/unittests/ExecutionEngine/Invoke.cpp

110

auto -> float

117–122

Test rank = 0 case as well? It may exercise corner cases.

[bondhugula]

Can you please add a commit summary as well - "manipulate ranked strided memref" doesn't capture all of this support accurately.

[ftynse]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

136

Would it make sense to accept any range rather than just an initializer list?

137

Nit: please add an assertion message

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

80

Nit: "mallocs" sounds like there is heap allocation happening, but it is not the case here.

103

The other functions return the descriptor by-value rather than by-pointer. I suppose this template is not instantiated, otherwise it should have triggered an error.

110

s/allow/allows one/

113

I would suggest casting the result back to pair<T *, T *> before returning. This will also let makeStridedMemRefDescriptor functions take T * rather than void * for better type checking. Furthermore, if the template arguments in makeStridedMemRefDescriptor are swapped to template <int N, typename T>, I expect type deduction to kick in for element types.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

81

Yes, but note that I'm just doing a mechanical namespace change here. This whole file is a bit messy in my opinion, I'd like to revamp the whole directory in the near future actually.

136

As a template?

template<typename Range>
T &operator[](Range indices) {
 ...

[mehdi_amini]

Address comment and add more tests

[mehdi_amini]

Thanks for the review!

I added a bunch of more test to increase the coverage: rank-0, rank-1, and a test with a callback to a C++ native function.

PTAL!

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

214

Added a specialization for rank 0, and tests for rank0 and rank1

[bondhugula]

A few more minor comments. Please also address the commit summary comment above. This functionality is looking great to me. Please do go with Alex's deeper review comments.

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

81

I see - missed that. Would be great to have a subsequent NFC patch.

195

indices.empty() please.

196

no -> number

mlir/unittests/ExecutionEngine/Invoke.cpp

248

Micro-nit: drop reference to elt - isn't it more overhead FWIW?

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

1

I am assuming you are also refactoring usage in iree/experimental/ModelBuilder, in a separate revision?

[ftynse]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

136

Cool, thanks!

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

1

What do you mean exactly? (I haven't touched anything IREE related so far)

mlir/unittests/ExecutionEngine/Invoke.cpp

248

Will do. It was a copy paste from the loop above where I assigned to elt

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

196

used "empty"

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

195

It does not work with initializer_list actually :(

[bondhugula]

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

195

Oh, didn't realize that. Never mind.

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

1

I mean that this file (https://github.com/google/iree/blob/main/experimental/ModelBuilder/MemRefUtils.h) should disappear and its uses updated now that you have moved and improved the functionality in core.

Please make sure we have only 1 source of truth for all this (yours).

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

1

I'm a bit confused about how you're bringing downstream project management into MLIR core review here?
I may be missing something but it seems to me that it should be entirely independent: I see it as "up to IREE" to decide if they want to adopt new upstream capabilities as they see fit.
(and separately, it's a question for non-IREE google internal projects to switch to MLIR provided tools instead of IREE/experimental/...)

[nicolasvasilache]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

1

When code gets moved around I feel it's only natural to gently ask for an update of the former source of truth.
Yes, ModelBuilder wants to use this improved version, can you please also update IREE separately to avoid what is now copy-pasta of core? Thanks much!

[RKSimon]

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

195

@mehdi_amini incorrect argument order for memset

memset(&other.descriptor, 0, sizeof(other.descriptor));

Reported here https://pvs-studio.com/en/blog/posts/cpp/1003/ (N10)

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 <array>
+#include <cassert>
 #include <cstdint>
+#include <initializer_list>
 
 //===----------------------------------------------------------------------===//
 // Codegen-compatible structures for Vector type.
 //===----------------------------------------------------------------------===//
+namespace mlir {
 namespace detail {
 
 constexpr bool isPowerOf2(int N) { return (!(N & (N - 1))); }
 
 constexpr unsigned nextPowerOf2(int N) {
   return (N <= 1) ? 1 : (isPowerOf2(N) ? N : (2 * nextPowerOf2((N + 1) / 2)));
 }
 
private:
   T vector[Dim];
 };
 
 // 1-D vector, padded to the next power of 2 allocation.
 // Specialization occurs to avoid zero size arrays (which fail in -Werror).
 template <typename T, int Dim>
 struct Vector1D<T, Dim, /*IsPowerOf2=*/false> {
   Vector1D() {
-    static_assert(detail::nextPowerOf2(sizeof(T[Dim])) > sizeof(T[Dim]),
-                  "size error");
-    static_assert(detail::nextPowerOf2(sizeof(T[Dim])) < 2 * sizeof(T[Dim]),
+    static_assert(nextPowerOf2(sizeof(T[Dim])) > sizeof(T[Dim]), "size error");
+    static_assert(nextPowerOf2(sizeof(T[Dim])) < 2 * sizeof(T[Dim]),
                   "size error");
   }
   inline T &operator[](unsigned i) { return vector[i]; }
   inline const T &operator[](unsigned i) const { return vector[i]; }
 
 private:
   T vector[Dim];
-  char padding[detail::nextPowerOf2(sizeof(T[Dim])) - sizeof(T[Dim])];
+  char padding[nextPowerOf2(sizeof(T[Dim])) - sizeof(T[Dim])];
 };
 } // 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 {
     return vector[i];
   }
 
 private:
   Vector<T, Dims...> vector[Dim];
 };
 
 // 1-D vectors in LLVM are automatically padded to the next power of 2.
 // We insert explicit padding in to account for this.
 template <typename T, int Dim>
 struct Vector<T, Dim>
-    : public detail::Vector1D<T, Dim, detail::isPowerOf2(sizeof(T[Dim]))> {};
+    : public mlir::detail::Vector1D<T, Dim,
+                                    mlir::detail::isPowerOf2(sizeof(T[Dim]))> {
+};
 
 template <int D1, typename T>
 using Vector1D = Vector<T, D1>;
 template <int D1, int D2, typename T>
 using Vector2D = Vector<T, D1, D2>;
 template <int D1, int D2, int D3, typename T>
 using Vector3D = Vector<T, D1, D2, D3>;
 template <int D1, int D2, int D3, int D4, typename T>
 using Vector4D = Vector<T, D1, D2, D3, D4>;
 
 template <int N>
 void dropFront(int64_t arr[N], int64_t *res) {
   for (unsigned i = 1; i < N; ++i)
     *(res + i - 1) = arr[i];
 }
 
 //===----------------------------------------------------------------------===//
 // Codegen-compatible structures for StridedMemRef type.
 //===----------------------------------------------------------------------===//
+template <typename T, int Rank>
+class StridedMemrefIterator;
+
 /// StridedMemRef descriptor type with static rank.
 template <typename T, int N>
 struct StridedMemRefType {
   T *basePtr;
   T *data;
   int64_t offset;
   int64_t sizes[N];
   int64_t strides[N];
+
+  template <typename Range>
+  T &operator[](Range indices) {
+    assert(indices.size() == N &&
+           "indices should match rank in memref subscript");
+    int64_t curOffset = offset;
+    for (int dim = N - 1; dim >= 0; --dim) {
+      int64_t currentIndex = *(indices.begin() + dim);
+      assert(currentIndex < sizes[dim] && "Index overflow");
+      curOffset += currentIndex * strides[dim];
+    }
+    return data[curOffset];
+  }
+
+  StridedMemrefIterator<T, N> begin() { return {*this}; }
+  StridedMemrefIterator<T, N> end() { return {*this, -1}; }
+
   // This operator[] is extremely slow and only for sugaring purposes.
   StridedMemRefType<T, N - 1> operator[](int64_t idx) {
     StridedMemRefType<T, N - 1> res;
     res.basePtr = basePtr;
     res.data = data;
     res.offset = offset + idx * strides[0];
     dropFront<N>(sizes, res.sizes);
     dropFront<N>(strides, res.strides);
     return res;
   }
 };
 
 /// StridedMemRef descriptor type specialized for rank 1.
 template <typename T>
 struct StridedMemRefType<T, 1> {
   T *basePtr;
   T *data;
   int64_t offset;
   int64_t sizes[1];
   int64_t strides[1];
+
+  template <typename Range>
+  T &operator[](Range indices) {
+    assert(indices.size() == 1 &&
+           "indices should match rank in memref subscript");
+    return (*this)[*indices.begin()];
+  }
+
+  StridedMemrefIterator<T, 1> begin() { return {*this}; }
+  StridedMemrefIterator<T, 1> end() { return {*this, -1}; }
+
   T &operator[](int64_t idx) { return *(data + offset + idx * strides[0]); }
 };
 
 /// StridedMemRef descriptor type specialized for rank 0.
 template <typename T>
 struct StridedMemRefType<T, 0> {
   T *basePtr;
   T *data;
   int64_t offset;
+
+  template <typename Range>
+  T &operator[](Range indices) {
+    assert(indices.size() == 0 &&
+           "Expect no indices for 0-rank memref subscript");
+    return data[offset];
+  }
+
+  StridedMemrefIterator<T, 0> begin() { return {*this}; }
+  StridedMemrefIterator<T, 0> end() { return {*this, 1}; }
+};
+
+/// 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];
+    return *this;
+  }
+
+  T &operator*() { return descriptor.data[offset]; }
+  T *operator->() { return &descriptor.data[offset]; }
+
+  const std::array<int64_t, Rank> &getIndices() { return indices; }
+
+  bool operator==(const StridedMemrefIterator &other) const {
+    return other.offset == offset && &other.descriptor == &descriptor;
+  }
+
+  bool operator!=(const StridedMemrefIterator &other) const {
+    return !(*this == other);
+  }
+
+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.
+  std::array<int64_t, Rank> indices = {};
+  /// Descriptor for the strided memref.
+  StridedMemRefType<T, Rank> &descriptor;
+};
+
+/// Iterate over all elements in a 0-ranked strided memref.
+template <typename T>
+class StridedMemrefIterator<T, 0> {
+public:
+  StridedMemrefIterator(StridedMemRefType<T, 0> &descriptor, int64_t offset = 0)
+      : elt(descriptor.data + offset) {}
+
+  StridedMemrefIterator<T, 0> &operator++() {
+    ++elt;
+    return *this;
+  }
+
+  T &operator*() { return *elt; }
+  T *operator->() { return elt; }
+
+  // There are no indices for a 0-ranked memref, but this API is provided for
+  // consistency with the general case.
+  const std::array<int64_t, 0> &getIndices() {
+    // Since this is a 0-array of indices we can keep a single global const
+    // copy.
+    static const std::array<int64_t, 0> indices = {};
+    return indices;
+  }
+
+  bool operator==(const StridedMemrefIterator &other) const {
+    return other.elt == elt;
+  }
+
+  bool operator!=(const StridedMemrefIterator &other) const {
+    return !(*this == other);
+  }
+
+private:
+  /// Pointer to the single element in the zero-ranked memref.
+  T *elt;
 };
 
 //===----------------------------------------------------------------------===//
 // Codegen-compatible structure for UnrankedMemRef type.
 //===----------------------------------------------------------------------===//
 // Unranked MemRef
 template <typename T>
 struct UnrankedMemRefType {

mlir/include/mlir/ExecutionEngine/MemRefUtils.h

+//===- MemRefUtils.h - Memref helpers to invoke MLIR JIT code ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utils for MLIR ABI interfacing with frameworks.
+//
+// The templated free functions below make it possible to allocate dense
+// contiguous buffers with shapes that interoperate properly with the MLIR
+// codegen ABI.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/ExecutionEngine/CRunnerUtils.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include "llvm/Support/raw_ostream.h"
+
+#include <algorithm>
+#include <array>
+#include <cassert>
+#include <functional>
+#include <initializer_list>
+#include <memory>
+
+#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, 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;
+}
+
+/// 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>
+typename std::enable_if<(N >= 1), StridedMemRefType<T, N>>::type
+makeStridedMemRefDescriptor(T *ptr, T *alignedPtr, ArrayRef<int64_t> shape,
+                            ArrayRef<int64_t> shapeAlloc) {
+  assert(shape.size() == N);
+  assert(shapeAlloc.size() == N);
+  StridedMemRefType<T, N> descriptor;
+  descriptor.basePtr = static_cast<T *>(ptr);
+  descriptor.data = static_cast<T *>(alignedPtr);
+  descriptor.offset = 0;
+  std::copy(shape.begin(), shape.end(), descriptor.sizes);
+  auto strides = makeStrides<N>(shapeAlloc);
+  std::copy(strides.begin(), strides.end(), descriptor.strides);
+  return descriptor;
+}
+
+/// Build a `StridedMemRefDescriptor<T, 0>` 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>
+typename std::enable_if<(N == 0), StridedMemRefType<T, 0>>::type
+makeStridedMemRefDescriptor(T *ptr, T *alignedPtr, ArrayRef<int64_t> shape = {},
+                            ArrayRef<int64_t> shapeAlloc = {}) {
+  assert(shape.size() == N);
+  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;
+}
+
+/// 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>()) {
+  assert(sizeof(T) < (1ul << 32) && "Elemental type overflows");
+  auto size = nElements * sizeof(T);
+  auto desiredAlignment = alignment.getValueOr(nextPowerOf2(sizeof(T)));
+  assert((desiredAlignment & (desiredAlignment - 1)) == 0);
+  assert(desiredAlignment >= sizeof(T));
+  T *data = reinterpret_cast<T *>(allocFun(size + desiredAlignment));
+  uintptr_t addr = reinterpret_cast<uintptr_t>(data);
+  uintptr_t rem = addr % desiredAlignment;
+  T *alignedData = (rem == 0)
+                       ? data
+                       : reinterpret_cast<T *>(addr + (desiredAlignment - rem));
+  assert(reinterpret_cast<uintptr_t>(alignedData) % desiredAlignment == 0);
+  return std::make_pair(data, alignedData);
+}
+
+} // namespace detail
+
+//===----------------------------------------------------------------------===//
+// Public API
+//===----------------------------------------------------------------------===//
+
+/// Convenient callback to "visit" a memref element by element.
+/// This takes a reference to an individual element as well as the coordinates.
+/// It can be used in conjuction with a StridedMemrefIterator.
+template <typename T>
+using ElementWiseVisitor = llvm::function_ref<void(T &ptr, ArrayRef<int64_t>)>;
+
+/// Owning MemRef type that abstracts over the runtime type for ranked strided
+/// memref.
+template <typename T, int Rank>
+class OwningMemRef {
+public:
+  using DescriptorType = StridedMemRefType<T, Rank>;
+  using FreeFunType = std::function<void(DescriptorType)>;
+
+  /// Allocate a new dense StridedMemrefRef with a given `shape`. An optional
+  /// `shapeAlloc` array can be supplied to "pad" every dimension individually.
+  /// If an ElementWiseVisitor is provided, it will be used to initialize the
+  /// data, else the memory will be zero-initialized. The alloc and free method
+  /// used to manage the data allocation can be optionally provided, and default
+  /// to malloc/free.
+  OwningMemRef(
+      ArrayRef<int64_t> shape, ArrayRef<int64_t> shapeAlloc = {},
+      ElementWiseVisitor<T> init = {},
+      llvm::Optional<uint64_t> alignment = llvm::Optional<uint64_t>(),
+      AllocFunType allocFun = &::malloc,
+      std::function<void(StridedMemRefType<T, Rank>)> freeFun =
+          [](StridedMemRefType<T, Rank> descriptor) {
+            ::free(descriptor.data);
+          })
+      : freeFunc(freeFun) {
+    if (shapeAlloc.empty())
+      shapeAlloc = shape;
+    assert(shape.size() == Rank);
+    assert(shapeAlloc.size() == Rank);
+    for (unsigned i = 0; i < Rank; ++i)
+      assert(shape[i] <= shapeAlloc[i] &&
+             "shapeAlloc must be greater than or equal to shape");
+    int64_t nElements = 1;
+    for (int64_t s : shapeAlloc)
+      nElements *= s;
+    T *data, *alignedData;
+    std::tie(data, alignedData) =
+        detail::allocAligned<T>(nElements, allocFun, alignment);
+    descriptor = detail::makeStridedMemRefDescriptor<Rank>(data, alignedData,
+                                                           shape, shapeAlloc);
+    if (init) {
+      for (StridedMemrefIterator<T, Rank> 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 descriptor with a custom deleter.
+  OwningMemRef(DescriptorType descriptor, FreeFunType freeFunc)
+      : freeFunc(freeFunc), descriptor(descriptor) {}
+  ~OwningMemRef() {
+    if (freeFunc)
+      freeFunc(descriptor);
+  }
+  OwningMemRef(const OwningMemRef &) = delete;
+  OwningMemRef &operator=(const OwningMemRef &) = delete;
+  OwningMemRef &operator=(const OwningMemRef &&other) {
+    freeFunc = other.freeFunc;
+    descriptor = other.descriptor;
+    other.freeFunc = nullptr;
+    memset(0, &other.descriptor, sizeof(other.descriptor));
+  }
+  OwningMemRef(OwningMemRef &&other) { *this = std::move(other); }
+
+  DescriptorType &operator*() { return descriptor; }
+  DescriptorType *operator->() { return &descriptor; }
+  T &operator[](std::initializer_list<int64_t> indices) {
+    return descriptor[std::move(indices)];
+  }
+
+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;
+};
+
+} // namespace mlir
+
+#endif // MLIR_EXECUTIONENGINE_MEMREFUTILS_H_

mlir/unittests/ExecutionEngine/Invoke.cpp

 //===- Invoke.cpp ------------------------------------*- C++ -*-===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
+#include "mlir/ExecutionEngine/MemRefUtils.h"
 #include "mlir/ExecutionEngine/RunnerUtils.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/Parser.h"
 #include "mlir/Pass/PassManager.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
 
TEST(MLIRExecutionEngine, SubtractFloat) {
   // The result of the function must be passed as output argument.
   float result = -1;
   llvm::Error error =
       jit->invoke("foo", 43.0f, 1.0f, ExecutionEngine::result(result));
   ASSERT_TRUE(!error);
   ASSERT_EQ(result, 42.f);
 }
 
+TEST(NativeMemRefJit, ZeroRankMemref) {
+  OwningMemRef<float, 0> A({});
+  A[{}] = 42.;
+  ASSERT_EQ(*A->data, 42);
+  A[{}] = 0;
+  std::string moduleStr = R"mlir(
+  func @zero_ranked(%arg0 : memref<f32>) attributes { llvm.emit_c_interface } {
+    %cst42 = constant 42.0 : f32
+    store %cst42, %arg0[] : memref<f32>
+    return
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  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("zero_ranked", &*A);
+  ASSERT_TRUE(!error);
+  EXPECT_EQ((A[{}]), 42.);
+  for (float &elt : *A)
+    EXPECT_EQ(&elt, &(A[{}]));
+}
+
+TEST(NativeMemRefJit, RankOneMemref) {
+  int64_t shape[] = {9};
+  OwningMemRef<float, 1> A(shape);
+  int count = 1;
+  for (float &elt : *A) {
+    EXPECT_EQ(&elt, &(A[{count - 1}]));
+    elt = count++;
+  }
+
+  std::string moduleStr = R"mlir(
+  func @one_ranked(%arg0 : memref<?xf32>) attributes { llvm.emit_c_interface } {
+    %cst42 = constant 42.0 : f32
+    %cst5 = constant 5 : index
+    store %cst42, %arg0[%cst5] : memref<?xf32>
+    return
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  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("one_ranked", &*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, BasicMemref) {
+  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};
+  int64_t shapeAlloc[] = {K + 1, M + 1};
+  OwningMemRef<float, 2> A(shape, shapeAlloc, init);
+  ASSERT_EQ(A->sizes[0], K);
+  ASSERT_EQ(A->sizes[1], M);
+  ASSERT_EQ(A->strides[0], M + 1);
+  ASSERT_EQ(A->strides[1], 1);
+  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 @rank2_memref(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>) attributes { llvm.emit_c_interface } {
+    %x = constant 2 : index
+    %y = constant 1 : index
+    %cst42 = constant 42.0 : f32
+    store %cst42, %arg0[%y, %x] : memref<?x?xf32>
+    store %cst42, %arg1[%x, %y] : memref<?x?xf32>
+    return
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  OwningModuleRef module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
+
+  llvm::Error error = jit->invoke("rank2_memref", &*A, &*A);
+  ASSERT_TRUE(!error);
+  EXPECT_EQ((A[{1, 2}]), 42.);
+  EXPECT_EQ((A[{2, 1}]), 42.);
+}
+
+// A helper function that will be called from the JIT
+static void memref_multiply(::StridedMemRefType<float, 2> *memref,
+                            int32_t coefficient) {
+  for (float &elt : *memref)
+    elt *= coefficient;
+}
+
+TEST(NativeMemRefJit, JITCallback) {
+  constexpr int K = 2;
+  constexpr int M = 2;
+  int64_t shape[] = {K, M};
+  int64_t shapeAlloc[] = {K + 1, M + 1};
+  OwningMemRef<float, 2> A(shape, shapeAlloc);
+  int count = 1;
+  for (float &elt : *A)
+    elt = count++;
+
+  std::string moduleStr = R"mlir(
+  func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32)  attributes { llvm.emit_c_interface } 
+  func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32) attributes { llvm.emit_c_interface } {
+    %unranked = memref_cast %arg0: memref<?x?xf32> to memref<*xf32>
+    call @callback(%arg0, %coefficient) : (memref<?x?xf32>, i32) -> ()
+    return
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  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(memref_multiply);
+    return symbolMap;
+  });
+
+  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++);
+}
+
 #endif // _WIN32