[Mlir-commits] [mlir] [mlir][ptr] Use allocation size for type offsets (PR #236864)

[Akimasa Watanuki]

https://github.com/Men-cotton created https://github.com/llvm/llvm-project/pull/236864

Return the allocation stride from `ptr::TypeOffsetOp::getTypeSize()` so the C++ query agrees with the operation's element-offset semantics and its GEP-based lowering. For example, `i24` occupies 3 storage bytes but has a 4-byte allocation stride under the default data layout.

Assisted-by: Codex / GPT-6

>From c587b4b2b475d8997f7028e60e646c86cb083aa4 Mon Sep 17 00:00:00 2001
From: mencotton <mencotton0410 at gmail.com>
Date: Thu, 29 Oct 2026 11:41:09 -0500
Subject: [PATCH] [mlir][ptr] Use allocation size for type offsets

Return the allocation stride from `ptr::TypeOffsetOp::getTypeSize()` so the C++ query agrees with the operation's element-offset semantics and its GEP-based lowering. For example, `i24` occupies 3 storage bytes but has a 4-byte allocation stride under the default data layout.

Assisted-by: Codex / GPT-6
---
 mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp    |  4 +--
 mlir/unittests/Dialect/CMakeLists.txt     |  1 +
 mlir/unittests/Dialect/Ptr/CMakeLists.txt |  9 +++++++
 mlir/unittests/Dialect/Ptr/PtrTest.cpp    | 33 +++++++++++++++++++++++
 4 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 mlir/unittests/Dialect/Ptr/CMakeLists.txt
 create mode 100644 mlir/unittests/Dialect/Ptr/PtrTest.cpp

diff --git a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
index 6d61c11416733c..1f39c87359f75d 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
@@ -466,9 +466,9 @@ LogicalResult ToPtrOp::verify() {
 
 llvm::TypeSize TypeOffsetOp::getTypeSize(std::optional<DataLayout> layout) {
   if (layout)
-    return layout->getTypeSize(getElementType());
+    return layout->getTypeAllocSize(getElementType());
   DataLayout dl = DataLayout::closest(*this);
-  return dl.getTypeSize(getElementType());
+  return dl.getTypeAllocSize(getElementType());
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 5fd835189ccc76..3f2da1000de94d 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -14,6 +14,7 @@ add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
 add_subdirectory(OpenACC)
+add_subdirectory(Ptr)
 add_subdirectory(SCF)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
diff --git a/mlir/unittests/Dialect/Ptr/CMakeLists.txt b/mlir/unittests/Dialect/Ptr/CMakeLists.txt
new file mode 100644
index 00000000000000..e1153fb0771bbd
--- /dev/null
+++ b/mlir/unittests/Dialect/Ptr/CMakeLists.txt
@@ -0,0 +1,9 @@
+add_mlir_unittest(MLIRPtrTests
+  PtrTest.cpp
+)
+mlir_target_link_libraries(MLIRPtrTests
+  PRIVATE
+  MLIRIR
+  MLIRPtrDialect
+  MLIRDataLayoutInterfaces
+)
diff --git a/mlir/unittests/Dialect/Ptr/PtrTest.cpp b/mlir/unittests/Dialect/Ptr/PtrTest.cpp
new file mode 100644
index 00000000000000..a675d8cdfcb187
--- /dev/null
+++ b/mlir/unittests/Dialect/Ptr/PtrTest.cpp
@@ -0,0 +1,33 @@
+//===- PtrTest.cpp - Ptr dialect tests ------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Ptr/IR/PtrOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+TEST(Ptr, TypeOffsetIncludesAllocationPadding) {
+  MLIRContext context;
+  context.loadDialect<ptr::PtrDialect>();
+  OpBuilder builder(&context);
+  OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
+  builder.setInsertionPointToStart(module->getBody());
+  auto check = [&](Type type, uint64_t expected) {
+    auto offset = ptr::TypeOffsetOp::create(builder, builder.getUnknownLoc(),
+                                            builder.getIndexType(), type);
+    DataLayout layout(module.get());
+    EXPECT_EQ(offset.getTypeSize(layout).getFixedValue(), expected);
+    EXPECT_EQ(offset.getTypeSize().getFixedValue(), expected);
+  };
+  check(builder.getIntegerType(24), 4);
+  check(VectorType::get({3}, builder.getI32Type()), 16);
+}