[MLIR][NFC] Adding a missing option in OneShotBufferizePass #130691

lialan · 2025-03-11T00:22:52Z

#129850 removed let constructor for OneShotBufferizePass, however the list of tablegen options missed one inherited from BufferizationOptions.

it missed one option which is inherited from BufferizationOptions.

llvmbot · 2025-03-11T00:23:23Z

@llvm/pr-subscribers-mlir

Author: Alan Li (lialan)

Changes

#129850 removed let constructor for OneShotBufferizePass, however the list of tablegen options missed one inherited from BufferizationOptions.

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

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+4)
(modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+1)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 3bbb8b02c644e..b979f9df620d2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -450,6 +450,10 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> {
        Option<"bufferAlignment", "buffer-alignment", "uint64_t",
               /*default=*/"64",
               "Sets the alignment of newly allocated buffers.">,
+       Option<"enforceAliasingInvariants", "enforce-aliasing-invariants", "bool",
+              /*default=*/"true",
+              "Enforce aliasing Op/Operand/OpResult invariants with buffer"
+              " copies.">,
   ];
 
   let statistics = [
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index c0e3fca428376..f4891a317476d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -81,6 +81,7 @@ struct OneShotBufferizePass
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
+      opt.enforceAliasingInvariants = enforceAliasingInvariants;
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.dumpAliasSets = dumpAliasSets;
       opt.setFunctionBoundaryTypeConversion(

llvmbot · 2025-03-11T00:23:24Z

@llvm/pr-subscribers-mlir-bufferization

Author: Alan Li (lialan)

Changes

#129850 removed let constructor for OneShotBufferizePass, however the list of tablegen options missed one inherited from BufferizationOptions.

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

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+4)
(modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+1)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 3bbb8b02c644e..b979f9df620d2 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -450,6 +450,10 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> {
        Option<"bufferAlignment", "buffer-alignment", "uint64_t",
               /*default=*/"64",
               "Sets the alignment of newly allocated buffers.">,
+       Option<"enforceAliasingInvariants", "enforce-aliasing-invariants", "bool",
+              /*default=*/"true",
+              "Enforce aliasing Op/Operand/OpResult invariants with buffer"
+              " copies.">,
   ];
 
   let statistics = [
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index c0e3fca428376..f4891a317476d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -81,6 +81,7 @@ struct OneShotBufferizePass
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.analysisHeuristic = parseHeuristicOption(analysisHeuristic);
+      opt.enforceAliasingInvariants = enforceAliasingInvariants;
       opt.copyBeforeWrite = copyBeforeWrite;
       opt.dumpAliasSets = dumpAliasSets;
       opt.setFunctionBoundaryTypeConversion(

joker-eph

Don't we have test coverage for this option?

lialan · 2025-03-11T13:07:47Z

Don't we have test coverage for this option?

I looked and I don't think we have one.

joker-eph · 2025-03-11T13:29:08Z

Can we get some?

lialan · 2025-03-11T15:34:54Z

Can we get some?

To be fair most (if not all) of those options in the bufferization passes do not have corresponding tests, so #129850 skipped a lot of checks and got merged. If we want tests we would need to also cover all those. But historically no tests were in place, and adding tests is beyond the scope of this quick patch, and requires some coordination.

Our downstream does require to set this specific option, so this is just a quick patch to unblock downstream.

chelini · 2025-03-11T19:19:37Z

Can we get some?

To be fair most (if not all) of those options in the bufferization passes do not have corresponding tests, so #129850 skipped a lot of checks and got merged. If we want tests we would need to also cover all those. But historically no tests were in place, and adding tests is beyond the scope of this quick patch, and requires some coordination.

Our downstream does require to set this specific option, so this is just a quick patch to unblock downstream.

I don’t think #129850 “skipped a lot of checks.” The issue isn’t with #129850 itself but rather the lack of testing for options. I did not check this pass option yet, so I don't know how hard it is to add a test, but I will do it tomorrow.

lialan · 2025-03-11T19:31:28Z

I did not check this pass option yet, so I don't know how hard it is to add a test, but I will do it tomorrow.

@chelini Meanwhile can you approve this PR so we can get unblocked.

joker-eph · 2025-03-11T23:05:50Z

I don't think it is OK to approve this PR because of downstream issue: if you have downstream issues you can always patch MLIR when you're consuming it downstream. Rushing PRs upstream because of quirks of your private downstream repo does not seem right to me.

We shouldn't expose options that are not tested.

joker-eph

(Marking as needed changes)

matthias-springer · 2025-03-12T07:44:13Z

Are you using this flag together with copyBeforeWrite = true? If not, this will probably make your program bufferize incorrectly.

  /// Certain ops have aliasing OpOperand/OpResult invariants (e.g., scf.for).
  /// If this flag is set to `false`, those invariants are no longer enforced
  /// with buffer copies.
  ///
  /// Note: Deactivating this flag can lead to incorrect bufferization results
  /// when used incorrectly. This flag is useful with
  /// `AlwaysCopyAnalysisState` which bufferizes all writing tensor
  /// OpOperands out-of-place.
  bool enforceAliasingInvariants = true;

If you set copyBeforeWrite = true, it is not necessary to set enforceAliasingInvariants = false. This is the only place where we read the flag:

    if (!state.getOptions().enforceAliasingInvariants ||
        state.getOptions().copyBeforeWrite)
      return success();

It's been a while since I looked at this, but maybe we can remove enforceAliasingInvariants entirely.

chelini · 2025-03-12T08:07:09Z

Are you using this flag together with copyBeforeWrite = true? If not, this will probably make your program bufferize incorrectly.
  /// Certain ops have aliasing OpOperand/OpResult invariants (e.g., scf.for).
  /// If this flag is set to `false`, those invariants are no longer enforced
  /// with buffer copies.
  ///
  /// Note: Deactivating this flag can lead to incorrect bufferization results
  /// when used incorrectly. This flag is useful with
  /// `AlwaysCopyAnalysisState` which bufferizes all writing tensor
  /// OpOperands out-of-place.
  bool enforceAliasingInvariants = true;
If you set copyBeforeWrite = true, it is not necessary to change enforceAliasingInvariants. This is the only place where we read the flag:
    if (!state.getOptions().enforceAliasingInvariants ||
        state.getOptions().copyBeforeWrite)
      return success();
It's been a while since I looked at this, but maybe we can remove enforceAliasingInvariants entirely.

I am trying to add a simple test case but I don't see any difference in the IR and the bufferization analysis state, here the simple test case:

func.func @loop_with_aliasing(%t: tensor<?xf32>, %idx: index, %n: index) -> tensor<?xf32> {
  %c1 = arith.constant 1 : index
  %v = arith.constant 1.0 : f32
  %0 = scf.for %i = %idx to %n step %c1 iter_args(%arg0 = %t) -> tensor<?xf32> {
    // This insert operation modifies the tensor in-place
    %1 = tensor.insert %v into %arg0[%i] : tensor<?xf32>
    scf.yield %1 : tensor<?xf32>
  }
  return %0 : tensor<?xf32>
}

It triggers the option but there is no IR difference. I would have expected a copy if I set it to false here.

matthias-springer · 2025-03-12T08:13:54Z

What flag did you set? copyBeforeWrite or enforceAliasingInvariants? Only copyBeforeWrite will force a copy. When we have copies everywhere, some of the aliasing invariants can be relaxed. That's what the enforceAliasingInvariants is for. enforceAliasingInvariants by itself does not force a copy everywhere.

chelini · 2025-03-12T08:17:04Z

What flag did you set? copyBeforeWrite or enforceAliasingInvariants? Only copyBeforeWrite will force a copy. When we have copies everywhere, some of the aliasing invariants can be relaxed. That's what the enforceAliasingInvariants is for. enforceAliasingInvariants by itself does not force a copy everywhere.

I am testing with:

-one-shot-bufferize="bufferize-function-boundaries enforce-aliasing-invariants=false"
-one-shot-bufferize="bufferize-function-boundaries enforce-aliasing-invariants=true"

matthias-springer · 2025-03-12T08:17:29Z

On a high level, enforceAliasingInvariants = true is adding additional copies for certain cases where you yield a value that is defined outside of the loop. (Whereas copyBeforeWrite will add copies every time you write to a buffer.)

chelini · 2025-03-12T09:26:05Z

@lialan FYI: I commented out the option in IREE and run ctest --test-dir ../iree-build/ and I don't see any test failures.

[MLIR][NFC] Adding a missing option in OneShotBufferizePass

34c458a

it missed one option which is inherited from BufferizationOptions.

lialan requested a review from matthias-springer as a code owner March 11, 2025 00:22

lialan requested review from chelini and removed request for matthias-springer March 11, 2025 00:22

llvmbot added mlir mlir:bufferization Bufferization infrastructure labels Mar 11, 2025

lialan requested a review from joker-eph March 11, 2025 00:23

lialan mentioned this pull request Mar 11, 2025

Integrates/llvm 20250310: Bump to llvm/llvm-project@967ab7e iree-org/iree#20207

Merged

joker-eph requested a review from matthias-springer March 11, 2025 09:53

joker-eph reviewed Mar 11, 2025

View reviewed changes

joker-eph requested changes Mar 11, 2025

View reviewed changes

[MLIR][NFC] Adding a missing option in OneShotBufferizePass #130691

Are you sure you want to change the base?

[MLIR][NFC] Adding a missing option in OneShotBufferizePass #130691

Conversation

lialan commented Mar 11, 2025

Uh oh!

llvmbot commented Mar 11, 2025

Uh oh!

llvmbot commented Mar 11, 2025

Uh oh!

joker-eph left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lialan commented Mar 11, 2025

Uh oh!

joker-eph commented Mar 11, 2025

Uh oh!

lialan commented Mar 11, 2025

Uh oh!

chelini commented Mar 11, 2025

Uh oh!

lialan commented Mar 11, 2025

Uh oh!

joker-eph commented Mar 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

matthias-springer commented Mar 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chelini commented Mar 12, 2025

Uh oh!

matthias-springer commented Mar 12, 2025

Uh oh!

chelini commented Mar 12, 2025

Uh oh!

matthias-springer commented Mar 12, 2025

Uh oh!

chelini commented Mar 12, 2025

Uh oh!

Uh oh!

joker-eph left a comment •

edited

Loading

joker-eph commented Mar 11, 2025 •

edited

Loading

matthias-springer commented Mar 12, 2025 •

edited

Loading