[Mlir-commits] [mlir] Added trait NoMemoryEffect to assume_alignment (PR #139520)

[Shay Kleiman]

https://github.com/shay-kl created https://github.com/llvm/llvm-project/pull/139520

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.

>From 4c694dafd63d16b0d95f6b1c132afcd039138ecc 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 trait NoMemoryEffect 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 +++++
 2 files changed, 6 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