[Mlir-commits] [mlir] [emitC]Pass in `mlir-opt` to wrap a func in class (PR #141158)

[Paul Kirth]

Valentin Clement =?utf-8?b?KOODkOODrOODsw=?=,Jaddyen <ajaden at google.com>,Jaddyen
 <ajaden at google.com>,Jaddyen <ajaden at google.com>,Jaddyen <ajaden at google.com>,Jaddyen
 <ajaden at google.com>,Jaddyen <ajaden at google.com>,Jaddyen <ajaden at google.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/141158 at github.com>


================
@@ -0,0 +1,150 @@
+//===- ConvertFuncToClass.cpp - Convert functions to classes -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Rewrite.h"
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/EmitC/Transforms/Passes.h"
+#include "mlir/Dialect/EmitC/Transforms/Transforms.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeRange.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/GraphWriter.h"
+#include "llvm/Support/LogicalResult.h"
+#include <string>
+
+namespace mlir {
+namespace emitc {
+
+#define GEN_PASS_DEF_WRAPFUNCINCLASSPASS
+#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
+
+namespace {
+
+struct WrapFuncInClassPass
+    : public impl::WrapFuncInClassPassBase<WrapFuncInClassPass> {
+  using WrapFuncInClassPassBase::WrapFuncInClassPassBase;
+  void runOnOperation() override {
+    Operation *rootOp = getOperation();
+    MLIRContext *context = rootOp->getContext();
+
+    RewritePatternSet patterns(context);
+    populateFuncPatterns(patterns, namedAttribute);
+
+    if (failed(applyPatternsGreedily(rootOp, std::move(patterns))))
+      return signalPassFailure();
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<emitc::EmitCDialect>();
+  }
+};
+
+} // namespace
+
+} // namespace emitc
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::emitc;
+
+class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {
+private:
+  std::string attributeName;
+
+public:
+  WrapFuncInClass(MLIRContext *context, const std::string &attrName)
+      : OpRewritePattern<emitc::FuncOp>(context), attributeName(attrName) {}
+
+  LogicalResult matchAndRewrite(emitc::FuncOp funcOp,
+                                PatternRewriter &rewriter) const override {
+    if (funcOp->getParentOfType<emitc::ClassOp>()) {
+      return failure();
+    }
+    auto className = funcOp.getSymNameAttr().str() + "Class";
+    mlir::emitc::ClassOp newClassOp =
+        rewriter.create<emitc::ClassOp>(funcOp.getLoc(), className);
+
+    SmallVector<std::pair<StringAttr, TypeAttr>> fields;
+    rewriter.createBlock(&newClassOp.getBody());
+    rewriter.setInsertionPointToStart(&newClassOp.getBody().front());
+
+    auto argAttrs = funcOp.getArgAttrs();
+    size_t idx = 0;
+
+    for (const BlockArgument &val : funcOp.getArguments()) {
----------------
ilovepi wrote:

I imagine `enumerate()` may be more natural? https://llvm.org/docs/ProgrammersManual.html#enumerate

https://github.com/llvm/llvm-project/pull/141158