[Mlir-commits] [mlir] ee2a225 - [mlir] Fix correct memset range in `OwningMemRef` zero-init (#158200)

[llvmlistbot at llvm.org]

Author: Ryan Kim
Date: 2025-09-13T12:04:55Z
New Revision: ee2a225a25fbc41fc7a47e089f09022f90eeaac3

URL: https://github.com/llvm/llvm-project/commit/ee2a225a25fbc41fc7a47e089f09022f90eeaac3
DIFF: https://github.com/llvm/llvm-project/commit/ee2a225a25fbc41fc7a47e089f09022f90eeaac3.diff

LOG: [mlir] Fix correct memset range in `OwningMemRef` zero-init (#158200)

`OwningMemref` allocates with overprovision + manual alignment.
This is fixing the zero-initialization of the data, the existing code
was potentially overrunning the allocation:

```cpp
memset(descriptor.data, 0, size + desiredAlignment); // ❌ may overrun
```

This is invalid because `descriptor.data` (the aligned pointer) **does
not point to the full allocated block** (`size + desiredAlignment`).
Zeroing that much from the aligned start can write past the end of the
allocation.

Instead we only initialize the data from the aligned pointer for the expected
buffer size. The padding from [allocatedPtr, alignedDataPtr] is left untouched.

Added: 
    

Modified: 
    mlir/include/mlir/ExecutionEngine/MemRefUtils.h
    mlir/unittests/ExecutionEngine/Invoke.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/ExecutionEngine/MemRefUtils.h b/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
index d66d757cb7a8e..e9471731afe13 100644
--- a/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
+++ b/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
@@ -164,19 +164,17 @@ class OwningMemRef {
     int64_t nElements = 1;
     for (int64_t s : shapeAlloc)
       nElements *= s;
-    auto [data, alignedData] =
+    auto [allocatedPtr, alignedData] =
         detail::allocAligned<T>(nElements, allocFun, alignment);
-    descriptor = detail::makeStridedMemRefDescriptor<Rank>(data, alignedData,
-                                                           shape, shapeAlloc);
+    descriptor = detail::makeStridedMemRefDescriptor<Rank>(
+        allocatedPtr, 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.value_or(detail::nextPowerOf2(sizeof(T))));
+      memset(alignedData, 0, nElements * sizeof(T));
     }
   }
   /// Take ownership of an existing descriptor with a custom deleter.

diff  --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index cdeeca20610f0..3161c7053f7a4 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -251,6 +251,24 @@ TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(BasicMemref)) {
   EXPECT_EQ((a[{2, 1}]), 42.);
 }
 
+TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(OwningMemrefZeroInit)) {
+  constexpr int k = 3;
+  constexpr int m = 7;
+  int64_t shape[] = {k, m};
+  // Use a large alignment to stress the case where the memref data/basePtr are
+  // disjoint.
+  int alignment = 8192;
+  OwningMemRef<float, 2> a(shape, {}, {}, alignment);
+  ASSERT_EQ(
+      (void *)(((uintptr_t)a->basePtr + alignment - 1) & ~(alignment - 1)),
+      a->data);
+  for (int i = 0; i < k; ++i) {
+    for (int j = 0; j < m; ++j) {
+      EXPECT_EQ((a[{i, j}]), 0.);
+    }
+  }
+}
+
 // A helper function that will be called from the JIT
 static void memrefMultiply(::StridedMemRefType<float, 2> *memref,
                            int32_t coefficient) {