aboutsummaryrefslogtreecommitdiff
path: root/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
blob: aa53f94fe839d9ec6212510d5ec3d37d20646182 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
//===- OneShotModuleBufferize.cpp - 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
//
//===----------------------------------------------------------------------===//
//
// Module Bufferization is an extension of One-Shot Bufferize that
// bufferizes function boundaries. It provides `BufferizableOpInterface`
// implementations for FuncOp, CallOp and ReturnOp. Although it is named
// Module Bufferization, it may operate on any SymbolTable.
//
// Module Bufferization is run via `runOneShotModuleBufferize(SymbolTableOp,
// ...)`. This function analyzes the given op and determines the order of
// analysis and bufferization: Functions that are called are processed before
// their respective callers.
//
// After analyzing a FuncOp, additional information about its bbArgs is
// gathered and stored in `FuncAnalysisState`.
//
// * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
// for
//   each tensor return value (if any).
// * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
//   read/written.
//
// Module Bufferization implements the following calling convention.
//
// * In the absence of conflicts within a FuncOp, the FuncOp's bbArgs may always
//   be written to in-place.
// * If a tensor operand of a CallOp is read after the CallOp, the operand of
//   the CallOp must bufferize out-of-place.
//
// Example: The tensor.insert op bufferizes in-place because it is allowed to
// modify the buffer of `%t1` directly. The CallOp in `caller` must bufferize
// out-of-place because `%t0` is modified by the callee but read by the
// tensor.extract op. The analysis of CallOps decides whether an OpOperand must
// bufferize out-of-place based on results of `funcOpBbArgReadWriteAnalysis`.
// ```
// func @callee(%t1 : tensor<?xf32>) -> tensor<?xf32> {
//   %f = ... : f32
//   %0 = tensor.insert %f into %t1[...] : tensor<?xf32>
//   return %0 : tensor<?xf32>
// }
//
// func @caller() -> () {
//   %t0 = ... : tensor<?xf32>
//   %1 = call @callee(%t0) : (tensor<?xf32>) -> (tensor<?xf32>)
//   %2 = tensor.extract %1[...]  : tensor<?xf32>
// }
// ```
//
// Note: If a function is external, `funcOpBbArgReadWriteAnalysis` cannot
// analyze the function body. In such a case, the CallOp analysis conservatively
// assumes that each tensor OpOperand is both read and written.
//
// TODO: Add FuncOp attributes so that bbArgs of external FuncOps can be marked
// as "not reading" and/or "not writing".

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

#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"

using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::bufferization::func_ext;

/// A mapping of FuncOps to their callers.
using FuncCallerMap = DenseMap<func::FuncOp, DenseSet<Operation *>>;

/// Get or create FuncAnalysisState.
static FuncAnalysisState &
getOrCreateFuncAnalysisState(OneShotAnalysisState &state) {
  auto *result = state.getExtension<FuncAnalysisState>();
  if (result)
    return *result;
  return state.addExtension<FuncAnalysisState>();
}

namespace {

/// Annotate IR with the results of the analysis. For testing purposes only.
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
                                          BlockArgument bbArg) {
  const char *kEquivalentArgsAttr = "__equivalent_func_args__";
  Operation *op = returnVal.getOwner();

  SmallVector<int64_t> equivBbArgs;
  if (op->hasAttr(kEquivalentArgsAttr)) {
    auto attr = cast<ArrayAttr>(op->getAttr(kEquivalentArgsAttr));
    equivBbArgs = llvm::to_vector<4>(llvm::map_range(attr, [](Attribute a) {
      return cast<IntegerAttr>(a).getValue().getSExtValue();
    }));
  } else {
    equivBbArgs.append(op->getNumOperands(), -1);
  }
  equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();

  OpBuilder b(op->getContext());
  op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
}

