[mlir] Add bufferization option for parallel region check

[Max191]

Handling parallel region RaW conflicts should usually be the responsibility of the source program, rather than bufferization analysis. However, to preserve current functionality, checks on parallel regions is put behind a bufferization in this PR, which is on by default. Default functionality will not change, but this PR enables the option to leave parallelism checks out of the bufferization analysis.

[llvmbot]

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

@llvm/pr-subscribers-mlir

Author: None (Max191)

Changes

Handling parallel region RaW conflicts should usually be the responsibility of the source program, rather than bufferization analysis. However, to preserve current functionality, checks on parallel regions is put behind a bufferization in this PR, which is on by default. Default functionality will not change, but this PR enables the option to leave parallelism checks out of the bufferization analysis.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+5)
• (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+1-1)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 2d8add82383be..2fda091e412ae 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -309,6 +309,11 @@ struct BufferizationOptions {
   /// bufferized or not.
   bool bufferizeFunctionBoundaries = false;
 
+  // Specifies whether to account for parallel regions in RaW analysis. If true,
+  // then writes inside of parallel regions that write to buffers defined
+  // outside of the parallel region will be given a new buffer.
+  bool checkParallelRegions = true;
+
   /// 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.
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 2d329a1f3d889..d0b4e0dd4383e 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -611,7 +611,7 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
   // Before going through the main RaW analysis, find cases where a buffer must
   // be privatized due to parallelism. If the result of a write is never read,
   // privatization is not necessary (and large parts of the IR are likely dead).
-  if (!usesRead.empty()) {
+  if (options.checkParallelRegions && !usesRead.empty()) {
     for (OpOperand *uConflictingWrite : usesWrite) {
       // Find the allocation point or last write (definition) of the buffer.
       // Note: In contrast to `findDefinitions`, this also returns results of

[MaheshRavishankar]

This looks reasonable to me. Maybe wait for Matthias to chime in, but to me this is OK.

[MaheshRavishankar]

Also, possible to add a test here?

[matthias-springer]

Generally speaking, the simpler the analysis the better.

Handling parallel region RaW conflicts should usually be the responsibility of the source program, rather than bufferization analysis.

I'm wondering how you're going to do that. You need to know how tensor ops are going to bufferize. (E.g., by querying the BufferizableOpInterface.) And you also need to predict alias sets, which are only decided during bufferization. So I think you will basically have to reimplement the bufferization RaW analysis.

I'm skeptical whether this is useful in practice. But it could certainly be useful as a testing flag (to debug bugs in the analysis), so feel free to merge.