[mlir][linalg][bufferize] Bufferize using PostOrder traversal

[springerm]

Summary

This is required for bufferization of scf::IfOp, which is added in a subsequent commit.

Some ops (scf::ForOp, TiledLoopOp) require PreOrder traversal to make sure that bbArgs are mapped before bufferizing the loop body.

[springerm]

no change

[springerm]

rebase

[springerm]

fix bug with nested for loop

[nicolasvasilache]

Bufferization crashes in various ways on the following example.
It may be that the IR is not in a good form to start with, still we should not crash at callOp->erase (and if I make that removal conditional I get assert.h assertion failed at third_party/llvm/llvm-project/llvm/include/llvm/ADT/EquivalenceClasses.h:172 in const ElemTy &llvm::EquivalenceClasses<mlir::linalg::BufferizationAliasInfo::ValueWrapper>::getLeaderValue(const ElemTy &) const [ElemTy = mlir::linalg::BufferizationAliasInfo::ValueWrapper]: MI != member_end() && "Value is not in the set!").

module  {
  func @reduction_2d_on_tensors(%arg0: tensor<4096x1024xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false}, %arg1: tensor<4096xf32> {linalg.buffer_layout = affine_map<(d0) -> (d0)>, linalg.inplaceable = true}) -> tensor<4096xf32> attributes {passthrough = ["noinline", ["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]} {
    %cst = arith.constant 0.000000e+00 : f32
    %c16 = arith.constant 16 : index
    %c4096 = arith.constant 4096 : index
    %c1024 = arith.constant 1024 : index
    %c0 = arith.constant 0 : index
    %0 = linalg.fill(%cst, %arg1) : f32, tensor<4096xf32> -> tensor<4096xf32> 
    %1 = linalg.init_tensor [16] : tensor<16xf32>
    %2 = linalg.fill(%cst, %1) : f32, tensor<16xf32> -> tensor<16xf32> 
    %3 = vector.transfer_read %2[%c0], %cst {in_bounds = [true]} : tensor<16xf32>, vector<16xf32>
    %4 = scf.for %arg2 = %c0 to %c4096 step %c16 iter_args(%arg3 = %0) -> (tensor<4096xf32>) {
      %5 = vector.transfer_read %0[%arg2], %cst {in_bounds = [true]} : tensor<4096xf32>, vector<16xf32>
      %6 = scf.for %arg4 = %c0 to %c1024 step %c16 iter_args(%arg5 = %arg3) -> (tensor<4096xf32>) {
        %7 = vector.transfer_read %arg0[%arg2, %arg4], %cst {in_bounds = [true, true]} : tensor<4096x1024xf32>, vector<16x16xf32>
        %8 = vector.multi_reduction #vector.kind<add>, %7 [1] : vector<16x16xf32> to vector<16xf32>
        %9 = arith.addf %8, %3 : vector<16xf32>
        %10 = arith.addf %5, %9 : vector<16xf32>
        %11 = vector.transfer_write %10, %arg5[%arg2] {in_bounds = [true]} : vector<16xf32>, tensor<4096xf32>
        scf.yield %11 : tensor<4096xf32>
      }
      scf.yield %6 : tensor<4096xf32>
    }
    return %4 : tensor<4096xf32>
  }
  func public @main(%arg0: tensor<4096x1024xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false}, %arg1: tensor<4096xf32> {linalg.buffer_layout = affine_map<(d0) -> (d0)>, linalg.inplaceable = true}, %arg2: index) -> tensor<4096xf32> attributes {llvm.emit_c_interface} {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %0 = scf.for %arg3 = %c0 to %arg2 step %c1 iter_args(%arg4 = %arg1) -> (tensor<4096xf32>) {
      %1 = call @reduction_2d_on_tensors(%arg0, %arg4) : (tensor<4096x1024xf32>, tensor<4096xf32>) -> tensor<4096xf32>
      scf.yield %1 : tensor<4096xf32>
    }
    return %0 : tensor<4096xf32>
  }
}

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

for (OpOperand &opOperand : callOp->getOpOperands()) {
     }
     newOperands.push_back(buffer);
   }
 
   // 4. Create the new CallOp.
   Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
                                           resultTypes, newOperands);
   newCallOp->setAttrs(callOp->getAttrs());
