Skip to content

[mlir][MemRef] Changed AssumeAlignment into a Pure ViewLikeOp #139521

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

shay-kl
Copy link
Contributor

@shay-kl shay-kl commented May 12, 2025

Made AssumeAlignment a ViewLikeOp that returns a new SSA memref equal
to its memref argument and made it have Pure trait. This
gives it a defined memory effect that matches what it does in practice
and makes it behave nicely with optimizations which won't get rid of it
unless its result isn't being used.

Copy link

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

@llvmbot
Copy link
Member

llvmbot commented May 12, 2025

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Shay Kleiman (shay-kl)

Changes

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 and assume_alignment gets lowered to it in RuntimeOpVerification. I'm not sure if this is correct and would appreciate the opinion of someone more experienced.


Full diff: https://github.com/llvm/llvm-project/pull/139521.diff

2 Files Affected:

  • (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+1-1)
  • (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+5)
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

@llvmbot
Copy link
Member

llvmbot commented May 12, 2025

@llvm/pr-subscribers-mlir-memref

Author: Shay Kleiman (shay-kl)

Changes

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 and assume_alignment gets lowered to it in RuntimeOpVerification. I'm not sure if this is correct and would appreciate the opinion of someone more experienced.


Full diff: https://github.com/llvm/llvm-project/pull/139521.diff

2 Files Affected:

  • (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+1-1)
  • (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+5)
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

@matthias-springer
Copy link
Member

I'm not quite sure how we handle undefined behavior and side effects in this case. Maybe @joker-eph or @kuhar have some opinion.

@shay-kl shay-kl force-pushed the memrefops/add-memory-effect-to-assume-aligned branch from 4c694da to 3aae81c Compare May 12, 2025 13:52
@llvmbot llvmbot added the mlir:bufferization Bufferization infrastructure label May 12, 2025
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.
@shay-kl shay-kl force-pushed the memrefops/add-memory-effect-to-assume-aligned branch from 3aae81c to bddbbe9 Compare May 12, 2025 14:07
@shay-kl
Copy link
Contributor Author

shay-kl commented May 12, 2025

I added a lit test to demonstrate the issue with the operation not passing OwnershipBasedBufferDeallocation.

@joker-eph
Copy link
Collaborator

I'm not quite sure how we handle undefined behavior and side effects in this case.

NoMemoryEffect means no side effects beyond undefined behavior, and Pure means no side-effects and no UB.

Copy link
Contributor

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Saying it has write semantics is a big pessimization. This is more a compiler hint. If an optimization is removing it, that needs to be fixed in the optimization AFAICS. For example CSE does no CSE operation with no results (even if it has no side-effects).

@matthias-springer
Copy link
Member

@shay-kl Is the buffer deallocation pass happy if you use NoMemoryEffect?

@shay-kl
Copy link
Contributor Author

shay-kl commented May 13, 2025

@matthias-springer
The buffer deallocation pass is happy if I use NoMemoryEffect, however in that case the operation is considered trivially dead and gets optimized away by cse.

@shay-kl
Copy link
Contributor Author

shay-kl commented May 13, 2025

@MaheshRavishankar
So you're saying that implicitly if an operation has no memory effects and no results then it makes sense for it to never be optimized away?
Would it be a better idea to explicitly define a trait for this type of operation or should we just deduct that they're an annotation and shouldn't be removed?

@matthias-springer
Copy link
Member

matthias-springer commented May 13, 2025

I see... The canonicalizer pass would then also remove the op because it is "trivially dead". This is not just a CSE problem.

I'd be in favor of giving the op a result and NoMemoryEffect. We've been using such a design internally and it seems to work well.

@shay-kl
Copy link
Contributor Author

shay-kl commented May 13, 2025

@matthias-springer
So the intent would be something like this
%res1 = memref.alloc() : memref<128xi8>
%res2 = memref.assume_alignment %res1, 64 : memref<128xi8>

and from now on we use res2 instead of res1?

@joker-eph
Copy link
Collaborator

joker-eph commented May 13, 2025

Saying it has write semantics is a big pessimization.

Right now without any traits, it has "unknown effect" which is most conservative situation (it has to be read and write to anything).

Using dataflow is a cleaner design.

Made AssumeAlignment a ViewLikeOp that returns a new SSA memref equal
to its memref argument and made it have NoMemoryEffect trait. This
gives it a defined memory effect that matches what it does in practice
and makes it behave nicely with optimizations which won't get rid of it
unless its result isn't being used.
@shay-kl
Copy link
Contributor Author

shay-kl commented May 14, 2025

@matthias-springer
I gave assume alignment a result and changed all its uses accordingly and it does seem cleaner this way. Thanks for the suggestion.

@joker-eph
Copy link
Collaborator

Can you update the title and description?

@shay-kl shay-kl changed the title Added Memory Effect to assume_alignment [mlir][MemRef] Changed AssumeAlignment into a ViewLikeOp May 14, 2025
Assume Alignment is now pure with its ub now deferred with poison
@shay-kl shay-kl changed the title [mlir][MemRef] Changed AssumeAlignment into a ViewLikeOp [mlir][MemRef] Changed AssumeAlignment into a Pure ViewLikeOp May 15, 2025
Set AsmResultName for the operation to be assume_align
@joker-eph
Copy link
Collaborator

LGTM but let's wait for others to double check again.

@shay-kl shay-kl requested a review from MaheshRavishankar May 15, 2025 13:04
Copy link
Member

