[mlir][CRunnerUtils] Fix iterators accessing MemRefs with non-zero offset

[ubfx]

Summary

MemRef descriptors contain - among others - a field called alignedPtr or data and a field called offset. The actual buffer of the MemRef starts at offset elements after alignedPtr. In the CRunnerUtils, there exist helper classes to iterate over MemRefs' elements but the offset is not handled consistently so that accessing a MemRef with an offset != 0 via an iterator will lead to incorrect results.

The problem is that "offset" can be understood in two ways, firstly as the offset of the beginning of the MemRef with respect to the alignedPtr, ie what the offset field means in the MemRef descriptor, and secondly as the offset of some element within the MemRef relative to the first element of the MemRef, which could more accurately be called something like linearIndex.

The offset field within StridedMemRefIterator and DynamicMemRefIterator are interpreted the first way, therefore the offsets passed to the constructors of these classes need to account for the already existing offset in the descriptor on top of any potential "shift" within the MemRef.
This patch takes care of that and adds some basic tests that catch problems with indexing MemRefs with an offset.

[ubfx]

add newline at end of file

[ubfx]

Tagging @mscuttari and @mehdi_amini who have last worked on this

[mscuttari]

Thanks for the patch. I am away from home though, I will try giving a look next week!

[mehdi_amini]

Thanks! What we floated with @ftynse recently was whether this offset was really needed or if we could just remove it entirely, we haven't articulated a good use case why it exists: thoughts?

[ubfx]

In D157008#4571671, @mehdi_amini wrote:

we haven't articulated a good use case why it exists: thoughts?

If I had to make a case for the offset, I would argue that it is the only way to allow the user to do custom arithmetic on the memref "pointer", before it actually is lowered to a pointer. If we want to support something like extract_strided_metadata -> do something -> reinterpret_cast, it seems like offset fits in with sizes and strides because it's also given in the denomination of the memref element size. I'm not sure whether such a use case should be encouraged but if we want to be able to move around the pointer via something like reinterpret_cast, then we need something like the offset to hide away the element size.

[ubfx]

friendly ping 🙏

[mehdi_amini]

The direction in https://github.com/llvm/llvm-project/issues/64334 is to remove the offset.
That said we can fix this in the meantime.

[ubfx]

Thank you, Mehdi. While I can't really weigh in on the very in-depth discussion about the consequences of restructuring the Memref descriptors, I have found a few bugs around the MemrefToLLVM / Memref interface to C / Python and submitted patches for them. I noticed that these bugs are in pretty core functionality (i.e. LLVM lowering of memref.copy, memref.transpose, etc.) but they had not been caught earlier because we don't seem to have tests that check (dynamically) whether the generated LLVMIR actually does what the Memref Ops demand.

We have unit tests for the ExecutionEngine that run generated LLVMIR but they are checking core functionality of the ExecutionEngine / Runtime and not the correctness of the lowering of higher level Ops (in this case memref). Do you think it would make sense to add such unittests, which lower to LLVMIR, run the IR using the ExecutionEngine and then check the behavior? I could provide a few of these and maybe they would be useful to ensure the transition to removing the offset field goes smoothly. I just wasn't sure whether such unittests are even desired because the unittests/Conversion folder is pretty empty.

btw: could you commit this revision for me?

bugs I'm talking about:
https://reviews.llvm.org/D156126
https://reviews.llvm.org/D158494
https://reviews.llvm.org/D159290

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

T &operator[](Range &&indices) {
     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> begin() { return {*this, offset}; }
   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];