/// Store function BlockArguments that are equivalent to/aliasing a returned
/// value in FuncAnalysisState.
static LogicalResult
aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                             FuncAnalysisState &funcState) {
  if (funcOp.getBody().empty()) {
    // No function body available. Conservatively assume that every tensor
    // return value may alias with any tensor bbArg.
    FunctionType type = funcOp.getFunctionType();
    for (const auto &inputIt : llvm::enumerate(type.getInputs())) {
      if (!isa<TensorType>(inputIt.value()))
        continue;
      for (const auto &resultIt : llvm::enumerate(type.getResults())) {
        if (!isa<TensorType>(resultIt.value()))
          continue;
        int64_t returnIdx = resultIt.index();
        int64_t bbArgIdx = inputIt.index();
        funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
      }
    }
    return success();
  }

  // Find all func.return ops.
  SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
  assert(!returnOps.empty() && "expected at least one ReturnOp");

  // Build alias sets. Merge all aliases from all func.return ops.
  for (BlockArgument bbArg : funcOp.getArguments()) {
    if (isa<RankedTensorType>(bbArg.getType())) {
      int64_t bbArgIdx = bbArg.getArgNumber();
      // Store aliases in a set, so that we don't add the same alias twice.
      SetVector<int64_t> aliases;
      for (func::ReturnOp returnOp : returnOps) {
        for (OpOperand &returnVal : returnOp->getOpOperands()) {
          if (isa<RankedTensorType>(returnVal.get().getType())) {
            int64_t returnIdx = returnVal.getOperandNumber();
            if (state.areAliasingBufferizedValues(returnVal.get(), bbArg))
              aliases.insert(returnIdx);
          }
        }
      }
      for (int64_t alias : aliases)
        funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(alias);
    }
  }

  // Build equivalence sets.
  // Helper function that finds an equivalent block argument index for the
  // given OpOperand. Return std::nullopt if no equivalent block argument could
  // be found.
  auto findEquivalentBlockArgIdx =
      [&](OpOperand &opOperand) -> std::optional<int64_t> {
    Value v = opOperand.get();
    if (!isa<TensorType>(v.getType()))
      return std::nullopt;
    for (BlockArgument bbArg : funcOp.getArguments()) {
      if (isa<RankedTensorType>(bbArg.getType())) {
        if (state.areEquivalentBufferizedValues(v, bbArg)) {
          if (state.getOptions().testAnalysisOnly)
            annotateEquivalentReturnBbArg(opOperand, bbArg);
          return bbArg.getArgNumber();
        }
      }
    }
    return std::nullopt;
  };

  int64_t numResults = returnOps.front()->getNumOperands();
  for (int64_t i = 0; i < numResults; ++i) {
    // Find the equivalent block argument index for the i-th operand of the
    // first func.return op.
    std::optional<int64_t> maybeEquiv =
        findEquivalentBlockArgIdx(returnOps.front()->getOpOperand(i));
    if (!maybeEquiv.has_value())
      continue;
    int64_t bbArgIdx = *maybeEquiv;
    bool allEquiv = true;

    // Check if all other func.return ops have the same equivalent block
    // argument for the i-th operand. In contrast to aliasing information,
    // which is just "merged", equivalence information must match across all
    // func.return ops.
    for (func::ReturnOp returnOp : ArrayRef(returnOps).drop_front()) {
      std::optional<int64_t> maybeEquiv =
          findEquivalentBlockArgIdx(returnOp->getOpOperand(i));
      if (maybeEquiv != bbArgIdx) {
        allEquiv = false;
        break;
      }
    }

    // All func.return ops have the same equivalent block argument for the i-th
    // operand.
    if (allEquiv)
      funcState.equivalentFuncArgs[funcOp][i] = bbArgIdx;
  }

  return success();
}

static void annotateFuncArgAccess(func::FuncOp funcOp, int64_t idx, bool isRead,
                                  bool isWritten) {
  OpBuilder b(funcOp.getContext());
  Attribute accessType;
  if (isRead && isWritten) {
    accessType = b.getStringAttr("read-write");
  } else if (isRead) {
    accessType = b.getStringAttr("read");
  } else if (isWritten) {
    accessType = b.getStringAttr("write");
  } else {
    accessType = b.getStringAttr("none");
  }
  funcOp.setArgAttr(idx, BufferizationDialect::kBufferAccessAttrName,
                    accessType);
}

/// Determine which FuncOp bbArgs are read and which are written. When run on a
/// function with unknown ops, we conservatively assume that such ops bufferize
/// to a read + write.
static LogicalResult
funcOpBbArgReadWriteAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                             FuncAnalysisState &funcState) {
  for (int64_t idx = 0, e = funcOp.getFunctionType().getNumInputs(); idx < e;
       ++idx) {
    // Skip non-tensor arguments.
    if (!isa<TensorType>(funcOp.getFunctionType().getInput(idx)))
      continue;
    bool isRead;
    bool isWritten;
    if (auto accessAttr = funcOp.getArgAttrOfType<StringAttr>(
            idx, BufferizationDialect::kBufferAccessAttrName)) {
      // Buffer access behavior is specified on the function. Skip the analysis.
      StringRef str = accessAttr.getValue();
      isRead = str == "read" || str == "read-write";
      isWritten = str == "write" || str == "read-write";
    } else if (funcOp.getBody().empty()) {
      // If the function has no body, conservatively assume that all args are
      // read + written.
      isRead = true;
      isWritten = true;
    } else {
      // Analyze the body of the function.
      BlockArgument bbArg = funcOp.getArgument(idx);
      isRead = state.isValueRead(bbArg);
      isWritten = state.isValueWritten(bbArg);
    }

    if (state.getOptions().testAnalysisOnly)
      annotateFuncArgAccess(funcOp, idx, isRead, isWritten);
    if (isRead)
      funcState.readBbArgs[funcOp].insert(idx);
    if (isWritten)
      funcState.writtenBbArgs[funcOp].insert(idx);
  }

  return success();
}
} // namespace

