110 lines · cpp
1//===- WrapFuncInClass.cpp - Wrap Emitc Funcs in classes -------------===//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/EmitC/IR/EmitC.h"10#include "mlir/Dialect/EmitC/Transforms/Passes.h"11#include "mlir/Dialect/EmitC/Transforms/Transforms.h"12#include "mlir/IR/Attributes.h"13#include "mlir/IR/Builders.h"14#include "mlir/IR/BuiltinAttributes.h"15#include "mlir/IR/PatternMatch.h"16#include "mlir/Transforms/WalkPatternRewriteDriver.h"17 18using namespace mlir;19using namespace emitc;20 21namespace mlir {22namespace emitc {23#define GEN_PASS_DEF_WRAPFUNCINCLASSPASS24#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"25 26namespace {27struct WrapFuncInClassPass28 : public impl::WrapFuncInClassPassBase<WrapFuncInClassPass> {29 using WrapFuncInClassPassBase::WrapFuncInClassPassBase;30 void runOnOperation() override {31 Operation *rootOp = getOperation();32 33 RewritePatternSet patterns(&getContext());34 populateFuncPatterns(patterns);35 36 walkAndApplyPatterns(rootOp, std::move(patterns));37 }38};39 40} // namespace41} // namespace emitc42} // namespace mlir43 44class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {45public:46 WrapFuncInClass(MLIRContext *context)47 : OpRewritePattern<emitc::FuncOp>(context) {}48 49 LogicalResult matchAndRewrite(emitc::FuncOp funcOp,50 PatternRewriter &rewriter) const override {51 52 auto className = funcOp.getSymNameAttr().str() + "Class";53 ClassOp newClassOp = ClassOp::create(rewriter, funcOp.getLoc(), className);54 55 SmallVector<std::pair<StringAttr, TypeAttr>> fields;56 rewriter.createBlock(&newClassOp.getBody());57 rewriter.setInsertionPointToStart(&newClassOp.getBody().front());58 59 auto argAttrs = funcOp.getArgAttrs();60 for (auto [idx, val] : llvm::enumerate(funcOp.getArguments())) {61 StringAttr fieldName =62 rewriter.getStringAttr("fieldName" + std::to_string(idx));63 64 TypeAttr typeAttr = TypeAttr::get(val.getType());65 fields.push_back({fieldName, typeAttr});66 67 FieldOp fieldop = emitc::FieldOp::create(rewriter, funcOp->getLoc(),68 fieldName, typeAttr, nullptr);69 70 if (argAttrs && idx < argAttrs->size()) {71 fieldop->setDiscardableAttrs(funcOp.getArgAttrDict(idx));72 }73 }74 75 rewriter.setInsertionPointToEnd(&newClassOp.getBody().front());76 FunctionType funcType = funcOp.getFunctionType();77 Location loc = funcOp.getLoc();78 FuncOp newFuncOp =79 emitc::FuncOp::create(rewriter, loc, ("execute"), funcType);80 81 rewriter.createBlock(&newFuncOp.getBody());82 newFuncOp.getBody().takeBody(funcOp.getBody());83 84 rewriter.setInsertionPointToStart(&newFuncOp.getBody().front());85 std::vector<Value> newArguments;86 newArguments.reserve(fields.size());87 for (auto &[fieldName, attr] : fields) {88 GetFieldOp arg =89 emitc::GetFieldOp::create(rewriter, loc, attr.getValue(), fieldName);90 newArguments.push_back(arg);91 }92 93 for (auto [oldArg, newArg] :94 llvm::zip(newFuncOp.getArguments(), newArguments)) {95 rewriter.replaceAllUsesWith(oldArg, newArg);96 }97 98 llvm::BitVector argsToErase(newFuncOp.getNumArguments(), true);99 if (failed(newFuncOp.eraseArguments(argsToErase)))100 newFuncOp->emitOpError("failed to erase all arguments using BitVector");101 102 rewriter.replaceOp(funcOp, newClassOp);103 return success();104 }105};106 107void mlir::emitc::populateFuncPatterns(RewritePatternSet &patterns) {108 patterns.add<WrapFuncInClass>(patterns.getContext());109}110