[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