Added tests for the wrap-emitc-func-in-class

Jaddyen · Jaddyen · commit 317691048206 · 2025-06-20T18:29:54.000Z
diff --git a/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td b/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td
@@ -28,6 +28,9 @@ def WrapFuncInClassPass : Pass<"wrap-emitc-func-in-class"> {
     to a new `execute` method within the class.
   }];
   let dependentDialects = ["emitc::EmitCDialect"];
+  let options = [Option<
+      "namedAttribute", "named-attribute", "std::string", "\"\"",
+      "Name of the attribute to look for field names on function arguments">];
 }
 
 #endif // MLIR_DIALECT_EMITC_TRANSFORMS_PASSES
diff --git a/mlir/include/mlir/Dialect/EmitC/Transforms/Transforms.h b/mlir/include/mlir/Dialect/EmitC/Transforms/Transforms.h
@@ -29,7 +29,8 @@ ExpressionOp createExpression(Operation *op, OpBuilder &builder);
 void populateExpressionPatterns(RewritePatternSet &patterns);
 
 /// Populates 'patterns' with func-related patterns.
-void populateFuncPatterns(RewritePatternSet &patterns);
+void populateFuncPatterns(RewritePatternSet &patterns,
+                          const std::string &namedAttribute);
 
 } // namespace emitc
 } // namespace mlir
diff --git a/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp b/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp
@@ -12,14 +12,13 @@
 #include "mlir/Dialect/EmitC/Transforms/Transforms.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeRange.h"
