[mlir][OneShotModuleBufferize] Add a new flag: no-analysis-func-filter

[mamrami]

Summary

OneShotModuleBufferize fails if the input IR cannot be analyzed.
One can set CopyBeforeWrite=true in order to skip analysis.
In that case, a buffer copy is inserted on every write.
This leads to many copies, also in FuncOps that could be analyzed.

This change aims to copy buffers only when it is a must.
When running OneShotModuleBufferize with CopyBeforeWrite=false,
FuncOps whose names are specified in noAnalysisFuncFilter will not be
analyzed. Ops in these FuncOps will not be analyzed as well.
They will be bufferized with CopyBeforeWrite=true,
while the other ops will be bufferized with CopyBeforeWrite=false.

[mamrami]

edit title

[springerm]

Instead of collecting FuncOps that could not be bufferized, can we pass a list of excluded FuncOps to the bufferization? Does that work in your flow?

We already have some infrastructure for that: BufferizationOptions::isOpAllowed. It returns false for ops that are excluded from the bufferization. I think it currently has no effect for FuncOps. Then you can bufferize with One-Shot Bufferize first, excluding all FuncOps of which you know that they cannot be bufferized. Then run the bufferization with copyBeforeWrite = true on all excluded FuncOps.

I think only small changes to mlir::bufferization::bufferizeModuleOp, mlir::bufferization::analyzeModuleOp and mlir::bufferization::insertTensorCopies should be needed.

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

1007 ↗

(On Diff #492448)

Ideally, no changes are needed to this file.

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

384

Put a check here:

if (!state.getOptions().isOpAllowed(funcOp)) continue;

[mamrami]

Addressing code review

[mamrami]

Then you can bufferize with One-Shot Bufferize first, excluding all FuncOps of which you know that they cannot be bufferized. Then run the bufferization with copyBeforeWrite = true on all excluded FuncOps.

There are cases where a function is called by another function, so bufferizing the caller is impossible when the callee is excluded.
Now the bufferization happens once on the entire IR so I don't have this problem.

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h

 //===- OneShotModuleBufferize.h - Bufferization across Func. Boundaries ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 namespace mlir {
 
 struct LogicalResult;
 class ModuleOp;
 
 namespace bufferization {
 struct BufferizationStatistics;
 class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state,
                               BufferizationStatistics *statistics = nullptr);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 ///
 /// Note: This function does not run One-Shot Analysis. No buffer copies are
-/// inserted unless `options.copyBeforeWrite` is set, in which case buffers are
-/// copied before every write.
-LogicalResult bufferizeModuleOp(ModuleOp moduleOp,
-                                const OneShotBufferizationOptions &options,
-                                BufferizationStatistics *statistics = nullptr);
+/// inserted except two cases:
+/// - `options.copyBeforeWrite` is set, in which case buffers are copied before
+///   every write.
+/// - `options.copyBeforeWrite` is not set and `analysisFilterFn` returns true
+///   for some FuncOps. These FuncOps were not analyzed. Buffer copies will be
+///   inserted only to these FuncOps.
+LogicalResult
+bufferizeModuleOp(ModuleOp moduleOp, const OneShotBufferizationOptions &options,
+                  BufferizationStatistics *statistics = nullptr,
+                  OpFilter::Entry::FilterFn analysisFilterFn = nullptr);
 
 /// Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
 void removeBufferizationAttributesInModule(ModuleOp moduleOp);
 
 /// Run One-Shot Module Bufferization on the given module. Performs a simple
 /// function call analysis to determine which function arguments are
 /// inplaceable. Then analyzes and bufferizes FuncOps one-by-one with One-Shot
 /// Bufferize.
 LogicalResult runOneShotModuleBufferize(
     ModuleOp moduleOp,
     const bufferization::OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics = nullptr);
+    BufferizationStatistics *statistics = nullptr,
+    OpFilter::Entry::FilterFn analysisFilterFn = nullptr);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

Option<"bufferizeFunctionBoundaries", "bufferize-function-boundaries",
            "Bufferize function boundaries (experimental).">,
     Option<"copyBeforeWrite", "copy-before-write", "bool", /*default=*/"false",
            "Skip the analysis. Make a buffer copy on every write.">,
     Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
+    ListOption<"noAnalysisFuncFilter", "no-analysis-func-filter", "std::string",
+               "Skip analysis of functions with these symbol names."
+               "Set copyBeforeWrite to true when bufferizing them.">,
     Option<"functionBoundaryTypeConversion",
            "function-boundary-type-conversion", "std::string",
            /*default=*/"\"infer-layout-map\"",
            "Controls layout maps when bufferizing function signatures.">,
     Option<"mustInferMemorySpace", "must-infer-memory-space", "bool",
            /*default=*/"false",
            "The memory space of an memref types must always be inferred. If "
            "unset, a default memory space of 0 is used otherwise.">,

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

