Add facilities to print/parse a pass pipeline through the C API

[mehdi_amini]

Summary

This also includes and exercise a register function for individual
passes.

mlir/include/mlir-c/Pass.h

 #undef DEFINE_C_API_STRUCT
 
 /** Create a new top-level PassManager. */
 MlirPassManager mlirPassManagerCreate(MlirContext ctx);
 
 /** Destroy the provided PassManager. */
 void mlirPassManagerDestroy(MlirPassManager passManager);
 
+/** Cast a top-level PassManager to a generic OpPassManager. */
+MlirOpPassManager
+mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
+
 /** Run the provided `passManager` on the given `module`. */
 MlirLogicalResult mlirPassManagerRun(MlirPassManager passManager,
                                      MlirModule module);
 
 /** Nest an OpPassManager under the top-level PassManager, the nested
  * passmanager will only run on operations matching the provided name.
  * The returned OpPassManager will be destroyed when the parent is destroyed.
  * To further nest more OpPassManager under the newly returned one, see

 
 /** Add a pass and transfer ownership to the provided mlirOpPassManager. If the
  * pass is not a generic operation pass or matching the type of the provided
  * PassManager, a new OpPassManager is implicitly nested under the provided
  * PassManager. */
 void mlirOpPassManagerAddOwnedPass(MlirOpPassManager passManager,
                                    MlirPass pass);
 
+/** Print a textual MLIR pass pipeline by sending chunks of the string
+ * representation and forwarding `userData to `callback`. Note that the callback
+ * may be called several times with consecutive chunks of the string. */
+void mlirPrintPassPipeline(MlirOpPassManager passManager,
+                           MlirStringCallback callback, void *userData);
+
+/** Parse a textual MLIR pass pipeline and add it to the provided OpPassManager.
+ */
+MlirLogicalResult mlirParsePassPipeline(MlirOpPassManager passManager,
+                                        MlirStringRef pipeline);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_PASS_H

mlir/lib/CAPI/IR/Pass.cpp

 MlirPassManager mlirPassManagerCreate(MlirContext ctx) {
   return wrap(new PassManager(unwrap(ctx)));
 }
 
 void mlirPassManagerDestroy(MlirPassManager passManager) {
   delete unwrap(passManager);
 }
 
+MlirOpPassManager
+mlirPassManagerGetAsOpPassManager(MlirPassManager passManager) {
+  return wrap(static_cast<OpPassManager *>(unwrap(passManager)));
+}
+
 MlirLogicalResult mlirPassManagerRun(MlirPassManager passManager,
                                      MlirModule module) {
   return wrap(unwrap(passManager)->run(unwrap(module)));
 }
 
 MlirOpPassManager mlirPassManagerGetNestedUnder(MlirPassManager passManager,
                                                 MlirStringRef operationName) {
   return wrap(&unwrap(passManager)->nest(unwrap(operationName)));
 }
 
 MlirOpPassManager mlirOpPassManagerGetNestedUnder(MlirOpPassManager passManager,
                                                   MlirStringRef operationName) {
   return wrap(&unwrap(passManager)->nest(unwrap(operationName)));
 }
 
 void mlirPassManagerAddOwnedPass(MlirPassManager passManager, MlirPass pass) {
   unwrap(passManager)->addPass(std::unique_ptr<Pass>(unwrap(pass)));
 }
 
 void mlirOpPassManagerAddOwnedPass(MlirOpPassManager passManager,
                                    MlirPass pass) {
   unwrap(passManager)->addPass(std::unique_ptr<Pass>(unwrap(pass)));
 }
+
+void mlirPrintPassPipeline(MlirOpPassManager passManager,
+                           MlirStringCallback callback, void *userData) {
+  detail::CallbackOstream stream(callback, userData);
+  unwrap(passManager)->printAsTextualPipeline(stream);
+}
+
+MlirLogicalResult mlirParsePassPipeline(MlirOpPassManager passManager,
+                                        MlirStringRef pipeline) {
+  // TODO: errors are sent to std::errs() at the moment, we should pass in a
+  // stream and redirect to a diagnostic.
+  return wrap(mlir::parsePassPipeline(unwrap(pipeline), *unwrap(passManager)));
+}

mlir/test/CAPI/pass.c

void testRunPassOnModule() {
   MlirModule module =
       mlirModuleCreateParse(ctx,
                             // clang-format off
 "func @foo(%arg0 : i32) -> i32 {                                            \n"
 "  %res = addi %arg0, %arg0 : i32                                           \n"
 "  return %res : i32                                                        \n"
 "}");
   // clang-format on
-  if (mlirModuleIsNull(module))
+  if (mlirModuleIsNull(module)) {
+    fprintf(stderr, "Unexpected failure parsing module.\n");
     exit(EXIT_FAILURE);
+  }
 
   // Run the print-op-stats pass on the top-level module:
   // CHECK-LABEL: Operations encountered:
   // CHECK: func              , 1
   // CHECK: module_terminator , 1
   // CHECK: std.addi          , 1
   // CHECK: std.return        , 1
   {
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
     mlirPassManagerAddOwnedPass(pm, printOpStatPass);
     MlirLogicalResult success = mlirPassManagerRun(pm, module);
-    if (mlirLogicalResultIsFailure(success))
+    if (mlirLogicalResultIsFailure(success)) {
+      fprintf(stderr, "Unexpected failure running pass manager.\n");
       exit(EXIT_FAILURE);
+    }
     mlirPassManagerDestroy(pm);
   }
   mlirModuleDestroy(module);
   mlirContextDestroy(ctx);
 }
 
 void testRunPassOnNestedModule() {
   MlirContext ctx = mlirContextCreate();
if (mlirLogicalResultIsFailure(success))
       exit(2);
     mlirPassManagerDestroy(pm);
   }
 
   mlirModuleDestroy(module);
   mlirContextDestroy(ctx);
 }
 