/// Remove bufferization attributes on FuncOp arguments.
static void removeBufferizationAttributes(BlockArgument bbArg) {
  auto funcOp = cast<func::FuncOp>(bbArg.getOwner()->getParentOp());
  funcOp.removeArgAttr(bbArg.getArgNumber(),
                       BufferizationDialect::kBufferLayoutAttrName);
  funcOp.removeArgAttr(bbArg.getArgNumber(),
                       BufferizationDialect::kWritableAttrName);
}

/// Return the func::FuncOp called by `callOp`.
static func::FuncOp
getCalledFunction(func::CallOp callOp,
                  mlir::SymbolTableCollection &symbolTable) {
  SymbolRefAttr sym =
      llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
  if (!sym)
    return nullptr;
  return dyn_cast_or_null<func::FuncOp>(
      symbolTable.lookupNearestSymbolFrom(callOp, sym));
}

/// Return "true" if the given function signature has tensor semantics.
static bool hasTensorSignature(func::FuncOp funcOp) {
  return llvm::any_of(funcOp.getFunctionType().getInputs(),
                      llvm::IsaPred<TensorType>) ||
         llvm::any_of(funcOp.getFunctionType().getResults(),
                      llvm::IsaPred<TensorType>);
}

/// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
/// callee-caller order (i.e., callees without callers first). Store all
/// remaining functions (i.e., the ones that call each other recursively) in
/// `remainingFuncOps`. Does not traverse nested symbol tables.
///
/// Store the map of FuncOp to all its callers in `callerMap`.
///
/// Return `failure()` if we are unable to retrieve the called FuncOp from
/// any func::CallOp.
static LogicalResult getFuncOpsOrderedByCalls(
    Operation *moduleOp, SmallVectorImpl<func::FuncOp> &orderedFuncOps,
    SmallVectorImpl<func::FuncOp> &remainingFuncOps, FuncCallerMap &callerMap,
    SymbolTableCollection &symbolTables) {
  // For each FuncOp, the set of functions called by it (i.e. the union of
  // symbols of all nested func::CallOp).
  DenseMap<func::FuncOp, DenseSet<func::FuncOp>> calledBy;
  // For each FuncOp, the number of func::CallOp it contains.
  DenseMap<func::FuncOp, unsigned> numberCallOpsContainedInFuncOp;
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (func::FuncOp funcOp : block.getOps<func::FuncOp>()) {
        // Collect function calls and populate the caller map.
        numberCallOpsContainedInFuncOp[funcOp] = 0;
        WalkResult res = funcOp.walk([&](func::CallOp callOp) -> WalkResult {
          func::FuncOp calledFunction = getCalledFunction(callOp, symbolTables);
          assert(calledFunction && "could not retrieved called func::FuncOp");
          // If the called function does not have any tensors in its signature,
          // then it is not necessary to bufferize the callee before the caller.
          if (!hasTensorSignature(calledFunction))
            return WalkResult::skip();

          callerMap[calledFunction].insert(callOp);
          if (calledBy[calledFunction].insert(funcOp).second) {
            numberCallOpsContainedInFuncOp[funcOp]++;
          }
          return WalkResult::advance();
        });
        if (res.wasInterrupted())
          return failure();
      }
    }
  }

  // Iteratively remove function operations that do not call any of the
  // functions remaining in the callCounter map and add them to ordered list.
  SmallVector<func::FuncOp> worklist;

  for (const auto &entry : numberCallOpsContainedInFuncOp) {
    if (entry.second == 0)
      worklist.push_back(entry.first);
  }

  while (!worklist.empty()) {
    func::FuncOp func = worklist.pop_back_val();
    orderedFuncOps.push_back(func);

    for (func::FuncOp caller : calledBy[func]) {
      auto &count = numberCallOpsContainedInFuncOp[caller];

      if (--count == 0)
        worklist.push_back(caller);
    }

    numberCallOpsContainedInFuncOp.erase(func);
  }

  // Put all other functions in the list of remaining functions. These are
  // functions that call each other circularly.
  for (auto it : numberCallOpsContainedInFuncOp)
    remainingFuncOps.push_back(it.first);

  return success();
}

