100 lines · cpp
1//===- TestSlicing.cpp - Testing slice functionality ----------------------===//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// This file implements a simple testing pass for slicing.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Analysis/SliceAnalysis.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/Linalg/IR/Linalg.h"16#include "mlir/IR/BuiltinOps.h"17#include "mlir/IR/IRMapping.h"18#include "mlir/IR/PatternMatch.h"19#include "mlir/Pass/Pass.h"20#include "mlir/Support/LLVM.h"21 22using namespace mlir;23 24/// Create a function with the same signature as the parent function of `op`25/// with name being the function name and a `suffix`.26static LogicalResult createBackwardSliceFunction(Operation *op,27 StringRef suffix,28 bool omitBlockArguments) {29 func::FuncOp parentFuncOp = op->getParentOfType<func::FuncOp>();30 OpBuilder builder(parentFuncOp);31 Location loc = op->getLoc();32 std::string clonedFuncOpName = parentFuncOp.getName().str() + suffix.str();33 func::FuncOp clonedFuncOp = func::FuncOp::create(34 builder, loc, clonedFuncOpName, parentFuncOp.getFunctionType());35 IRMapping mapper;36 builder.setInsertionPointToEnd(clonedFuncOp.addEntryBlock());37 for (const auto &arg : enumerate(parentFuncOp.getArguments()))38 mapper.map(arg.value(), clonedFuncOp.getArgument(arg.index()));39 SetVector<Operation *> slice;40 BackwardSliceOptions options;41 options.omitBlockArguments = omitBlockArguments;42 // TODO: Make this default.43 options.omitUsesFromAbove = false;44 LogicalResult result = getBackwardSlice(op, &slice, options);45 assert(result.succeeded() && "expected a backward slice");46 (void)result;47 for (Operation *slicedOp : slice)48 builder.clone(*slicedOp, mapper);49 func::ReturnOp::create(builder, loc);50 return success();51}52 53namespace {54/// Pass to test slice generated from slice analysis.55struct SliceAnalysisTestPass56 : public PassWrapper<SliceAnalysisTestPass, OperationPass<ModuleOp>> {57 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SliceAnalysisTestPass)58 59 StringRef getArgument() const final { return "slice-analysis-test"; }60 StringRef getDescription() const final {61 return "Test Slice analysis functionality.";62 }63 64 Option<bool> omitBlockArguments{65 *this, "omit-block-arguments",66 llvm::cl::desc("Test Slice analysis with multiple blocks but slice "67 "omiting block arguments"),68 llvm::cl::init(true)};69 70 void runOnOperation() override;71 SliceAnalysisTestPass() = default;72 SliceAnalysisTestPass(const SliceAnalysisTestPass &) {}73};74} // namespace75 76void SliceAnalysisTestPass::runOnOperation() {77 ModuleOp module = getOperation();78 auto funcOps = module.getOps<func::FuncOp>();79 unsigned opNum = 0;80 for (auto funcOp : funcOps) {81 // TODO: For now this is just looking for Linalg ops. It can be generalized82 // to look for other ops using flags.83 funcOp.walk([&](Operation *op) {84 if (!isa<linalg::LinalgOp>(op))85 return WalkResult::advance();86 std::string append =87 std::string("__backward_slice__") + std::to_string(opNum);88 (void)createBackwardSliceFunction(op, append, omitBlockArguments);89 opNum++;90 return WalkResult::advance();91 });92 }93}94 95namespace mlir {96void registerSliceAnalysisTestPass() {97 PassRegistration<SliceAnalysisTestPass>();98}99} // namespace mlir100