[Mlir-commits] [mlir] [mlir][ptr] Use allocation size for type offsets (PR #236864)
Akimasa Watanuki
llvmlistbot at llvm.org
Thu Oct 29 12:22:45 PDT 2026
https://github.com/Men-cotton updated https://github.com/llvm/llvm-project/pull/236864
>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 1/2] [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);
+}
>From 52aecb9c1da5dbe5580c6c7501f0fdface7856e7 Mon Sep 17 00:00:00 2001
From: mencotton <mencotton0410 at gmail.com>
Date: Thu, 29 Oct 2026 13:49:47 -0500
Subject: [PATCH 2/2] fix: [mlir][ptr] Test type offsets with custom data
layouts
Replace the type-offset unit test with lit coverage for default, custom, and inherited data layouts. Check allocation padding, scalable offsets, and explicit and implicit layout queries.
Assisted-by: Codex / GPT-6
---
mlir/test/Dialect/Ptr/type-offset.mlir | 51 +++++++++++++++++++
mlir/test/lib/Dialect/CMakeLists.txt | 1 +
mlir/test/lib/Dialect/Ptr/CMakeLists.txt | 11 +++++
mlir/test/lib/Dialect/Ptr/TestTypeOffset.cpp | 52 ++++++++++++++++++++
mlir/tools/mlir-opt/CMakeLists.txt | 1 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
mlir/unittests/Dialect/CMakeLists.txt | 1 -
mlir/unittests/Dialect/Ptr/CMakeLists.txt | 9 ----
mlir/unittests/Dialect/Ptr/PtrTest.cpp | 33 -------------
9 files changed, 118 insertions(+), 43 deletions(-)
create mode 100644 mlir/test/Dialect/Ptr/type-offset.mlir
create mode 100644 mlir/test/lib/Dialect/Ptr/CMakeLists.txt
create mode 100644 mlir/test/lib/Dialect/Ptr/TestTypeOffset.cpp
delete mode 100644 mlir/unittests/Dialect/Ptr/CMakeLists.txt
delete mode 100644 mlir/unittests/Dialect/Ptr/PtrTest.cpp
diff --git a/mlir/test/Dialect/Ptr/type-offset.mlir b/mlir/test/Dialect/Ptr/type-offset.mlir
new file mode 100644
index 00000000000000..f6269bf6176db7
--- /dev/null
+++ b/mlir/test/Dialect/Ptr/type-offset.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-opt %s --test-ptr-set-type-offset --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --test-ptr-set-type-offset=use-explicit-layout --split-input-file | FileCheck %s
+
+// CHECK-LABEL: @default_layout
+func.func @default_layout() {
+ // Include allocation padding, rather than just the bytes storing the value.
+ // CHECK: ptr.type_offset i24 {type_offset = 4 : index}
+ %0 = ptr.type_offset i24 : index
+ // CHECK: ptr.type_offset vector<3xi32> {type_offset = 16 : index}
+ %1 = ptr.type_offset vector<3xi32> : index
+ // CHECK: ptr.type_offset vector<[4]xi32> {type_offset = 16 : index, type_offset_scalable}
+ %2 = ptr.type_offset vector<[4]xi32> : index
+ return
+}
+
+// -----
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+ #dlti.dl_entry<i24, dense<[64, 128]> : vector<2xi64>>
+>} {
+ // CHECK-LABEL: @aligned_to_eight
+ func.func @aligned_to_eight() {
+ // Use the ABI alignment of 8 bytes, not the preferred alignment of 16.
+ // CHECK: ptr.type_offset i24 {type_offset = 8 : index}
+ %0 = ptr.type_offset i24 : i64
+ return
+ }
+
+ module @nested {
+ // CHECK-LABEL: @inherited_layout
+ func.func @inherited_layout() {
+ // A scope without a layout specification inherits its enclosing layout.
+ // CHECK: ptr.type_offset i24 {type_offset = 8 : index}
+ %0 = ptr.type_offset i24 : index
+ return
+ }
+ }
+}
+
+// -----
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<
+ #dlti.dl_entry<i24, dense<[128, 256]> : vector<2xi64>>
+>} {
+ // CHECK-LABEL: @aligned_to_sixteen
+ func.func @aligned_to_sixteen() {
+ // CHECK: ptr.type_offset i24 {type_offset = 16 : index}
+ %0 = ptr.type_offset i24 : index
+ return
+ }
+}
diff --git a/mlir/test/lib/Dialect/CMakeLists.txt b/mlir/test/lib/Dialect/CMakeLists.txt
index e31140a6f6ba78..a48292b39ca464 100644
--- a/mlir/test/lib/Dialect/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/CMakeLists.txt
@@ -13,6 +13,7 @@ add_subdirectory(MemRef)
add_subdirectory(Shard)
add_subdirectory(NVGPU)
add_subdirectory(OpenACC)
+add_subdirectory(Ptr)
add_subdirectory(SCF)
add_subdirectory(Shape)
add_subdirectory(SPIRV)
diff --git a/mlir/test/lib/Dialect/Ptr/CMakeLists.txt b/mlir/test/lib/Dialect/Ptr/CMakeLists.txt
new file mode 100644
index 00000000000000..776093348f795a
--- /dev/null
+++ b/mlir/test/lib/Dialect/Ptr/CMakeLists.txt
@@ -0,0 +1,11 @@
+add_mlir_library(MLIRPtrTestPasses
+ TestTypeOffset.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+)
+mlir_target_link_libraries(MLIRPtrTestPasses PUBLIC
+ MLIRDataLayoutInterfaces
+ MLIRIR
+ MLIRPass
+ MLIRPtrDialect
+)
diff --git a/mlir/test/lib/Dialect/Ptr/TestTypeOffset.cpp b/mlir/test/lib/Dialect/Ptr/TestTypeOffset.cpp
new file mode 100644
index 00000000000000..c07137777c1f72
--- /dev/null
+++ b/mlir/test/lib/Dialect/Ptr/TestTypeOffset.cpp
@@ -0,0 +1,52 @@
+//===- TestTypeOffset.cpp - Test Ptr type offsets ------------------------===//
+//
+// 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/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+struct TestPtrSetTypeOffset
+ : public PassWrapper<TestPtrSetTypeOffset, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPtrSetTypeOffset)
+
+ TestPtrSetTypeOffset() = default;
+ TestPtrSetTypeOffset(const TestPtrSetTypeOffset &pass) : PassWrapper(pass) {}
+
+ StringRef getArgument() const final { return "test-ptr-set-type-offset"; }
+ StringRef getDescription() const final {
+ return "Attach the queried type offset to each ptr.type_offset operation";
+ }
+
+ void runOnOperation() override {
+ Builder builder(&getContext());
+ getOperation()->walk([&](ptr::TypeOffsetOp op) {
+ llvm::TypeSize size = useExplicitLayout
+ ? op.getTypeSize(DataLayout::closest(op))
+ : op.getTypeSize();
+ op->setAttr("type_offset", builder.getIndexAttr(size.getKnownMinValue()));
+ if (size.isScalable())
+ op->setAttr("type_offset_scalable", builder.getUnitAttr());
+ });
+ }
+
+ Option<bool> useExplicitLayout{
+ *this, "use-explicit-layout",
+ llvm::cl::desc("Pass the enclosing data layout explicitly to the query"),
+ llvm::cl::init(false)};
+};
+} // namespace
+
+namespace mlir::test {
+void registerTestPtrSetTypeOffset() {
+ PassRegistration<TestPtrSetTypeOffset>();
+}
+} // namespace mlir::test
diff --git a/mlir/tools/mlir-opt/CMakeLists.txt b/mlir/tools/mlir-opt/CMakeLists.txt
index c607ccfa80e3ce..b34f7fdbe484f4 100644
--- a/mlir/tools/mlir-opt/CMakeLists.txt
+++ b/mlir/tools/mlir-opt/CMakeLists.txt
@@ -29,6 +29,7 @@ if(MLIR_INCLUDE_TESTS)
MLIRShardTest
MLIRNVGPUTestPasses
MLIROpenACCTestPasses
+ MLIRPtrTestPasses
MLIRSCFTestPasses
MLIRShapeTestPasses
MLIRSPIRVTestPasses
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index f4f8c5943580ec..bf8efb86372eb6 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -140,6 +140,7 @@ void registerTestOpaqueLoc();
void registerTestOpLoweringPasses();
void registerTestPadFusion();
void registerTestParallelLoopUnrollingPass();
+void registerTestPtrSetTypeOffset();
void registerTestRecursiveTypesPass();
void registerTestSCFUpliftWhileToFor();
void registerTestSCFUtilsPass();
@@ -289,6 +290,7 @@ static void registerTestPasses() {
mlir::test::registerTestOpLoweringPasses();
mlir::test::registerTestPadFusion();
mlir::test::registerTestParallelLoopUnrollingPass();
+ mlir::test::registerTestPtrSetTypeOffset();
mlir::test::registerTestRecursiveTypesPass();
mlir::test::registerTestSCFUpliftWhileToFor();
mlir::test::registerTestSCFUtilsPass();
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 3f2da1000de94d..5fd835189ccc76 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -14,7 +14,6 @@ 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
deleted file mode 100644
index e1153fb0771bbd..00000000000000
--- a/mlir/unittests/Dialect/Ptr/CMakeLists.txt
+++ /dev/null
@@ -1,9 +0,0 @@
-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
deleted file mode 100644
index a675d8cdfcb187..00000000000000
--- a/mlir/unittests/Dialect/Ptr/PtrTest.cpp
+++ /dev/null
@@ -1,33 +0,0 @@
-//===- 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);
-}
More information about the Mlir-commits
mailing list