if (!options) {
       opt.opFilter.allowOperation(filterFn);
     } else {
       opt = *options;
     }
 
     BufferizationStatistics statistics;
     ModuleOp moduleOp = getOperation();
     if (opt.bufferizeFunctionBoundaries) {
-      if (failed(runOneShotModuleBufferize(moduleOp, opt, &statistics))) {
+      OpFilter::Entry::FilterFn analysisFilterFn = nullptr;
+      // FuncOps whose names are specified in noAnalysisFuncFilter will not be
+      // analyzed. Ops in these FuncOps will not be analyzed as well.
+      if (this->noAnalysisFuncFilter.hasValue())
+        analysisFilterFn = [=](Operation *op) {
+          auto func = dyn_cast<func::FuncOp>(op);
+          if (!func)
+            func = op->getParentOfType<func::FuncOp>();
+          if (func)
+            return llvm::is_contained(noAnalysisFuncFilter, func.getSymName());
+          return false;
+        };
+      if (failed(runOneShotModuleBufferize(moduleOp, opt, &statistics,
+                                           analysisFilterFn))) {
         signalPassFailure();
         return;
       }
     } else {
+      assert(!this->noAnalysisFuncFilter.hasValue() &&
+             "invalid combination of bufferization flags");
       if (failed(runOneShotBufferize(moduleOp, opt, &statistics))) {
         signalPassFailure();
         return;
       }
     }
 
     // Set pass statistics.
     this->numBufferAlloc = statistics.numBufferAlloc;

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
+    if (!state.getOptions().isOpAllowed(funcOp))
+      continue;
+
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, state, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state, statistics)))
       return failure();
void mlir::bufferization::removeBufferizationAttributesInModule(
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 }
 
 LogicalResult mlir::bufferization::bufferizeModuleOp(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics) {
+    BufferizationStatistics *statistics,
+    OpFilter::Entry::FilterFn analysisFilterFn) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   IRRewriter rewriter(moduleOp.getContext());
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Bufferize functions.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
-    if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite,
+    bool copyBeforeWrite = options.copyBeforeWrite ||
+                           (analysisFilterFn && analysisFilterFn(funcOp));
+    if (failed(bufferizeOp(funcOp, options, copyBeforeWrite,
                            /*opFilter=*/nullptr, statistics)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
         LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }
 
 LogicalResult mlir::bufferization::runOneShotModuleBufferize(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics) {
+    BufferizationStatistics *statistics,
+    OpFilter::Entry::FilterFn analysisFilterFn) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&
          "invalid combination of bufferization flags");
   if (!options.copyBeforeWrite) {
+    if (!analysisFilterFn) {
       if (failed(insertTensorCopies(moduleOp, options, statistics)))
         return failure();
+    } else {
+      OneShotBufferizationOptions updatedOptions(options);
+      updatedOptions.opFilter.denyOperation(analysisFilterFn);
+      if (failed(insertTensorCopies(moduleOp, updatedOptions, statistics)))
+        return failure();
+    }
   }
   if (options.testAnalysisOnly)
     return success();
-  if (failed(bufferizeModuleOp(moduleOp, options, statistics)))
+
+  if (failed(
+          bufferizeModuleOp(moduleOp, options, statistics, analysisFilterFn)))
     return failure();
   return success();
 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-force-copy-before-write.mlir

+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 no-analysis-func-filter=contains_to_memref_op" -drop-equivalent-buffer-results --split-input-file | FileCheck %s
+
+// ToMemref ops do not pass analysis step. CopyBeforeWrite will be true only for the
+// FuncOp "contains_to_memref_op" since it is specified in no-analysis-func-filter.
+
+module {
+  // CHECK-LABEL:   func.func @foo(
+  // CHECK-SAME:                   %[[arg0:.*]]: memref<?xf32, strided<[?], offset: ?>>) {
+  func.func @foo(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK-NEXT:      %[[c0:.*]] = arith.constant 0 : index
+    %cst = arith.constant 1.000000e+00 : f32
+
+  // CHECK-NEXT:      %[[c1:.*]] = arith.constant 1.000000e+00 : f32
+    %c0 = arith.constant 0 : index
+
+  // CHECK-NEXT:      memref.store %[[c1]], %[[arg0]]{{\[}}%[[c0]]] : memref<?xf32, strided<[?], offset: ?>>
+    %inserted = tensor.insert %cst into %arg0[%c0] : tensor<?xf32>
+
+    return %inserted : tensor<?xf32>
+  }
+
+  // CHECK-LABEL:   func.func @contains_to_memref_op(
+  // CHECK-SAME:                                     %[[arg0:.*]]: memref<?xf32, strided<[?], offset: ?>>,
+  // CHECK-SAME:                                     %[[arg1:.*]]: index) -> vector<5xf32> {
+  func.func @contains_to_memref_op(%arg0: tensor<?xf32> {bufferization.writable = true}, %arg1: index) -> vector<5xf32> {
+
+    %0 = bufferization.to_memref %arg0 : memref<?xf32>
+
+    // CHECK:           %[[c0:.*]] = arith.constant 0 : index
+    %cst = arith.constant 0.000000e+00 : f32
+
+    // CHECK:           %[[dim:.*]] = memref.dim %[[arg0]], %[[c0]] : memref<?xf32, strided<[?], offset: ?>>
+    // CHECK:           %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?xf32>
+    // CHECK:           memref.copy %[[arg0]], %[[alloc]] : memref<?xf32, strided<[?], offset: ?>> to memref<?xf32>
+    // CHECK:           vector.transfer_read
+    %1 = vector.transfer_read %0[%arg1], %cst : memref<?xf32>, vector<5xf32>
+    return %1 : vector<5xf32>
+  }
+}
 No newline at end of file