[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.

 // Bufferization entry-point for functions.
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::linalg::bufferizeOp(
     Operation *op, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes,
     GlobalCreator *globalCreator) {
   OpBuilder b(op->getContext());
+
+  // Return the first child op or nullptr for ops without children.
+  // TODO: Is there a better way?
+  auto getFirstChildOp = [](Operation *op) {
+    Operation *firstChildOp = nullptr;
+    op->walk([&](Operation *childOp) -> WalkResult {
+      firstChildOp = childOp;
+      return WalkResult::interrupt();
+    });
+    return firstChildOp;
+  };
+
+  // Some ops require a pre-order traversal. Bufferize such ops right before
+  // bufferizing their first child op.
+  if (getFirstChildOp(op->getParentOp()) == op) {
+    // op is the first child inside op's parent.
+    LogicalResult preOrder =
+        TypeSwitch<Operation *, LogicalResult>(op->getParentOp())
+            // bbArgs must be mapped in bvm before bufferizing the loop body.
+            .Case<scf::ForOp, TiledLoopOp>([&](auto op) {
+              LDBG("Begin bufferize:\n" << op << '\n');
+              return bufferize(b, op, bvm, aliasInfo);
+            })
+            .Default([](Operation *op) { return success(); });
+    if (failed(preOrder))
+      return failure();
+  }
+
+  // Handling of ops with regular post-order traversal.
   return TypeSwitch<Operation *, LogicalResult>(op)
       // Skip BufferCast and TensorLoad ops.
       .Case<memref::BufferCastOp, memref::TensorLoadOp>(
           [&](auto) { return success(); })
-      .Case<tensor::CastOp, tensor::DimOp, ExtractSliceOp, scf::ForOp,
-            InitTensorOp, InsertSliceOp, tensor::ExtractOp, LinalgOp, ReturnOp,
-            TiledLoopOp, VectorTransferOpInterface, linalg::YieldOp,
-            scf::YieldOp>([&](auto op) {
+      // Handle pre-order traversal ops that have no children.
+      .Case<scf::ForOp, TiledLoopOp>([&](auto op) {
+        if (!getFirstChildOp(op))
+          return bufferize(b, op, bvm, aliasInfo);
+        return success();
+      })
+      .Case<tensor::CastOp, tensor::DimOp, ExtractSliceOp, InitTensorOp,
+            InsertSliceOp, tensor::ExtractOp, LinalgOp, ReturnOp,
+            VectorTransferOpInterface, linalg::YieldOp, scf::YieldOp>(
+          [&](auto op) {
             LDBG("Begin bufferize:\n" << op << '\n');
             return bufferize(b, op, bvm, aliasInfo);
           })
       .Case([&](CallOpInterface op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         if (!bufferizedFunctionTypes)
           llvm_unreachable(
               "null bufferizedFunctionTypes when bufferizing CallOpInterface");
         return bufferize(b, op, bvm, aliasInfo, *bufferizedFunctionTypes);
       })
       .Case([&](ConstantOp op) {

 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.
+  auto walkFunc = [&](Operation *op) -> WalkResult {
+  };
+  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})
   %r0 = tensor.extract_slice %A[0][4][1] : tensor<?xf32> to tensor<4xf32>
 
   // 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>) {
-    scf.yield %A : tensor<4xf32>
-  } else {
-    scf.yield %B : tensor<4xf32>
-  }
-  return %r: tensor<4xf32>
-}
-
-// -----
-
 func @unknown_op(%A : tensor<4xf32>) -> tensor<4xf32>
 {
   // expected-error @+1 {{unsupported op with tensors}}
   %r = "marklar"(%A) : (tensor<4xf32>) -> (tensor<4xf32>)
   return %r: tensor<4xf32>
 }
 
 // -----
 
 // expected-error @+1 {{memref return type is unsupported}}
 func @mini_test_case1() -> tensor<10x20xf32> {
   %f0 = constant 0.0 : f32
   %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>
 }
 
 // -----