-#include "mlir/IR/ValueRange.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/GraphWriter.h"
+#include "llvm/Support/LogicalResult.h"
 
 namespace mlir {
 namespace emitc {
@@ -31,12 +30,13 @@ 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);
+    populateFuncPatterns(patterns, namedAttribute);
 
     if (failed(applyPatternsGreedily(rootOp, std::move(patterns))))
       return signalPassFailure();
@@ -54,16 +54,13 @@ struct WrapFuncInClassPass
 using namespace mlir;
 using namespace mlir::emitc;
 
-static bool validOp(Operation &opToClone) {
-  return isa<emitc::ConstantOp>(opToClone) ||
-         isa<emitc::SubscriptOp>(opToClone) || isa<emitc::LoadOp>(opToClone) ||
-         isa<emitc::AddOp>(opToClone) || isa<emitc::AssignOp>(opToClone) ||
-         isa<emitc::ReturnOp>(opToClone);
-}
-
 class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {
+private:
+  std::string attributeName;
+
 public:
-  using OpRewritePattern<emitc::FuncOp>::OpRewritePattern;
+  WrapFuncInClass(MLIRContext *context, const std::string &attrName)
+      : OpRewritePattern<emitc::FuncOp>(context), attributeName(attrName) {}
 
   LogicalResult matchAndRewrite(emitc::FuncOp funcOp,
                                 PatternRewriter &rewriter) const override {
@@ -79,23 +76,25 @@ class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {
     rewriter.setInsertionPointToStart(&newClassOp.getBody().front());
 
     auto argAttrs = funcOp.getArgAttrs();
-
-    for (const auto &[arg, val] : (zip(*argAttrs, funcOp.getArguments()))) {
-      // FIXME:How can we avoid hardcoding this name?
-      // Should we loop through the dictionary and check for each named
-      // attribute if attr.getName().getValue().contains("tf_saved_model")
-      if (auto namedAttr = dyn_cast<mlir::DictionaryAttr>(arg).getNamed(
-              "tf_saved_model.index_path")) {
-        Attribute nv = namedAttr->getValue();
-        StringAttr fieldName =
-            cast<mlir::StringAttr>(cast<mlir::ArrayAttr>(nv)[0]);
-        TypeAttr typeAttr = TypeAttr::get(val.getType());
-        fields.push_back({fieldName, typeAttr});
-
-        rewriter.create<emitc::FieldOp>(funcOp.getLoc(), fieldName, typeAttr,
-                                        /* attributes*/ arg);
-      } else
-        funcOp->emitOpError("Only Covers TF models");
+    if (argAttrs) {
+      for (const auto &[arg, val] :
+           llvm::zip(*argAttrs, funcOp.getArguments())) {
+        if (auto namedAttr =
+                dyn_cast<mlir::DictionaryAttr>(arg).getNamed(attributeName)) {
+          Attribute nv = namedAttr->getValue();
+          StringAttr fieldName =
+              cast<mlir::StringAttr>(cast<mlir::ArrayAttr>(nv)[0]);
+          TypeAttr typeAttr = TypeAttr::get(val.getType());
+          fields.push_back({fieldName, typeAttr});
+
+          rewriter.create<emitc::FieldOp>(funcOp.getLoc(), fieldName, typeAttr,
+                                          /* attributes*/ arg);
+        }
+      }
+    } else {
+      funcOp->emitOpError("arguments should have attributes so we can "
+                          "initialize class fields.");
+      return failure();
     }
 
     rewriter.setInsertionPointToEnd(&newClassOp.getBody().front());
@@ -107,19 +106,20 @@ class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {
     FuncOp newFuncOp = rewriter.create<emitc::FuncOp>(
         loc, rewriter.getStringAttr("execute"), funcType);
 
-    rewriter.setInsertionPointToStart(newFuncOp.addEntryBlock());
+    rewriter.createBlock(&newFuncOp.getBody());
+    newFuncOp.getBody().takeBody(funcOp.getBody());
 
+    rewriter.setInsertionPointToStart(&newFuncOp.getBody().front());
     std::vector<Value> newArguments;
     for (auto [fieldName, attr] : fields) {
       auto arg =
           rewriter.create<emitc::GetFieldOp>(loc, attr.getValue(), fieldName);
       newArguments.push_back(arg);
     }
 
-    IRMapping mapper;
     for (auto [oldArg, newArg] :
-         llvm::zip(funcOp.getArguments(), newArguments)) {
-      mapper.map(oldArg, newArg);
+         llvm::zip(newFuncOp.getArguments(), newArguments)) {
+      rewriter.replaceAllUsesWith(oldArg, newArg);
     }
 
     while (!newFuncOp.getArguments().empty()) {
@@ -128,23 +128,12 @@ class WrapFuncInClass : public OpRewritePattern<emitc::FuncOp> {
       }
     }
 
-    // TODO: The mapper is easier to use but cloning is more expensive than
-    // moving the body. Working on changing this portion to move the body
-    // instead
-    auto body = llvm::make_early_inc_range(funcOp.getBody().front());
-    for (Operation &opToClone : body) {
-      if (validOp(opToClone)) {
-        rewriter.clone(opToClone, mapper);
-      } else {
-        opToClone.emitOpError("Unsupported operation found");
-      }
-    }
-
     rewriter.replaceOp(funcOp, newClassOp);
     return funcOp->use_empty() ? success() : failure();
   }
 };
 
-void mlir::emitc::populateFuncPatterns(RewritePatternSet &patterns) {
-  patterns.add<WrapFuncInClass>(patterns.getContext());
+void mlir::emitc::populateFuncPatterns(RewritePatternSet &patterns,
+                                       const std::string &namedAttribute) {
+  patterns.add<WrapFuncInClass>(patterns.getContext(), namedAttribute);
 }
diff --git a/mlir/test/Dialect/EmitC/wrap_emitc_func_in_class.mlir b/mlir/test/Dialect/EmitC/wrap_emitc_func_in_class.mlir
@@ -1,7 +1,7 @@
-// RUN: mlir-opt  %s --wrap-emitc-func-in-class
+// RUN: mlir-opt --wrap-emitc-func-in-class='named-attribute=tf_saved_model.index_path' %s | FileCheck %s
 
 module attributes {tf_saved_model.semantics, tfl.description = "MLIR Converted.", tfl.schema_version = 3 : i32} {
-  emitc.func @main(%arg0: !emitc.array<1xf32> {tf_saved_model.index_path = ["another_feature"]}, %arg1: !emitc.array<1xf32> {tf_saved_model.index_path = ["some_feature"]}, %arg2: !emitc.array<1xf32> {tf_saved_model.index_path = ["output_0"]}) attributes {tf.entry_function = {inputs = "serving_default_another_feature:0,serving_default_some_feature:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+  emitc.func @Model(%arg0: !emitc.array<1xf32> {tf_saved_model.index_path = ["another_feature"]}, %arg1: !emitc.array<1xf32> {tf_saved_model.index_path = ["some_feature"]}, %arg2: !emitc.array<1xf32> {tf_saved_model.index_path = ["output_0"]}) attributes {tf.entry_function = {inputs = "serving_default_another_feature:0,serving_default_some_feature:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
     %0 = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
     %1 = subscript %arg1[%0] : (!emitc.array<1xf32>, !emitc.size_t) -> !emitc.lvalue<f32>
     %2 = load %1 : <f32>
@@ -13,3 +13,25 @@ module attributes {tf_saved_model.semantics, tfl.description = "MLIR Converted."
     return
   }
 }
+
+// CHECK: module attributes {tf_saved_model.semantics, tfl.description = "MLIR Converted.", tfl.schema_version = 3 : i32} {
+// CHECK:   emitc.class @MyModelClass {
+// CHECK:     emitc.field @another_feature : !emitc.array<1xf32> = {tf_saved_model.index_path = ["another_feature"]}
+// CHECK:     emitc.field @some_feature : !emitc.array<1xf32> = {tf_saved_model.index_path = ["some_feature"]}
+// CHECK:     emitc.field @output_0 : !emitc.array<1xf32> = {tf_saved_model.index_path = ["output_0"]}
+// CHECK:     emitc.func @execute() {
+// CHECK:       %{{[0-9]+}} = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
+// CHECK:       %{{[0-9]+}} = get_field @another_feature : !emitc.array<1xf32>
+// CHECK:       %{{[0-9]+}} = get_field @some_feature : !emitc.array<1xf32>
+// CHECK:       %{{[0-9]+}} = get_field @output_0 : !emitc.array<1xf32>
+// CHECK:       %{{[0-9]+}} = subscript %{{[0-9]+}}[%{{[0-9]+}}] : (!emitc.array<1xf32>, !emitc.size_t) -> !emitc.lvalue<f32>
+// CHECK:       %{{[0-9]+}} = load %{{[0-9]+}} : <f32>
+// CHECK:       %{{[0-9]+}} = subscript %{{[0-9]+}}[%{{[0-9]+}}] : (!emitc.array<1xf32>, !emitc.size_t) -> !emitc.lvalue<f32>
+// CHECK:       %{{[0-9]+}} = load %{{[0-9]+}} : <f32>
+// CHECK:       %{{[0-9]+}} = add %{{[0-9]+}}, %{{[0-9]+}} : (f32, f32) -> f32
+// CHECK:       %{{[0-9]+}} = subscript %{{[0-9]+}}[%{{[0-9]+}}] : (!emitc.array<1xf32>, !emitc.size_t) -> !emitc.lvalue<f32>
+// CHECK:       assign %{{[0-9]+}} : f32 to %{{[0-9]+}} : <f32>
+// CHECK:       return
+// CHECK:     }
+// CHECK:   }
+// CHECK: }
diff --git a/mlir/test/Dialect/EmitC/wrap_emitc_func_in_class_neg.mlir b/mlir/test/Dialect/EmitC/wrap_emitc_func_in_class_neg.mlir
@@ -0,0 +1,8 @@
+// RUN: mlir-opt --wrap-emitc-func-in-class='named-attribute=tf_saved_model.index_path' %s 2>&1 | FileCheck %s
+
+emitc.func @foo(%arg0 : i32) {
+  emitc.call_opaque "bar" (%arg0) : (i32) -> ()
+  emitc.return
+}
+
+// CHECK: error: 'emitc.func' op arguments should have attributes so we can initialize class fields.

Original file line number	Diff line number	Diff line change
`@@ -28,6 +28,9 @@ def WrapFuncInClassPass : Pass<"wrap-emitc-func-in-class"> {`
`28`	`28`	to a new `execute` method within the class.
`29`	`29`	`}];`
`30`	`30`	`let dependentDialects = ["emitc::EmitCDialect"];`
	`31`	`+ let options = [Option<`
	`32`	`+ "namedAttribute", "named-attribute", "std::string", "\"\"",`
	`33`	`+ "Name of the attribute to look for field names on function arguments">];`
`31`	`34`	`}`
`32`	`35`
`33`	`36`	`#endif // MLIR_DIALECT_EMITC_TRANSFORMS_PASSES`