[Mlir-commits] [mlir] Added Memory Effect to assume_alignment (PR #139521)

[Shay Kleiman]

https://github.com/shay-kl updated https://github.com/llvm/llvm-project/pull/139521

>From bddbbe9a382963a94cf1d598b66111b8658a0048 Mon Sep 17 00:00:00 2001
From: Shay Kleiman <shayk at epgd045.me-corp.lan>
Date: Sun, 11 May 2025 16:22:44 +0300
Subject: [PATCH] Added MemoryEffect to assume_alignment

Assume_alignment has no trait which specifies how it interacts with
memory, this causes an issue in OwnershipBasedBufferDeallocation,
which require all operations which operate on buffers to have explicit
traits defining how the operation interacts with memory.

>From my understanding, technically the operation is pure, however to make
sure the operation doesn't get optimized away it has to have some side
effect. I defined it to have similar side effects to CF AssertOp as both
are asserts. I'm not sure if this is correct and would appreciate the
opinion of someone more experienced.
---
 mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td           | 2 +-
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp                   | 5 +++++
 .../OwnershipBasedBufferDeallocation/misc-other.mlir       | 7 +++++++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index d6d8161d3117b..856b033f401a0 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -142,7 +142,7 @@ class AllocLikeOp<string mnemonic,
 // AssumeAlignmentOp
 //===----------------------------------------------------------------------===//
 
-def AssumeAlignmentOp : MemRef_Op<"assume_alignment"> {
+def AssumeAlignmentOp : MemRef_Op<"assume_alignment",[DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary =
       "assertion that gives alignment information to the input memref";
   let description = [{
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index a0237c18cf2fe..1872a63f93c93 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -526,6 +526,11 @@ LogicalResult AssumeAlignmentOp::verify() {
     return emitOpError("alignment must be power of 2");
   return success();
 }
+void AssumeAlignmentOp::getEffects(
+  SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+      &effects) {
+effects.emplace_back(MemoryEffects::Write::get());
+}
 
 //===----------------------------------------------------------------------===//
 // CastOp
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
index 05e52848ca877..4e0e743fc3ed9 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/misc-other.mlir
@@ -10,4 +10,11 @@ func.func @func_with_assert(%arg0: index, %arg1: index) {
   %0 = arith.cmpi slt, %arg0, %arg1 : index
   cf.assert %0, "%arg0 must be less than %arg1"
   return
+}
+
+// CHECK-LABEL: func @func_with_assume_alignment(
+//       CHECK: memref.assume_alignment %arg0, 64 : memref<128xi8>
+func.func @func_with_assume_alignment(%arg0: memref<128xi8>) {
+  memref.assume_alignment %arg0, 64 : memref<128xi8>
+  return
 }
\ No newline at end of file