+static void printToStderr(const char *str, intptr_t len, void *userData) {
+  (void)userData;
+  fwrite(str, 1, len, stderr);
+}
+
+void testPrintPassPipeline() {
+  MlirContext ctx = mlirContextCreate();
+  MlirPassManager pm = mlirPassManagerCreate(ctx);
+  // Populate the pass-manager
+  MlirOpPassManager nestedModulePm = mlirPassManagerGetNestedUnder(
+      pm, mlirStringRefCreateFromCString("module"));
+  MlirOpPassManager nestedFuncPm = mlirOpPassManagerGetNestedUnder(
+      nestedModulePm, mlirStringRefCreateFromCString("func"));
+  MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
+  mlirOpPassManagerAddOwnedPass(nestedFuncPm, printOpStatPass);
+
+  // Print the top level pass manager
+  // CHECK: Top-level: module(func(print-op-stats))
+  fprintf(stderr, "Top-level: ");
+  mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
+                        NULL);
+  fprintf(stderr, "\n");
+
+  // Print the pipeline nested one level down
+  // CHECK: Nested Module: func(print-op-stats)
+  fprintf(stderr, "Nested Module: ");
+  mlirPrintPassPipeline(nestedModulePm, printToStderr, NULL);
+  fprintf(stderr, "\n");
+
+  // Print the pipeline nested two levels down
+  // CHECK: Nested Module>Func: print-op-stats
+  fprintf(stderr, "Nested Module>Func: ");
+  mlirPrintPassPipeline(nestedFuncPm, printToStderr, NULL);
+  fprintf(stderr, "\n");
+
+  mlirPassManagerDestroy(pm);
+  mlirContextDestroy(ctx);
+}
+
+void testParsePassPipeline() {
+  MlirContext ctx = mlirContextCreate();
+  MlirPassManager pm = mlirPassManagerCreate(ctx);
+  // Try parse a pipeline.
+  MlirLogicalResult status = mlirParsePassPipeline(
+      mlirPassManagerGetAsOpPassManager(pm),
+      mlirStringRefCreateFromCString(
+          "module(func(print-op-stats), func(print-op-stats))"));
+  // Expect a failure, we haven't registered the print-op-stats pass yet.
+  if (mlirLogicalResultIsSuccess(status)) {
+    fprintf(stderr, "Unexpected success parsing pipeline without registering the pass\n");
+    exit(EXIT_FAILURE);
+  }
+  // Try again after registrating the pass.
+  mlirRegisterTransformsPrintOpStats();
+  status = mlirParsePassPipeline(
+      mlirPassManagerGetAsOpPassManager(pm),
+      mlirStringRefCreateFromCString(
+          "module(func(print-op-stats), func(print-op-stats))"));
+  // Expect a failure, we haven't registered the print-op-stats pass yet.
+  if (mlirLogicalResultIsFailure(status)) {
+    fprintf(stderr, "Unexpected failure parsing pipeline after registering the pass\n");
+    exit(EXIT_FAILURE);
+  }
+
+  // CHECK: Round-trip: module(func(print-op-stats), func(print-op-stats))
+  fprintf(stderr, "Round-trip: ");
+  mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
+                        NULL);
+  fprintf(stderr, "\n");
+}
+
 int main() {
   testRunPassOnModule();
   testRunPassOnNestedModule();
+  testPrintPassPipeline();
+  testParsePassPipeline();
   return 0;
 }

mlir/tools/mlir-tblgen/PassCAPIGen.cpp

groupName("prefix",
               llvm::cl::desc("The prefix to use for this group of passes. The "
                              "form will be mlirCreate<prefix><passname>, the "
                              "prefix can avoid conflicts across libraries."),
               llvm::cl::cat(passGenCat));
 
 const char *const passDecl = R"(
 /* Create {0} Pass. */
 MlirPass mlirCreate{0}{1}();
+void mlirRegister{0}{1}();
 
 )";
 
 const char *const fileHeader = R"(
 /* Autogenerated by mlir-tblgen; don't manually edit. */
 
 #include "mlir-c/Pass.h"
 
static bool emitCAPIHeader(const llvm::RecordKeeper &records, raw_ostream &os) {
   os << fileFooter;
   return false;
 }
 
 const char *const passCreateDef = R"(
 MlirPass mlirCreate{0}{1}() {
   return wrap({2}.release());
 }
+void mlirRegister{0}{1}() {
+  register{1}Pass();
+}
 
 )";
 
 static bool emitCAPIImpl(const llvm::RecordKeeper &records, raw_ostream &os) {
   os << "/* Autogenerated by mlir-tblgen; don't manually edit. */";
   for (const auto *def : records.getAllDerivedDefinitions("PassBase")) {
     Pass pass(def);
     StringRef defName = pass.getDef()->getName();