[Mlir-commits] [mlir] [mlir][EmitC] Create a pass to add a reflection map to a class (PR #205464)

[Andrzej Warzyński]

================
@@ -0,0 +1,240 @@
+//===----------------------------------------------------------------------===//
+//
+// 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/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/Transforms/WalkPatternRewriteDriver.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+using namespace emitc;
+
+namespace mlir {
+namespace emitc {
+#define GEN_PASS_DEF_MLGOADDREFLECTIONMAPPASS
+#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
+
+namespace {
+constexpr const char *mapLibraryHeader = "map";
+constexpr const char *stringLibraryHeader = "string";
+
+struct PatternMatchListener : public RewriterBase::Listener {
+  bool patternApplied = false;
+
+  void notifyOperationInserted(Operation *op,
+                               OpBuilder::InsertPoint previous) override {
+    patternApplied = true;
+  }
+};
+
+IncludeOp addHeader(OpBuilder &builder, ModuleOp module, StringRef headerName) {
+  StringAttr includeAttr = builder.getStringAttr(headerName);
+  return IncludeOp::create(builder, module.getLoc(), includeAttr,
+                           /*is_standard_include=*/builder.getUnitAttr());
+}
+
+class MLGOAddReflectionMapPass
+    : public impl::MLGOAddReflectionMapPassBase<MLGOAddReflectionMapPass> {
+  using MLGOAddReflectionMapPassBase::MLGOAddReflectionMapPassBase;
+  void runOnOperation() override {
+    mlir::ModuleOp moduleOp = getOperation();
+
+    RewritePatternSet patterns(&getContext());
+    populateMLGOAddReflectionMapPatterns(patterns, includedFieldAttrs);
+
+    PatternMatchListener listener;
+    walkAndApplyPatterns(moduleOp, std::move(patterns), &listener);
+
+    // If nothing was matched, no reflection maps were added, removing the need
+    // to add include headers for map and string
+    if (!listener.patternApplied)
+      return;
+
+    // Check if the map and/or string headers are already present
+    bool hasMapHdr = false;
+    bool hasStringHdr = false;
+    for (auto &op : *moduleOp.getBody()) {
+      IncludeOp includeOp = llvm::dyn_cast<IncludeOp>(op);
+      if (!includeOp)
+        continue;
+
+      if (includeOp.getIsStandardInclude()) {
+        auto include = includeOp.getInclude();
+
+        hasMapHdr |= include == mapLibraryHeader;
+        hasStringHdr |= include == stringLibraryHeader;
+      }
+
+      if (hasMapHdr && hasStringHdr)
+        return;
+    }
+
+    mlir::OpBuilder builder(moduleOp.getBody(), moduleOp.getBody()->begin());
+    if (!hasMapHdr)
+      addHeader(builder, moduleOp, mapLibraryHeader);
+
+    if (!hasStringHdr)
+      addHeader(builder, moduleOp, stringLibraryHeader);
+  }
+};
+
+} // namespace
+} // namespace emitc
+} // namespace mlir
+
+/// Rewrites a `emitc::ClassOp` to generate a reflection map of member fields
+/// and a lookup method.
+///
+/// Fields to be mapped are identified via `includedFieldAttrs` attributes
+/// (e.g., `emitc.field_ref`). Fields that do not have a matching attribute
+/// are omitted from the reflection map.
+///
+/// Before:
+/// ```mlir
+/// emitc.class @foo {
+///   emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref =
+///   ["another_feature"]} emitc.func @bar() { return }
+/// }
+/// ```
+///
+/// After:
+/// ```mlir
+/// emitc.class @foo {
+///   emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref =
+///   ["another_feature"]} emitc.field @reflectionMap : !emitc.opaque<"const
+///   std::map<std::string, char*>"> =
+///     #emitc.opaque<"{ { \"another_feature\",
+///     reinterpret_cast<char*>(&fieldName0) } }">
+///   emitc.func @getBufferForName(%name: !emitc.opaque<"std::string">) ->
+///   !emitc.ptr<!emitc.opaque<"char">> {
+///     %map = get_field @reflectionMap : !emitc.opaque<"const
+///     std::map<std::string, char*>"> %ptr = member_call_opaque %map
+///     "at"(%name) : ... return %ptr : !emitc.ptr<!emitc.opaque<"char">>
+///   }
+///   emitc.func @bar() { return }
+/// }
+/// ```
+class MLGOAddReflectionMapClass : public OpRewritePattern<ClassOp> {
+public:
+  MLGOAddReflectionMapClass(MLIRContext *context,
+                            llvm::ArrayRef<std::string> includedFieldAttrs)
+      : OpRewritePattern<ClassOp>(context),
+        includedFieldAttrs(includedFieldAttrs) {}
+
+  LogicalResult matchAndRewrite(ClassOp classOp,
+                                PatternRewriter &rewriter) const override {
+    MLIRContext *context = rewriter.getContext();
+
+    emitc::OpaqueType mapType = mlir::emitc::OpaqueType::get(
+        context, "const std::map<std::string, char*>");
+
+    // Collect the names of all FieldOps that have one of the attributes in
+    // includedFieldAttrs to use the first element of the array attribute
+    // as the reflection map key
+    std::vector<std::pair<StringRef, StringRef>> fieldNames;
+    classOp.walk([&](FieldOp fieldOp) {
+      for (const auto &attr : includedFieldAttrs) {
+        auto arrayAttr = dyn_cast_if_present<ArrayAttr>(fieldOp->getAttr(attr));
+
+        if (!arrayAttr)
+          continue;
+
+        if (!arrayAttr.empty() && isa<StringAttr>(arrayAttr[0])) {
+          fieldNames.emplace_back(cast<StringAttr>(arrayAttr[0]).getValue(),
+                                  fieldOp.getName());
+          return;
+        }
+      }
+    });
+
+    if (fieldNames.empty())
+      return failure();
+
+    // Create reflection map contents
+    std::string reflectionMapContents;
+    reflectionMapContents += "{ ";
+    bool first = true;
+    for (const auto &[name, value] : fieldNames) {
+      if (!first)
+        reflectionMapContents += ", ";
+
+      first = false;
+      reflectionMapContents += llvm::formatv(
+          "{ \"{0}\", reinterpret_cast<char*>(&{1}) }", name, value);
+    }
+    reflectionMapContents += " }";
+
+    // Set insertion point before the first function or after all fields
+    // if there are no functions within the class
+    auto funcs = classOp.getBlock().getOps<FuncOp>();
+    auto it = funcs.begin();
+    if (it != funcs.end())
+      rewriter.setInsertionPoint(*it);
+    else
+      rewriter.setInsertionPointToEnd(&classOp.getBlock());
----------------
banach-space wrote:

What's the benefit of this conditionality? Why not always use `rewriter.setInsertionPointToEnd(&classOp.getBlock());`?

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