[mlir][EmitC] Add class-name-format option to wrap-emitc-func-in-class pass (#211147)

Added the `class-name-format` option that takes a format string used to
generate the wrapper class name.

This provides greater control over the generated `ClassOp` name. When
multiple identically structured IR files need to be programmatically
lowered with this pass and linked, there will be naming conflicts which
this patch provides a method of resolution for.

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

@@ -56,7 +56,11 @@ def WrapFuncInClassPass : Pass<"wrap-emitc-func-in-class", "ModuleOp"> {
       Option<"funcName", "func-name", "std::string",
       /*default=*/[{"operator()"}],
       "The name of the newly generated member function with body "
-      "matching the input function.">
+      "matching the input function.">,
+      Option<"classNameFormat", "class-name-format", "std::string",
+      /*default=*/[{"{}Class"}],
+      "Format string for the generated class name where '{}' "
+      "is replaced with the function name.">
   ];
 }
 

mlir/include/mlir/Dialect/EmitC/Transforms/Transforms.h

@@ -35,7 +35,7 @@ void populateExpressionPatterns(RewritePatternSet &patterns);
 //===----------------------------------------------------------------------===//
 
 void populateWrapFuncInClass(
-    RewritePatternSet &patterns, StringRef funcName,
+    RewritePatternSet &patterns, StringRef funcName, StringRef classNameFormat,
     DenseMap<FuncOp, llvm::DenseSet<GlobalOp>> &globalsToMove);
 
 //===----------------------------------------------------------------------===//

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

@@ -17,6 +17,7 @@
 #include "mlir/Transforms/WalkPatternRewriteDriver.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 using namespace emitc;
@@ -47,7 +48,8 @@ struct WrapFuncInClassPass
     });
 
     RewritePatternSet patterns(&getContext());
-    populateWrapFuncInClass(patterns, funcName, globalsUsedByFuncs);
+    populateWrapFuncInClass(patterns, funcName, classNameFormat,
+                            globalsUsedByFuncs);
 
     walkAndApplyPatterns(moduleOp, std::move(patterns));
 
@@ -67,15 +69,16 @@ struct WrapFuncInClassPass
 class WrapFuncInClass : public OpRewritePattern<FuncOp> {
 public:
   WrapFuncInClass(
-      MLIRContext *context, StringRef funcName,
+      MLIRContext *context, StringRef funcName, StringRef classNameFormat,
       const DenseMap<FuncOp, llvm::DenseSet<GlobalOp>> &globalsToMove)
       : OpRewritePattern<FuncOp>(context), funcName(funcName),
-        globalsToMove(globalsToMove) {}
+        classNameFormat(classNameFormat), globalsToMove(globalsToMove) {}
 
   LogicalResult matchAndRewrite(FuncOp funcOp,
                                 PatternRewriter &rewriter) const override {
 
-    auto className = funcOp.getSymNameAttr().str() + "Class";
+    std::string className =
+        llvm::formatv(classNameFormat.c_str(), funcOp.getName());
     ClassOp newClassOp = ClassOp::create(rewriter, funcOp.getLoc(), className);
 
     SmallVector<std::pair<StringAttr, TypeAttr>> fields;
@@ -149,13 +152,18 @@ private:
   /// function.
   std::string funcName;
 
+  /// Format string used to create the wrapper class name where '{}' is
+  /// replaced with the original function name.
+  std::string classNameFormat;
+
   /// Map of FuncOp and the GlobalOps it uses which need to be moved into the
   /// ClassOp wrapper.
   DenseMap<FuncOp, llvm::DenseSet<GlobalOp>> globalsToMove;
 };
 
 void mlir::emitc::populateWrapFuncInClass(
-    RewritePatternSet &patterns, StringRef funcName,
+    RewritePatternSet &patterns, StringRef funcName, StringRef classNameFormat,
     DenseMap<FuncOp, DenseSet<GlobalOp>> &globalsToMove) {
-  patterns.add<WrapFuncInClass>(patterns.getContext(), funcName, globalsToMove);
+  patterns.add<WrapFuncInClass>(patterns.getContext(), funcName,
+                                classNameFormat, globalsToMove);
 }

mlir/test/Dialect/EmitC/wrap-func-in-class.mlir

@@ -1,5 +1,6 @@
 // RUN: mlir-opt %s -wrap-emitc-func-in-class -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -wrap-emitc-func-in-class=func-name=execute -split-input-file | FileCheck %s --check-prefixes=EXECUTE
+// RUN: mlir-opt %s -wrap-emitc-func-in-class=class-name-format=Custom_{} -split-input-file | FileCheck %s --check-prefixes=CLASS-NAME-FORMAT
 
 emitc.func @foo(%arg0 : !emitc.array<1xf32>) {
   emitc.call_opaque "bar" (%arg0) : (!emitc.array<1xf32>) -> ()
@@ -18,6 +19,8 @@ emitc.func @foo(%arg0 : !emitc.array<1xf32>) {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {
+
 // -----
 
 module attributes { } {
@@ -59,6 +62,8 @@ module attributes { } {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_model {
+
 // -----
 // Tests that GlobalOps are moved into the ClassOp wrapper correctly as fields
 
@@ -81,6 +86,8 @@ module attributes { } {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {
+
 // -----
 // Tests that only GlobalOps that are used within a function are moved into the
 // ClassOp wrapper as fields
@@ -108,6 +115,8 @@ module attributes { } {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {
+
 // -----
 // Tests that when multiple functions use different globals, only the used globals
 // are moved into their respective ClassOp wrappers as fields.
@@ -147,6 +156,9 @@ module attributes { } {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {
+// CLASS-NAME-FORMAT: emitc.class @Custom_bar {
+
 // -----
 // Tests that when multiple functions use the same global, the global is moved
 // into each ClassOp wrapper as a field and erased from the module.
@@ -183,6 +195,9 @@ module attributes { } {
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
 
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {
+// CLASS-NAME-FORMAT: emitc.class @Custom_bar {
+
 // -----
 // Tests that multiple uses of the same global in a function result in a single field.
 
@@ -208,3 +223,5 @@ module attributes { } {
 
 // EXECUTE-NOT: operator
 // EXECUTE: execute()
+
+// CLASS-NAME-FORMAT: emitc.class @Custom_foo {