diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -128,6 +128,11 @@ a fully dynamic layout is assumed for best compatibility. Users should use "memref.tensor_store" when possible. + "memref.alloc" is used for new buffer allocations. The buffer is deallocated + at the end of the block. Custom allocation ops can be specified via + `alloc_op`. Currently supported are "memref.alloc" and "memref.alloca". In + case of a "memref.alloca", the buffer is not deallocated. + #### Return modes This operation consumes the `target` handle and produces the @@ -137,7 +142,9 @@ let arguments = (ins TransformHandleTypeInterface:$target, OptionalAttr:$memory_space, DefaultValuedAttr: - $memcpy_op); + $memcpy_op, + DefaultValuedAttr: + $alloc_op); let hasVerifier = 1; let results = (outs Transform_AnyValue:$allocated_buffer, Transform_AnyOpType:$new_ops); diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -47,8 +47,10 @@ //===----------------------------------------------------------------------===// struct BufferizeToAllocationOptions { - enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 }; + enum class AllocOp { MemrefAlloc = 0, MemrefAlloca = 1 }; + AllocOp allocOp = AllocOp::MemrefAlloc; + enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 }; MemcpyOp memcpyOp = MemcpyOp::MemrefTensorStore; }; diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -248,6 +248,15 @@ } else { llvm_unreachable("invalid memcpy op"); } + if (getAllocOp() == "memref.alloc") { + options.allocOp = + linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloc; + } else if (getAllocOp() == "memref.alloca") { + options.allocOp = + linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloca; + } else { + llvm_unreachable("invalid alloc op"); + } // Bufferize ops. Attribute memorySpace = @@ -283,6 +292,8 @@ if (getMemcpyOp() != "memref.tensor_store" && getMemcpyOp() != "memref.copy" && getMemcpyOp() != "linalg.copy") return emitOpError() << "unsupported memcpy op"; + if (getAllocOp() != "memref.alloc" && getAllocOp() != "memref.alloca") + return emitOpError() << "unsupported alloc op"; return success(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp @@ -185,9 +185,10 @@ return dynSizes; } -static Value createAllocationForTensor(RewriterBase &rewriter, Location loc, - Value value, - Attribute memorySpace = {}) { +static Value +createAllocationForTensor(RewriterBase &rewriter, Location loc, Value value, + const linalg::BufferizeToAllocationOptions &options, + Attribute memorySpace = {}) { OpBuilder::InsertionGuard g(rewriter); auto tensorType = cast(value.getType()); @@ -196,11 +197,19 @@ cast(bufferization::getMemRefTypeWithStaticIdentityLayout( tensorType, memorySpace)); SmallVector dynamicSizes = reifyOrComputeDynamicSizes(rewriter, value); - Value alloc = rewriter.create(loc, memrefType, dynamicSizes); - // Place deallocation at the end of the block. - rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator()); - rewriter.create(loc, alloc); + Value alloc; + if (options.allocOp == + linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloc) { + alloc = rewriter.create(loc, memrefType, dynamicSizes); + // Place deallocation at the end of the block. + rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator()); + rewriter.create(loc, alloc); + } else if (options.allocOp == + linalg::BufferizeToAllocationOptions::AllocOp::MemrefAlloca) { + alloc = rewriter.create(loc, memrefType, dynamicSizes); + // No dealloc is needed. + } return alloc; } @@ -213,8 +222,8 @@ Location loc = padOp.getLoc(); // Create buffer allocation. - Value alloc = - createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace); + Value alloc = createAllocationForTensor(rewriter, loc, padOp.getResult(), + options, memorySpace); rewriter.setInsertionPoint(padOp); if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) { @@ -491,8 +500,8 @@ rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op); SmallVector allocs; for (OpOperand *operand : outOfPlaceOperands) { - Value alloc = createAllocationForTensor(rewriter, op->getLoc(), - operand->get(), memorySpace); + Value alloc = createAllocationForTensor( + rewriter, op->getLoc(), operand->get(), options, memorySpace); allocs.push_back(alloc); if (!state.findDefinitions(operand->get()).empty()) { // Initialize buffer with a copy of the operand data. Not needed if the diff --git a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir --- a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir +++ b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir @@ -50,6 +50,7 @@ // CHECK-LABEL: func @tensor_pad_constant_with_custom_copy( // CHECK-NOT: memref.tensor_store // CHECK-NOT: memref.copy +// CHECK: memref.alloca // CHECK: linalg.copy func.func @tensor_pad_constant_with_custom_copy( %t: tensor, %l2: index, %h1: index, %h2: index) @@ -66,7 +67,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !transform.any_op): %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op - %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, memcpy_op = "linalg.copy"}: !transform.any_op + %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, alloc_op = "memref.alloca", memcpy_op = "linalg.copy"}: !transform.any_op // Ensure that one linalg.fill was generated. %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op @@ -78,6 +79,11 @@ // expected-remark @below{{1}} test_print_number_of_associated_payload_ir_ops %linalg_copy : !transform.any_op + // Ensure that one memref.alloca was generated. + %alloca = transform.select "memref.alloca" in %new : (!transform.any_op) -> !transform.any_op + // expected-remark @below{{1}} + test_print_number_of_associated_payload_ir_ops %alloca : !transform.any_op + // Make sure that One-Shot Bufferize can bufferize the rest. %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op }