[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 }
----------------
banach-space wrote:

Here and in `After`, the formatting is off:
```suggestion
///   emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref =
///      ["another_feature"]} 
///   emitc.func @bar() { return }
```

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