84 lines · cpp
1//===- TestOneShotModuleBufferzation.cpp - Bufferization Test -----*- c++2//-*-===//3//4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5// See https://llvm.org/LICENSE.txt for license information.6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7//8//===----------------------------------------------------------------------===//9 10#include "mlir/Dialect/Bufferization/IR/Bufferization.h"11#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"12#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"13#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Pass/Pass.h"16 17#include "TestAttributes.h" // TestTensorEncodingAttr, TestMemRefLayoutAttr18#include "TestDialect.h"19 20using namespace mlir;21 22namespace {23MemRefLayoutAttrInterface24getMemRefLayoutForTensorEncoding(RankedTensorType tensorType) {25 if (auto encoding = dyn_cast_if_present<test::TestTensorEncodingAttr>(26 tensorType.getEncoding())) {27 return cast<MemRefLayoutAttrInterface>(test::TestMemRefLayoutAttr::get(28 tensorType.getContext(), encoding.getDummy()));29 }30 return {};31}32 33struct TestOneShotModuleBufferizePass34 : public PassWrapper<TestOneShotModuleBufferizePass, OperationPass<>> {35 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOneShotModuleBufferizePass)36 37 TestOneShotModuleBufferizePass() = default;38 TestOneShotModuleBufferizePass(const TestOneShotModuleBufferizePass &pass)39 : PassWrapper(pass) {}40 41 void getDependentDialects(DialectRegistry ®istry) const override {42 registry.insert<test::TestDialect>();43 registry.insert<bufferization::BufferizationDialect>();44 }45 StringRef getArgument() const final {46 return "test-one-shot-module-bufferize";47 }48 StringRef getDescription() const final {49 return "Pass to test One Shot Module Bufferization";50 }51 52 void runOnOperation() override {53 54 llvm::errs() << "Running TestOneShotModuleBufferize on: "55 << getOperation()->getName() << "\n";56 bufferization::OneShotBufferizationOptions opt;57 58 opt.bufferizeFunctionBoundaries = true;59 opt.functionArgTypeConverterFn =60 [&](bufferization::TensorLikeType tensor, Attribute memSpace,61 func::FuncOp, const bufferization::BufferizationOptions &) {62 assert(isa<RankedTensorType>(tensor) && "tests only builtin tensors");63 auto tensorType = cast<RankedTensorType>(tensor);64 auto layout = getMemRefLayoutForTensorEncoding(tensorType);65 return cast<bufferization::BufferLikeType>(66 MemRefType::get(tensorType.getShape(),67 tensorType.getElementType(), layout, memSpace));68 };69 70 bufferization::BufferizationState bufferizationState;71 72 if (failed(bufferization::runOneShotModuleBufferize(getOperation(), opt,73 bufferizationState)))74 signalPassFailure();75 }76};77} // namespace78 79namespace mlir::test {80void registerTestOneShotModuleBufferizePass() {81 PassRegistration<TestOneShotModuleBufferizePass>();82}83} // namespace mlir::test84