@matthias-springer matthias-springer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

Copy link
Contributor

@MaheshRavishankar MaheshRavishankar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

THis seems good to me, but I dont fully follow why the ViewLikeOpInterface here?

@MaheshRavishankar MaheshRavishankar dismissed their stale review May 15, 2025 23:48

Dismissing my review cause the main blocker that I had is resolved. Still have question(s) though.

@shay-kl
Copy link
Contributor Author

shay-kl commented May 16, 2025

THis seems good to me, but I dont fully follow why the ViewLikeOpInterface here?

I assumed that an identity operation would count as a view, and that since assume_alignment returns the same buffer, it could be beneficial for some generic passes. But since it doesn’t change shape or layout, I’m fine with removing the interface if that’s clearer.

@shay-kl shay-kl force-pushed the memrefops/add-memory-effect-to-assume-aligned branch from 637364b to 21bb4fd Compare May 16, 2025 11:32
Copy link

github-actions bot commented May 17, 2025

✅ With the latest revision this PR passed the C/C++ code formatter.

@shay-kl shay-kl force-pushed the memrefops/add-memory-effect-to-assume-aligned branch from 21bb4fd to 9a4e9be Compare May 18, 2025 05:45
Changed assertion to assumption and improved description of poison
@shay-kl shay-kl force-pushed the memrefops/add-memory-effect-to-assume-aligned branch from 9a4e9be to 1c158c1 Compare May 18, 2025 05:51
@shay-kl
Copy link
Contributor Author

shay-kl commented May 18, 2025

I don't have permissions, can you please land it?

@AviadCo AviadCo merged commit ffb9bbf into llvm:main May 18, 2025
11 checks passed
Copy link

@shay-kl Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

@llvm-ci
Copy link
Collaborator

llvm-ci commented May 18, 2025

LLVM Buildbot has detected a new failure on builder mlir-nvidia running on mlir-nvidia while building mlir at step 7 "test-build-check-mlir-build-only-check-mlir".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/138/builds/13299

Here is the relevant piece of the build log for the reference
Step 7 (test-build-check-mlir-build-only-check-mlir) failure: test (failure)
******************** TEST 'MLIR :: Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir' FAILED ********************
Exit Code: 1

Command Output (stdout):
--
# RUN: at line 1
/vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/mlir-opt /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir -generate-runtime-verification      -expand-strided-metadata      -test-cf-assert      -convert-to-llvm |  /vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/mlir-runner -e main -entry-point-result=void      -shared-libs=/vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/lib/libmlir_runner_utils.so 2>&1 |  /vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/FileCheck /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir
# executed command: /vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/mlir-opt /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir -generate-runtime-verification -expand-strided-metadata -test-cf-assert -convert-to-llvm
# executed command: /vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/mlir-runner -e main -entry-point-result=void -shared-libs=/vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/lib/libmlir_runner_utils.so
# executed command: /vol/worker/mlir-nvidia/mlir-nvidia/llvm.obj/bin/FileCheck /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir
# .---command stderr------------
# | /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir:31:17: error: CHECK-NEXT: expected string not found in input
# |  // CHECK-NEXT: "memref.assume_alignment"(%{{.*}}) <{alignment = 4 : i32}> : (memref<1xf32>) -> ()
# |                 ^
# | <stdin>:1:38: note: scanning from here
# | ERROR: Runtime op verification failed
# |                                      ^
# | <stdin>:2:1: note: possible intended match here
# | %23 = "memref.assume_alignment"(%17) <{alignment = 4 : i32}> : (memref<1xf32>) -> memref<1xf32>
# | ^
# | 
# | Input file: <stdin>
# | Check file: /vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir
# | 
# | -dump-input=help explains the following input dump.
# | 
# | Input was:
# | <<<<<<
# |            1: ERROR: Runtime op verification failed 
# | next:31'0                                          X error: no match found
# |            2: %23 = "memref.assume_alignment"(%17) <{alignment = 4 : i32}> : (memref<1xf32>) -> memref<1xf32> 
# | next:31'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# | next:31'1     ?                                                                                                possible intended match
# |            3: ^ memref is not aligned to 4 
# | next:31'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# |            4: Location: loc("/vol/worker/mlir-nvidia/mlir-nvidia/llvm.src/mlir/test/Integration/Dialect/MemRef/assume-alignment-runtime-verification.mlir":34:3) 
# | next:31'0     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# | >>>>>>
# `-----------------------------
# error: command failed with exit status: 1

--

********************


@AviadCo
Copy link
Contributor

AviadCo commented May 18, 2025

LLVM Buildbot has detected a new failure on builder mlir-nvidia running on mlir-nvidia while building mlir at step 7 "test-build-check-mlir-build-only-check-mlir".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/138/builds/13299

Here is the relevant piece of the build log for the reference

Fixed failing lit test here: #140447

Groverkss added a commit to iree-org/llvm-project that referenced this pull request May 21, 2025
Groverkss added a commit to iree-org/iree that referenced this pull request May 27, 2025
ajaden-codes pushed a commit to Jaddyen/llvm-project that referenced this pull request Jun 6, 2025
…39521)

Made AssumeAlignment a ViewLikeOp that returns a new SSA memref equal
to its memref argument and made it have Pure trait. This
gives it a defined memory effect that matches what it does in practice
and makes it behave nicely with optimizations which won't get rid of it
unless its result isn't being used.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

7 participants