/// Helper function that extracts the source from a memref.cast. If the given
/// value is not a memref.cast result, simply returns the given value.
static Value unpackCast(Value v) {
  auto castOp = v.getDefiningOp<memref::CastOp>();
  if (!castOp)
    return v;
  return castOp.getSource();
}

/// Helper function that returns the return types (skipping casts) of the given
/// func.return ops. This function returns as many types as the return ops have
/// operands. If the i-th operand is not the same for all func.return ops, then
/// the i-th returned type is an "empty" type.
static SmallVector<Type> getReturnTypes(SmallVector<func::ReturnOp> returnOps) {
  assert(!returnOps.empty() && "expected at least one ReturnOp");
  int numOperands = returnOps.front()->getNumOperands();

  // Helper function that unpacks memref.cast ops and returns the type.
  auto getSourceType = [&](Value v) { return unpackCast(v).getType(); };

  SmallVector<Type> result;
  for (int i = 0; i < numOperands; ++i) {
    // Get the type of the i-th operand of the first func.return ops.
    Type t = getSourceType(returnOps.front()->getOperand(i));

    // Check if all other func.return ops have a matching operand type.
    for (int j = 1; j < static_cast<int>(returnOps.size()); ++j)
      if (getSourceType(returnOps[j]->getOperand(i)) != t)
        t = Type();

    result.push_back(t);
  }

  return result;
}

/// Fold return values that are memref casts and update function return types.
///
/// During FuncOp bufferization, the exact type of the returned memrefs (if any)
/// is not known yet. Therefore, the bufferization uses memref types with the
/// most generic layout map as function return types. After bufferizing the
/// entire function body, a more concise memref type can potentially be used for
/// the return type of the function.
static void foldMemRefCasts(func::FuncOp funcOp) {
  // There is nothing to do for bodiless ops.
  if (funcOp.getBody().empty())
    return;

  // Compute the common result types of all return ops.
  SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
  SmallVector<Type> resultTypes = getReturnTypes(returnOps);

  // Remove direct casts.
  for (func::ReturnOp returnOp : returnOps) {
    for (OpOperand &operand : returnOp->getOpOperands()) {
      // Bail if no common result type was found.
      if (resultTypes[operand.getOperandNumber()]) {
        operand.set(unpackCast(operand.get()));
      }
    }
  }

  // Fill in the missing result types that were not the same among all
  // func.return ops.
  for (int i = 0; i < static_cast<int>(resultTypes.size()); ++i) {
    if (resultTypes[i])
      continue;
    resultTypes[i] = funcOp.getFunctionType().getResult(i);
  }

  // Update the function type.
  auto newFuncType = FunctionType::get(
      funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
  funcOp.setType(newFuncType);
}

LogicalResult
mlir::bufferization::analyzeModuleOp(Operation *moduleOp,
                                     OneShotAnalysisState &state,
                                     BufferizationStatistics *statistics) {
  assert(state.getOptions().bufferizeFunctionBoundaries &&
         "expected that function boundary bufferization is activated");
  FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state);

  // A list of non-circular functions in the order in which they are analyzed
  // and bufferized.
  SmallVector<func::FuncOp> orderedFuncOps;
  // A list of all other functions. I.e., functions that call each other
  // recursively. For these, we analyze the function body but not the function
  // boundary.
  SmallVector<func::FuncOp> remainingFuncOps;

  // A mapping of FuncOps to their callers.
  FuncCallerMap callerMap;

  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps,
                                      remainingFuncOps, callerMap,
                                      funcState.symbolTables)))
    return failure();

  // Analyze functions in order. Starting with functions that are not calling
  // any other functions.
  for (func::FuncOp funcOp : orderedFuncOps) {
    if (!state.getOptions().isOpAllowed(funcOp))
      continue;

    // Now analyzing function.
    funcState.startFunctionAnalysis(funcOp);

    // Analyze funcOp.
    if (failed(analyzeOp(funcOp, state, statistics)))
      return failure();

    // Run some extra function analyses.
    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) ||
        failed(funcOpBbArgReadWriteAnalysis(funcOp, state, funcState)))
      return failure();

    // Mark op as fully analyzed.
    funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
  }

  // Analyze all other functions. All function boundary analyses are skipped.
  for (func::FuncOp funcOp : remainingFuncOps) {
    if (!state.getOptions().isOpAllowed(funcOp))
      continue;

    // Analyze funcOp.
    if (failed(analyzeOp(funcOp, state, statistics)))
      return failure();

    // TODO: We currently skip all function argument analyses for functions
    // that call each other circularly. These analyses do not support recursive
    // calls yet. The `BufferizableOpInterface` implementations of `func`
    // dialect ops return conservative results in the absence of analysis
    // information.
  }

  return success();
}

