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 @@ -90,7 +90,25 @@ buffer. Furthermore, it returns a handle to the result of the `to_tensor` op. - Only `tensor.pad` targets are supported at the moment. + Only bufferizable ops are that bufferize to a memory write or have an + aliasing OpOperand (and do not themselves bufferize to an allocation) are + supported. They are bufferized using their BufferizableOpInterface + implementation. E.g.: + + ``` + %0 = tensor.insert %f into %dest[%pos] : tensor<10xf32> + ``` + + Is bufferized to: + + ``` + %alloc = memref.alloc() : memref<10xf32> + memref.tensor_store %dest, %alloc : memref<10xf32> + memref.store %f, %alloc[%pos] : memref<10xf32> + %0 = bufferization.to_tensor %alloc restrict writable : memref<10xf32> + ``` + + Selected ops that bufferize to an allocation are also supported: - `tensor.pad` is lowered to an allocation, followed by a `linalg.fill` and and a buffer copy (all on memrefs). 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 @@ -65,11 +65,15 @@ Attribute memorySpace = {}); /// Bufferize the given op with tensor semantics and materialize the result in -/// a newly allocated buffer. E.g.: +/// a newly allocated buffer. /// -/// Only tensor.pad is supported at the moment. +/// Only bufferizable ops that bufferize to a memory write or have an +/// aliasing OpOperand (and do not themselves bufferize to an allocation) are +/// supported. They are bufferized using their BufferizableOpInterface +/// implementation. /// -/// This function returns the newly allocated buffer. +/// Selected ops that bufferize to an allocation are also supported: +/// - tensor.pad Value bufferizeToAllocation(RewriterBase &rewriter, Operation *op, Attribute memorySpace = {}); diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -942,7 +942,7 @@ // attribute. Such tensors may alias any other tensor, which is currently // not handled in the analysis. if (auto toTensorOp = dyn_cast(op.getOperation())) { - if (!toTensorOp.getRestrict()) { + if (!toTensorOp.getRestrict() && !toTensorOp->getUses().empty()) { op->emitError("to_tensor ops without `restrict` are not supported by " "One-Shot Analysis"); return WalkResult::interrupt(); 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 @@ -330,12 +330,91 @@ Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op, Attribute memorySpace) { + using namespace bufferization; + // Call specialized overload for certain ops. if (auto padOp = dyn_cast(op)) return bufferizeToAllocation(rewriter, padOp, memorySpace); - // TODO: Support other ops. - return nullptr; + // Only bufferizable ops are supported. + auto bufferizableOp = dyn_cast(op); + if (!bufferizableOp) + return nullptr; + BufferizationOptions options; + AnalysisState state(options); + + // Gather tensor results. + SmallVector tensorResults; + for (OpResult result : op->getResults()) { + if (!result.getType().isa()) + continue; + // Unranked tensors are not supported + if (!isa(result.getType())) + return nullptr; + // Ops that bufferize to an allocation are not supported. + if (bufferizableOp.bufferizesToAllocation(result)) + return nullptr; + tensorResults.push_back(result); + } + + // Gather all operands that should bufferize to a new allocation. I.e., + // bufferize out-of-place. + SmallVector outOfPlaceOperands, resultUses; + auto addOutOfPlaceOperand = [&](OpOperand *operand) { + if (llvm::find(outOfPlaceOperands, operand) == outOfPlaceOperands.end()) + outOfPlaceOperands.push_back(operand); + }; + for (OpResult result : tensorResults) { + AliasingOpOperandList aliasingOperands = + state.getAliasingOpOperands(result); + for (const AliasingOpOperand &operand : aliasingOperands) { + addOutOfPlaceOperand(operand.opOperand); + for (OpOperand &resultUse : result.getUses()) + resultUses.push_back(&resultUse); + } + } + for (OpOperand &operand : op->getOpOperands()) { + if (!state.bufferizesToMemoryWrite(operand)) + continue; + if (!isa(operand.get().getType())) + return nullptr; + addOutOfPlaceOperand(&operand); + } + // TODO: Support multiple buffers. + if (outOfPlaceOperands.size() != 1) + return nullptr; + + // Allocate buffers. + OpBuilder::InsertionGuard g(rewriter); + rewriter.setInsertionPoint(op); + SmallVector allocs; + for (OpOperand *operand : outOfPlaceOperands) { + Value alloc = createAllocationForTensor(rewriter, op->getLoc(), + operand->get(), memorySpace); + allocs.push_back(alloc); + // Initialize buffer with a copy of the operand data. + // TODO: Do not copy uninitialized tensors such as tensor.empty. + rewriter.create(op->getLoc(), operand->get(), alloc); + rewriter.updateRootInPlace(op, [&]() { + operand->set(rewriter.create(op->getLoc(), alloc)); + }); + } + + // Bufferize the op. + if (failed(bufferizableOp.bufferize(rewriter, options))) + return nullptr; + + // Set "restrict" attribute, indicating that no other tensor aliases with + // this tensor. That is because we just allocated a new buffer for the tensor. + for (OpOperand *resultUse : resultUses) { + auto toTensorOp = resultUse->get().getDefiningOp(); + assert(toTensorOp && "expected to_tensor op"); + rewriter.updateRootInPlace(toTensorOp, [&]() { + toTensorOp.setRestrict(true); + toTensorOp.setWritable(true); + }); + } + return allocs.front(); } namespace { 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 @@ -1,4 +1,4 @@ -// RUN: mlir-opt -split-input-file \ +// RUN: mlir-opt -split-input-file -verify-diagnostics \ // RUN: -test-transform-dialect-interpreter -canonicalize \ // RUN: -allow-unregistered-dialect -split-input-file %s | FileCheck %s @@ -62,3 +62,41 @@ %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op } +// ----- + +// CHECK-LABEL: func @tensor_insert( +// CHECK-SAME: %[[t:.*]]: tensor +// CHECK: %[[m:.*]] = bufferization.to_memref %[[t]] +// CHECK: %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref +// CHECK: memref.copy %[[m]], %[[alloc]] +// CHECK: memref.store %{{.*}}, %[[alloc]] +// CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable +// CHECK: memref.dealloc %[[alloc]] +// CHECK: return %[[r]] +func.func @tensor_insert(%t: tensor, %idx: index, %v: index) -> tensor { + %r = tensor.insert %v into %t[%idx, %idx] : tensor + return %r : tensor +} + +transform.sequence failures(propagate) { +^bb1(%arg1: !transform.any_op): + %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op + %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !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 +} + +// ----- + +func.func @tensor_extract(%t: tensor, %idx: index) -> index { + // expected-note @below{{target payload op}} + %r = tensor.extract %t[%idx, %idx] : tensor + return %r : index +} + +transform.sequence failures(propagate) { +^bb1(%arg1: !transform.any_op): + %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op + // expected-error @below{{failed to bufferize operation}} + %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op +}