80 lines · cpp
1//===- TestTensorCopyInsertion.cpp - Bufferization Analysis -----*- c++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Bufferization/IR/Bufferization.h"10#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"11#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"12#include "mlir/Pass/Pass.h"13 14using namespace mlir;15 16namespace {17/// This pass runs One-Shot Analysis and inserts copies for all OpOperands that18/// were decided to bufferize out-of-place. After running this pass, a19/// bufferization can write to buffers directly (without making copies) and no20/// longer has to care about potential read-after-write conflicts.21///22/// Note: By default, all newly inserted tensor copies/allocs (i.e., newly23/// created `bufferization.alloc_tensor` ops) that do not escape block are24/// annotated with `escape = false`. If `create-allocs` is unset, all newly25/// inserted tensor copies/allocs are annotated with `escape = true`. In that26/// case, they are not getting deallocated when bufferizing the IR.27struct TestTensorCopyInsertionPass28 : public PassWrapper<TestTensorCopyInsertionPass, OperationPass<ModuleOp>> {29 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTensorCopyInsertionPass)30 31 TestTensorCopyInsertionPass() = default;32 TestTensorCopyInsertionPass(const TestTensorCopyInsertionPass &pass)33 : PassWrapper(pass) {}34 35 void getDependentDialects(DialectRegistry ®istry) const override {36 registry.insert<bufferization::BufferizationDialect>();37 }38 StringRef getArgument() const final { return "test-tensor-copy-insertion"; }39 StringRef getDescription() const final {40 return "Module pass to test Tensor Copy Insertion";41 }42 43 void runOnOperation() override {44 bufferization::OneShotBufferizationOptions options;45 options.allowReturnAllocsFromLoops = allowReturnAllocsFromLoops;46 options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;47 if (mustInferMemorySpace) {48 options.defaultMemorySpaceFn =49 [](TensorType t) -> std::optional<Attribute> { return std::nullopt; };50 }51 52 bufferization::BufferizationState bufferizationState;53 54 if (failed(bufferization::insertTensorCopies(getOperation(), options,55 bufferizationState)))56 signalPassFailure();57 }58 59 Option<bool> allowReturnAllocsFromLoops{60 *this, "allow-return-allocs-from-loops",61 llvm::cl::desc("Allows returning/yielding new allocations from a loop."),62 llvm::cl::init(false)};63 Option<bool> bufferizeFunctionBoundaries{64 *this, "bufferize-function-boundaries",65 llvm::cl::desc("Bufferize function boundaries."), llvm::cl::init(false)};66 Option<bool> mustInferMemorySpace{67 *this, "must-infer-memory-space",68 llvm::cl::desc(69 "The memory space of an memref types must always be inferred. If "70 "unset, a default memory space of 0 is used otherwise."),71 llvm::cl::init(false)};72};73} // namespace74 75namespace mlir::test {76void registerTestTensorCopyInsertionPass() {77 PassRegistration<TestTensorCopyInsertionPass>();78}79} // namespace mlir::test80