void mlir::bufferization::removeBufferizationAttributesInModule(
    Operation *moduleOp) {
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (func::FuncOp funcOp : block.getOps<func::FuncOp>()) {
        for (BlockArgument bbArg : funcOp.getArguments())
          removeBufferizationAttributes(bbArg);
      }
    }
  }
}

LogicalResult mlir::bufferization::bufferizeModuleOp(
    Operation *moduleOp, const OneShotBufferizationOptions &options,
    BufferizationState &state, BufferizationStatistics *statistics) {
  assert(options.bufferizeFunctionBoundaries &&
         "expected that function boundary bufferization is activated");
  IRRewriter rewriter(moduleOp->getContext());

  // A list of non-circular functions in the order in which they are analyzed
  // and bufferized.
  SmallVector<func::FuncOp> orderedFuncOps;
  // A list of all other functions. I.e., functions that call each other
  // recursively. For these, we analyze the function body but not the function
  // boundary.
  SmallVector<func::FuncOp> remainingFuncOps;

  // A mapping of FuncOps to their callers.
  FuncCallerMap callerMap;

  // Try to bufferize functions in calling order. I.e., first bufferize
  // functions that do not call other functions. This allows us to infer
  // accurate buffer types for function return values. Functions that call
  // each other recursively are bufferized in an unspecified order at the end.
  // We may use unnecessarily "complex" (in terms of layout map) buffer types.
  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps,
                                      remainingFuncOps, callerMap,
                                      state.getSymbolTables())))
    return failure();
  llvm::append_range(orderedFuncOps, remainingFuncOps);

  // 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 (llvm::is_contained(options.noAnalysisFuncFilter, funcOp.getSymName())) {
      // This function was not analyzed and RaW conflicts were not resolved.
      // Buffer copies must be inserted before every write.
      OneShotBufferizationOptions updatedOptions = options;
      updatedOptions.copyBeforeWrite = true;
      if (failed(bufferizeOp(funcOp, updatedOptions, state, statistics)))
        return failure();
    } else {
      if (failed(bufferizeOp(funcOp, options, state, statistics)))
        return failure();
    }

    // Change buffer return types to more precise layout maps.
    if (options.inferFunctionResultLayout)
      foldMemRefCasts(funcOp);
  }

  // Bufferize all other ops.
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (mlir::Operation &op :
           llvm::make_early_inc_range(block.getOperations())) {
        // Functions were already bufferized.
        if (isa<func::FuncOp>(&op) || op.hasTrait<OpTrait::SymbolTable>())
          continue;
        if (failed(bufferizeOp(&op, options, state, statistics)))
          return failure();
      }
    }
  }

  // Post-pass cleanup of function argument attributes.
  removeBufferizationAttributesInModule(moduleOp);

  return success();
}

LogicalResult mlir::bufferization::runOneShotModuleBufferize(
    Operation *moduleOp, const OneShotBufferizationOptions &options,
    BufferizationState &state, BufferizationStatistics *statistics) {
  assert(options.bufferizeFunctionBoundaries &&
         "expected that function boundary bufferization is activated");
  assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&
         "invalid combination of bufferization flags");
  if (!options.copyBeforeWrite) {
    if (options.noAnalysisFuncFilter.empty()) {
      if (failed(insertTensorCopies(moduleOp, options, state, statistics)))
        return failure();
    } else {
      // FuncOps whose names are specified in options.noAnalysisFuncFilter will
      // not be analyzed. Ops in these FuncOps will not be analyzed as well.
      OpFilter::Entry::FilterFn analysisFilterFn = [=](Operation *op) {
        auto func = dyn_cast<func::FuncOp>(op);
        if (!func)
          func = op->getParentOfType<func::FuncOp>();
        if (func)
          return llvm::is_contained(options.noAnalysisFuncFilter,
                                    func.getSymName());
        return false;
      };
      OneShotBufferizationOptions updatedOptions(options);
      updatedOptions.opFilter.denyOperation(analysisFilterFn);
      if (failed(
              insertTensorCopies(moduleOp, updatedOptions, state, statistics)))
        return failure();
    }
  }
  if (options.testAnalysisOnly)
    return success();
  if (failed(bufferizeModuleOp(moduleOp, options, state, statistics)))
    return failure();
  return success();
}