+  callOp->erase();
   return success();
 }
 
 /// tensor::CastOp bufferizes to memref::CastOp.
 static LogicalResult bufferize(OpBuilder &b, tensor::CastOp castOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   // Take a guard before anything else.

 static LogicalResult bufferizeFuncOpInternals(
     FuncOp funcOp, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
     DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes,
     GlobalCreator &globalCreator) {
 
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin BufferizeFuncOpInternals:\n" << funcOp << '\n');
   OpBuilder b(funcOp->getContext());
-  /// Start by bufferizing `funcOp` arguments.
+
+  // Start by bufferizing `funcOp` arguments.
   if (failed(bufferize(b, funcOp, bvm, aliasInfo)))
     return failure();
 
-  // Walk in PreOrder to ensure ops with regions are handled before their body.
-  // Since walk has to be PreOrder, we need to erase ops that require it
-  // separately: this is the case for CallOp
-  SmallVector<Operation *> toErase;
-  if (funcOp
-          .walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
-            if (failed(bufferizeOp(op, bvm, aliasInfo, &bufferizedFunctionTypes,
-                                   &globalCreator)))
-              return failure();
-            // Register post-walk erasure, if necessary.
-            if (isa<CallOpInterface>(op))
-              if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
-                  llvm::any_of(op->getResultTypes(), isaTensor))
-                toErase.push_back(op);
-            return success();
-          })
-          .wasInterrupted())
+  // Bufferize the function body. `bufferizedOps` keeps track ops that were
+  // already bufferized with pre-order traversal.
+  DenseSet<Operation *> bufferizedOps;
+  auto walkFunc = [&](Operation *op) -> WalkResult {
+    // Collect ops that need to be bufferized before `op`.
+    SmallVector<Operation *> preorderBufferize;
+    Operation *parentOp = op->getParentOp();
+    // scf::ForOp and TiledLoopOp must be bufferized before their blocks
+    // ("pre-order") because BBargs must be mapped when bufferizing children.
+    while (isa_and_nonnull<scf::ForOp, TiledLoopOp>(parentOp)) {
+      if (bufferizedOps.contains(parentOp))
+        break;
+      bufferizedOps.insert(parentOp);
+      preorderBufferize.push_back(parentOp);
+      parentOp = parentOp->getParentOp();
+    }
+
+    for (Operation *op : llvm::reverse(preorderBufferize))
+
+    if (!bufferizedOps.contains(op) &&
+        failed(bufferizeOp(op, bvm, aliasInfo, &bufferizedFunctionTypes,
+                           &globalCreator)))
+      return failure();
+  };
+  if (funcOp.walk(walkFunc).wasInterrupted())
     return failure();
 
   LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
 
-  for (Operation *op : toErase)
-    op->erase();
-
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for modules.
 //===----------------------------------------------------------------------===//
 
 /// Return the op with Allocate MemoryEffect if `v` is equivalent to such an

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

func @extract_slice_fun(%A : tensor<?xf32> {linalg.inplaceable = true})
   // expected-error @+1 {{buffer result #0 not produced by an alloc}}
   return %r0: tensor<4xf32>
 }
 
 // -----
 
 func @scf_yield(%b : i1, %A : tensor<4xf32>, %B : tensor<4xf32>) -> tensor<4xf32>
 {
-  // expected-error @+1 {{unsupported op with tensors}}
   %r = scf.if %b -> (tensor<4xf32>) {
+    // expected-error @+1 {{expected scf::ForOp parent for scf::YieldOp}}
     scf.yield %A : tensor<4xf32>
   } else {
     scf.yield %B : tensor<4xf32>
   }
   return %r: tensor<4xf32>
 }
 
 // -----
func @mini_test_case1() -> tensor<10x20xf32> {
   %t = linalg.init_tensor [10, 20] : tensor<10x20xf32>
   %r = linalg.fill(%f0, %t) : f32, tensor<10x20xf32> -> tensor<10x20xf32>
   return %r : tensor<10x20xf32>
 }
 
 // -----
 
 func @main() -> tensor<4xi32> {
-  // expected-error @+1 {{unsupported op with tensors}}
+  // expected-error @+1 {{expected result-less scf.execute_region containing op}}
   %r = scf.execute_region -> tensor<4xi32> {
     %A = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
     scf.yield %A: tensor<4xi32>
   }
   return %r: tensor<4xi32>
 }
 
 // -----

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @scf_for_yield_only(%A : tensor<?xf32>,
 
   //     CHECK:   memref.dealloc %[[ALLOC_FOR_A]] : memref<?xf32>
   //     CHECK:   return %[[ALLOC_FOR_A]] : memref<?xf32>
   return %r0, %r1: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
+// Ensure that the function bufferizes without error. This tests pre-order
+// traversal of scf.for loops during bufferization. No need to check the IR,
+// just want to make sure that it does not crash.
+
+// CHECK-LABEL: func @nested_scf_for
+func @nested_scf_for(%A : tensor<?xf32> {linalg.inplaceable = true},
+                     %v : vector<5xf32>) -> tensor<?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c10 = constant 10 : index
+  %r1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%B = %A) -> tensor<?xf32> {
+    %r2 = scf.for %j = %c0 to %c10 step %c1 iter_args(%C = %B) -> tensor<?xf32> {
+      %w = vector.transfer_write %v, %C[%c0] : vector<5xf32>, tensor<?xf32>
+      scf.yield %w : tensor<?xf32>
+    }
+    scf.yield %r2 : tensor<?xf32>
+  }
+  return %r1 : tensor<?xf32>
+}
+
+// -----
+
 // CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL: func @scf_for_with_tensor.insert_slice
 //  CHECK-SAME:   %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 //  CHECK-SAME:   %[[B:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>
 //  CHECK-SAME:   %[[C:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]>
 func @scf_for_with_tensor.insert_slice(
    %A : tensor<?xf32>,