[mlir] Fix correct memset range in OwningMemRef zero-init

[chokobole]

Problem: OwningMemRef zero-init can write past allocation

OwningMemref allocates with overprovision + manual alignment:

auto desiredAlignment = alignment.value_or(nextPowerOf2(sizeof(T))); 
void* data = malloc(size + desiredAlignment); // size = nElements * sizeof(T)

// Compute aligned view
auto addr = reinterpret_cast<uintptr_t>(data);
auto rem  = addr % desiredAlignment;
T* alignedData = (rem == 0)
                   ? static_cast<T*>(data)
                   : reinterpret_cast<T*>(addr + (desiredAlignment - rem));

The descriptor stores:

• descriptor.basePtr = data (the raw allocation)
• descriptor.data = alignedData (the aligned start)

When creating an OwningMemRef without init, we zero it. The buggy code does:

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.

Correct zero-init options

Choose one of the following:

1. Zero the entire allocated block from the raw base:

memset(descriptor.basePtr, 0,
       nElements * sizeof(T) + alignment.value_or(nextPowerOf2(sizeof(T))));

2. Zero only the logical payload (what the user sees) from the aligned start:

memset(descriptor.data, 0, nElements * sizeof(T));

Recommendation

I think one of the main root causes here is naming. At first, I assumed descriptor.data referred to the raw data pointer—but it doesn’t. Using more consistent naming could make the meaning clearer. For example: descriptor.data vs. descriptor.alignedData.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir-execution-engine

@llvm/pr-subscribers-mlir

Author: Ryan Kim (chokobole)

Changes

Problem: OwningMemRef zero-init can write past allocation

OwningMemref allocates with overprovision + manual alignment:

auto desiredAlignment = alignment.value_or(nextPowerOf2(sizeof(T))); 
void* data = malloc(size + desiredAlignment); // size = nElements * sizeof(T)

// Compute aligned view
auto addr = reinterpret_cast<uintptr_t>(data);
auto rem  = addr % desiredAlignment;
T* alignedData = (rem == 0)
                   ? static_cast<T*>(data)
                   : reinterpret_cast<T*>(addr + (desiredAlignment - rem));

The descriptor stores:

• descriptor.basePtr = data (the raw allocation)
• descriptor.data = alignedData (the aligned start)

When creating an OwningMemRef without init, we zero it. The buggy code does:

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.

Correct zero-init options

Choose one of the following:

1. Zero the entire allocated block from the raw base:

memset(descriptor.basePtr, 0,
       nElements * sizeof(T) + alignment.value_or(nextPowerOf2(sizeof(T))));

2. Zero only the logical payload (what the user sees) from the aligned start:

memset(descriptor.data, 0, nElements * sizeof(T));

Recommendation

I think one of the main root causes here is naming. At first, I assumed descriptor.data referred to the raw data pointer—but it doesn’t. Using more consistent naming could make the meaning clearer. For example: descriptor.data vs. descriptor.alignedData.

---

Full diff: https://github.com/llvm/llvm-project/pull/158200.diff

1 Files Affected:

• (modified) mlir/include/mlir/ExecutionEngine/MemRefUtils.h (+1-3)

diff --git a/mlir/include/mlir/ExecutionEngine/MemRefUtils.h b/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
index d66d757cb7a8e..035cd8c4593a9 100644
--- a/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
+++ b/mlir/include/mlir/ExecutionEngine/MemRefUtils.h
@@ -174,9 +174,7 @@ class OwningMemRef {
            it != end; ++it)
         init(*it, it.getIndices());
     } else {
-      memset(descriptor.data, 0,
-             nElements * sizeof(T) +
-                 alignment.value_or(detail::nextPowerOf2(sizeof(T))));
+      memset(descriptor.data, 0, nElements * sizeof(T));
     }
   }
   /// Take ownership of an existing descriptor with a custom deleter.

[joker-eph]

Is this an AI that detected this and wrote this PR? The way the description is phrased looks kind of verbose in the way AI agents tends to be.

Can you add a test please?

[joker-eph]

I suspect the confusion comes from auto [data, alignedData] = which should probably be auto [allocatedPtr, alignedData] = ... Can you update it line 167?

[chokobole]

I believe this has nothing to do with that line because this data comes from the descriptor. In my opinion, the root cause is a misuse of the descriptor.data. When you use descriptor.data, you need to use the size of your data. When you use the descriptor.basePtr, you need to use the sum of your data's size and its alignment.

[joker-eph]

Right, I know this isn't directly causing the bug, but the naming confusion is unfortunate and probably a contributing factor to the existence of the bug in the first place.

You also acknowledged it at the end of the description:

I think one of the main root causes here is naming. At first, I assumed descriptor.data referred to the raw data pointer—but it doesn’t.

[chokobole]

I suggest to change the member name in StridedMemRefType, if you see data member, it's hard to guess if it's aligned or even there's another pointer to free. That's why we have #153133.

template <typename T, int N>
struct StridedMemRefType {
  T *basePtr;
  T *data;

Maybe then I guess it's better to change them to allocatedPtr and alignedData? If you think it's good, then I'll proceed it

[joker-eph]

Why not if you'd like, but I was less ambitious and pointing that aligning the names to be self-consistent in this function was something that can easily be squeezed with this patch.

[chokobole]

Thanks—that makes sense. To keep this patch focused and avoid churn, I’ll make the naming self-consistent in this function and zero only the logical payload from the aligned start. Concretely:

Rename the local unpack to allocatedPtr / alignedData to avoid confusion.

When zero-init is requested, use alignedData with the payload size only.

- auto [data, alignedData] = allocateWithOverprovision(...);
+ auto [allocatedPtr, alignedData] = allocateWithOverprovision(...);
- descriptor = detail::makeStridedMemRefDescriptor<Rank>(data, alignedData,
+ descriptor = detail::makeStridedMemRefDescriptor<Rank>(allocatedPtr, alignedData,
                                                         shape, shapeAlloc);

  if (init) {

  } else {
-   memset(descriptor.data, 0, size + desiredAlignment); // could overrun
+   // Zero only the logical payload; do not write past the allocation.
+   memset(alignedData, 0, nElements * sizeof(T));
  }

This keeps the change scoped to this function, removes the overrun, and matches the mental model you suggested. If you’re good with this, I’ll push it here.

[joker-eph]

Yes, that's align with what I had in mind.

[chokobole]

Please review again!

[chokobole]

Is this an AI that detected this and wrote this PR? The way the description is phrased looks kind of verbose in the way AI agents tends to be.

Can you add a test please?

1. I used an AI to tidy up the PR description. However, the core fix is based on an issue I encountered during testing with OwnedMemRef which resulted in a free(): invalid pointer error.

2. I agree a test would be valuable. However, a proper unit test would require a complex memory simulation to prevent overwrites, which seems overly complicated for this situation.

[joker-eph]

Can we add a test along the line of: https://github.com/llvm/llvm-project/blob/main/mlir/unittests/ExecutionEngine/Invoke.cpp#L208-L214

Initializing the same kind of "widely unaligned memref", making it large enough, and then checking that is indeed zero-initialized?

[chokobole]

Added a unit test you suggested!

[joker-eph]

Can you add a comment? This test becomes quite non trivial now.

I wonder if it wouldn't just have been better to add a separate test that just check the 0 init and nothing else.

[chokobole]

I separated the test and limited to just check the 0 init only!

[github-actions]

@chokobole Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!