Specify the pass reqs

Author: Jaddyen

mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td

@@ -53,33 +53,41 @@ def WrapFuncInClassPass : Pass<"wrap-emitc-func-in-class"> {
   let dependentDialects = ["emitc::EmitCDialect"];
 }
 
-def AddReflectionMapPass : Pass<"add-reflection-map"> {
+def AddReflectionMapPass : Pass<"add-reflection-map", "ModuleOp"> {
   let summary =
       "Add a reflection map function to EmitC classes for runtime field lookup";
   let description = [{
         This pass adds a `getBufferForName` function to EmitC classes that enables 
         runtime lookup of field buffers by their string names. 
-        This enables runtime introspection and dynamic access to class fields by name,
-        which is useful for interfacing with external systems that need to access 
-        tensors/buffers by their semantic names.
+        This would require that the class has fields with attributes and a function named `execute`.
+        The `fieldop` attribute is expected to be a dictionary where:
+        - The keys are `namedAttribute`.
+        - The values are arrays containing a single string attribute.
+
+
+        Example:
 
-        Example transformation:
         ```mlir
         emitc.class @MyClass {
           emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref = ["another_feature"]}
           emitc.field @fieldName1 : !emitc.array<1xf32> {emitc.field_ref = ["some_feature"]}
           emitc.func @execute() { ... }
         }
-        ```
 
-        Becomes:
-        ```mlir
+        // becomes:
+        
         emitc.class @MyClass {
           emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref = ["another_feature"]}
           emitc.field @fieldName1 : !emitc.array<1xf32> {emitc.field_ref = ["some_feature"]}
           emitc.func @getBufferForName(%name : !emitc.opaque<"std::string_view">) -> !emitc.opaque<"char*"> {
-            %map = "emitc.constant"(){value = #emitc.opaque<"{"another_feature", reinterpret_cast<char*>(&another_feature)}, {"some_feature", reinterpret_cast<char*>(&some_feature)}">} : () -> !emitc.opaque<"std::map<std::string, char*>">
-            return %null : !emitc.opaque<"char*">
+            %0 = "emitc.constant"() <{value = #emitc.opaque<"{  { \22another_feature\22, reinterpret_cast<char*>(&fieldName0) },  { \22some_feature\22, reinterpret_cast<char*>(&fieldName1) } }">}> : () -> !emitc.opaque<"const std::map<std::string, char*>">
+            %1 = call_opaque "find"(%0, %arg0) : (!emitc.opaque<"const std::map<std::string, char*>">, !emitc.opaque<"std::string_view">) -> !emitc.opaque<"std::map<std::string, char*>::const_iterator">
+            %2 = call_opaque "end"(%0) : (!emitc.opaque<"const std::map<std::string, char*>">) -> !emitc.opaque<"std::map<std::string, char*>::const_iterator">
+            %3 = call_opaque "operator=="(%1, %2) : (!emitc.opaque<"std::map<std::string, char*>::const_iterator">, !emitc.opaque<"std::map<std::string, char*>::const_iterator">) -> i1
+            %4 = "emitc.constant"() <{value = #emitc.opaque<"nullptr">}> : () -> !emitc.opaque<"char">
+            %5 = call_opaque "second"(%1) : (!emitc.opaque<"std::map<std::string, char*>::const_iterator">) -> !emitc.opaque<"char">
+            %6 = conditional %3, %4, %5 : !emitc.opaque<"char">
+            return %6 : !emitc.opaque<"char">
           }
           emitc.func @execute() { ... }
         }

mlir/lib/Dialect/EmitC/Transforms/AddReflectionMap.cpp

@@ -23,16 +23,48 @@ namespace emitc {
 #include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
 
 namespace {
+constexpr const char *kMapLibraryHeader = "map";
+constexpr const char *kStringLibraryHeader = "string";
 class AddReflectionMapPass
     : public impl::AddReflectionMapPassBase<AddReflectionMapPass> {
   using AddReflectionMapPassBase::AddReflectionMapPassBase;
   void runOnOperation() override {
-    Operation *rootOp = getOperation();
+    mlir::ModuleOp module = getOperation();
 
     RewritePatternSet patterns(&getContext());
     populateAddReflectionMapPatterns(patterns, namedAttribute);
 
-    walkAndApplyPatterns(rootOp, std::move(patterns));
+    walkAndApplyPatterns(module, std::move(patterns));
+    bool hasMap = false;
+    bool hasString = false;
+    for (auto &op : *module.getBody()) {
+      emitc::IncludeOp includeOp = llvm::dyn_cast<mlir::emitc::IncludeOp>(op);
+      if (!includeOp)
+        continue;
+      if (includeOp.getIsStandardInclude()) {
+        if (includeOp.getInclude() == kMapLibraryHeader)
+          hasMap = true;
+        if (includeOp.getInclude() == kStringLibraryHeader)
+          hasString = true;
+      }
+    }
+
+    if (hasMap && hasString)
+      return;
+
+    mlir::OpBuilder builder(module.getBody(), module.getBody()->begin());
+    if (!hasMap) {
+      StringAttr includeAttr = builder.getStringAttr(kMapLibraryHeader);
+      builder.create<mlir::emitc::IncludeOp>(
+          module.getLoc(), includeAttr,
+          /*is_standard_include=*/builder.getUnitAttr());
+    }
+    if (!hasString) {
+      StringAttr includeAttr = builder.getStringAttr(kStringLibraryHeader);
+      builder.create<emitc::IncludeOp>(
+          module.getLoc(), includeAttr,
+          /*is_standard_include=*/builder.getUnitAttr());
+    }
   }
 };
 
@@ -50,16 +82,20 @@ class AddReflectionMapClass : public OpRewritePattern<emitc::ClassOp> {
     mlir::MLIRContext *context = rewriter.getContext();
     emitc::OpaqueType stringViewType =
         mlir::emitc::OpaqueType::get(rewriter.getContext(), "std::string_view");
-    emitc::OpaqueType charPtrType =
+    emitc::OpaqueType charType =
         mlir::emitc::OpaqueType::get(rewriter.getContext(), "char");
     emitc::OpaqueType mapType = mlir::emitc::OpaqueType::get(
         rewriter.getContext(), "const std::map<std::string, char*>");
 
     FunctionType funcType =
-        rewriter.getFunctionType({stringViewType}, {charPtrType});
+        rewriter.getFunctionType({stringViewType}, {charType});
     emitc::FuncOp executeFunc =
         classOp.lookupSymbol<mlir::emitc::FuncOp>("execute");
-    rewriter.setInsertionPoint(executeFunc);
+    if (executeFunc)
+      rewriter.setInsertionPoint(executeFunc);
+    else
+      classOp.emitError() << "ClassOp must contain a function named 'execute' "
+                             "to add reflection map";
 
     emitc::FuncOp getBufferFunc = rewriter.create<mlir::emitc::FuncOp>(
         classOp.getLoc(), "getBufferForName", funcType);
@@ -74,9 +110,8 @@ class AddReflectionMapClass : public OpRewritePattern<emitc::ClassOp> {
               fieldOp->getAttrDictionary().get("attrs")) {
         if (DictionaryAttr innerDictAttr =
                 dyn_cast<mlir::DictionaryAttr>(attrsAttr)) {
-          auto indexPathAttr = innerDictAttr.getNamed(attributeName);
-          ArrayAttr arrayAttr =
-              dyn_cast<mlir::ArrayAttr>(indexPathAttr->getValue());
+          ArrayAttr arrayAttr = dyn_cast<mlir::ArrayAttr>(
+              innerDictAttr.getNamed(attributeName)->getValue());
           if (!arrayAttr.empty()) {
             StringAttr stringAttr = dyn_cast<mlir::StringAttr>(arrayAttr[0]);
             std::string indexPath = stringAttr.getValue().str();
@@ -122,13 +157,12 @@ class AddReflectionMapClass : public OpRewritePattern<emitc::ClassOp> {
         classOp.getLoc(), rewriter.getI1Type(),
         "operator==", mlir::ValueRange{it.getResult(0), endIt.getResult(0)});
     emitc::ConstantOp nullPtr = rewriter.create<emitc::ConstantOp>(
-        classOp.getLoc(), charPtrType,
-        emitc::OpaqueAttr::get(context, "nullptr"));
+        classOp.getLoc(), charType, emitc::OpaqueAttr::get(context, "nullptr"));
     emitc::CallOpaqueOp second = rewriter.create<emitc::CallOpaqueOp>(
-        classOp.getLoc(), charPtrType, "second", it.getResult(0));
+        classOp.getLoc(), charType, "second", it.getResult(0));
 
     emitc::ConditionalOp result = rewriter.create<emitc::ConditionalOp>(
-        classOp.getLoc(), charPtrType, isEnd.getResult(0), nullPtr.getResult(),
+        classOp.getLoc(), charType, isEnd.getResult(0), nullPtr.getResult(),
         second.getResult(0));
 
     rewriter.create<emitc::ReturnOp>(classOp.getLoc(), result.getResult());

mlir/test/Dialect/EmitC/add_reflection_map.mlir

@@ -21,6 +21,8 @@ emitc.class @mainClass {
 }
 
 // CHECK: module {
+// CHECK-NEXT:   emitc.include <"map">
+// CHECK-NEXT:   emitc.include <"string">
 // CHECK-NEXT:   emitc.class @mainClass {
 // CHECK-NEXT:     emitc.field @fieldName0 : !emitc.array<1xf32> {emitc.field_ref = ["another_feature"]}
 // CHECK-NEXT:     emitc.field @fieldName1 : !emitc.array<1xf32> {emitc.field_ref = ["some_feature"]}