struct StridedMemRefType<T, 1> {
   template <typename Range,
             typename sfinae = decltype(std::declval<Range>().begin())>
   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> begin() { return {*this, offset}; }
   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,
             typename sfinae = decltype(std::declval<Range>().begin())>
   T &operator[](Range indices) {
     assert((indices.size() == 0) &&
            "Expect empty indices for 0-rank memref subscript");
     return data[offset];
   }
 
-  StridedMemrefIterator<T, 0> begin() { return {*this}; }
-  StridedMemrefIterator<T, 0> end() { return {*this, 1}; }
+  StridedMemrefIterator<T, 0> begin() { return {*this, offset}; }
+  StridedMemrefIterator<T, 0> end() { return {*this, offset + 1}; }
 };
 
 /// Iterate over all elements in a strided memref.
 template <typename T, int Rank>
 class StridedMemrefIterator {
 public:
   using iterator_category = std::forward_iterator_tag;
   using value_type = T;
T &operator[](Range &&indices) {
     for (int dim = rank - 1; dim >= 0; --dim) {
       int64_t currentIndex = *(indices.begin() + dim);
       assert(currentIndex < sizes[dim] && "Index overflow");
       curOffset += currentIndex * strides[dim];
     }
     return data[curOffset];
   }
 
-  DynamicMemRefIterator<T> begin() { return {*this}; }
+  DynamicMemRefIterator<T> begin() { return {*this, offset}; }
   DynamicMemRefIterator<T> end() { return {*this, -1}; }
 
   // This operator[] is extremely slow and only for sugaring purposes.
   DynamicMemRefType<T> operator[](int64_t idx) {
     assert(rank > 0 && "can't make a subscript of a zero ranked array");
 
     DynamicMemRefType<T> res(*this);
     --res.rank;

mlir/unittests/ExecutionEngine/CMakeLists.txt

 add_mlir_unittest(MLIRExecutionEngineTests
   DynamicMemRef.cpp
+  StridedMemRef.cpp
   Invoke.cpp
 )
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 target_link_libraries(MLIRExecutionEngineTests
   PRIVATE
   MLIRArithToLLVM
   MLIRExecutionEngine
   MLIRMemRefToLLVM
   MLIRReconcileUnrealizedCasts
   ${dialect_libs}
 
 )

mlir/unittests/ExecutionEngine/DynamicMemRef.cpp

TEST(DynamicMemRef, rankThree) {
   memRef.strides[0] = 12;
   memRef.strides[1] = 4;
   memRef.strides[2] = 1;
 
   DynamicMemRefType<int> dynamicMemRef(memRef);
 
   llvm::SmallVector<int, 24> values(dynamicMemRef.begin(), dynamicMemRef.end());
   EXPECT_THAT(values, ElementsAreArray(data));
-}
+}
+
+TEST(DynamicMemRef, rankOneWithOffset) {
+  constexpr int offset = 4;
+  std::array<int, 3 + offset> buffer;
+
+  for (size_t i = 0; i < buffer.size(); ++i) {
+    buffer[i] = i;
+  }
+
+  StridedMemRefType<int, 1> memRef;
+  memRef.basePtr = buffer.data();
+  memRef.data = buffer.data();
+  memRef.offset = offset;
+  memRef.sizes[0] = 3;
+  memRef.strides[0] = 1;
+
+  DynamicMemRefType<int> dynamicMemRef(memRef);
+
+  llvm::SmallVector<int, 3> values(dynamicMemRef.begin(), dynamicMemRef.end());
+
+  for (int64_t i = 0; i < 3; ++i) {
+    EXPECT_EQ(values[i], buffer[offset + i]);
+    EXPECT_EQ(*dynamicMemRef[i], buffer[offset + i]);
+  }
+}

mlir/unittests/ExecutionEngine/StridedMemRef.cpp

+//===- StridedMemRef.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/ExecutionEngine/CRunnerUtils.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include "gmock/gmock.h"
+
+using namespace ::mlir;
+using namespace ::testing;
+
+TEST(StridedMemRef, rankOneWithOffset) {
+  std::array<int, 15> data;
+
+  for (size_t i = 0; i < data.size(); ++i) {
+    data[i] = i;
+  }
+
+  StridedMemRefType<int, 1> memRefA;
+  memRefA.basePtr = data.data();
+  memRefA.data = data.data();
+  memRefA.offset = 0;
+  memRefA.sizes[0] = 10;
+  memRefA.strides[0] = 1;
+
+  StridedMemRefType<int, 1> memRefB = memRefA;
+  memRefB.offset = 5;
+
+  llvm::SmallVector<int, 10> valuesA(memRefA.begin(), memRefA.end());
+  llvm::SmallVector<int, 10> valuesB(memRefB.begin(), memRefB.end());
+
+  for (int64_t i = 0; i < 10; ++i) {
+    EXPECT_EQ(valuesA[i], i);
+    EXPECT_EQ(valuesA[i] + 5, valuesB[i]);
+  }
+}