[mlir][CAPI] Allow running pass manager on any operation

[rkayaith]

Summary

mlirPassManagerRun is currently restricted to running on
builtin.module ops, but this restriction doesn't exist on the C++
side. This renames it to mlirPassManagerRunOnOp and updates it to take
MlirOperation instead of MlirModule.

Depends on D143352

[rkayaith]

fix parse calls

[mehdi_amini]

mlir/include/mlir-c/Pass.h

75

Can we delete the old API and name the new one mlirPassManagerRunOnOp? That will make the symbol change name and avoid weird behavior for anyone using the C API for linking to shared library for example.

[rkayaith]

mlir/include/mlir-c/Pass.h

75

No objections from me, but I don't quite understand when we should change the symbol names for these API breaks. Here's a comment from a previous change: https://reviews.llvm.org/D136404#3885547

[rkayaith]

mlir/include/mlir-c/Pass.h

75

@ftynse any thoughts here?

[rkayaith]

rename to mlirPassManagerRunOnOp

mlir/include/mlir-c/Pass.h

 static inline bool mlirPassManagerIsNull(MlirPassManager passManager) {
   return !passManager.ptr;
 }
 
 /// Cast a top-level PassManager to a generic OpPassManager.
 MLIR_CAPI_EXPORTED MlirOpPassManager
 mlirPassManagerGetAsOpPassManager(MlirPassManager passManager);
 
-/// Run the provided `passManager` on the given `module`.
+/// Run the provided `passManager` on the given `op`.
 MLIR_CAPI_EXPORTED MlirLogicalResult
-mlirPassManagerRun(MlirPassManager passManager, MlirModule module);
+mlirPassManagerRunOnOp(MlirPassManager passManager, MlirOperation op);
 
 /// Enable mlir-print-ir-after-all.
 MLIR_CAPI_EXPORTED void
 mlirPassManagerEnableIRPrinting(MlirPassManager passManager);
 
 /// Enable / disable verify-each.
 MLIR_CAPI_EXPORTED void
 mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable);

mlir/lib/Bindings/Python/Pass.cpp

py::class_<PyPassManager>(m, "PassManager", py::module_local())
               throw SetPyError(PyExc_ValueError, std::string(errorMsg.join()));
           },
           py::arg("pipeline"),
           "Add textual pipeline elements to the pass manager. Throws a "
           "ValueError if the pipeline can't be parsed.")
       .def(
           "run",
           [](PyPassManager &passManager, PyModule &module) {
-            MlirLogicalResult status =
-                mlirPassManagerRun(passManager.get(), module.get());
+            MlirLogicalResult status = mlirPassManagerRunOnOp(
+                passManager.get(), mlirModuleGetOperation(module.get()));
             if (mlirLogicalResultIsFailure(status))
               throw SetPyError(PyExc_RuntimeError,
                                "Failure while executing pass pipeline.");
           },
           py::arg("module"),
           "Run the pass manager on the provided module, throw a RuntimeError "
           "on failure.")
       .def(

mlir/lib/CAPI/IR/Pass.cpp

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)));
+MlirLogicalResult mlirPassManagerRunOnOp(MlirPassManager passManager,
+                                         MlirOperation op) {
+  return wrap(unwrap(passManager)->run(unwrap(op)));
 }
 
 void mlirPassManagerEnableIRPrinting(MlirPassManager passManager) {
   return unwrap(passManager)->enableIRPrinting();
 }
 
 void mlirPassManagerEnableVerifier(MlirPassManager passManager, bool enable) {
   unwrap(passManager)->enableVerifier(enable);

mlir/test/CAPI/execution_engine.c

 
 void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
   MlirPassManager pm = mlirPassManagerCreate(ctx);
   MlirOpPassManager opm = mlirPassManagerGetNestedUnder(
       pm, mlirStringRefCreateFromCString("func.func"));
   mlirPassManagerAddOwnedPass(pm, mlirCreateConversionConvertFuncToLLVMPass());
   mlirOpPassManagerAddOwnedPass(
       opm, mlirCreateConversionArithToLLVMConversionPass());
-  MlirLogicalResult status = mlirPassManagerRun(pm, module);
+  MlirLogicalResult status =
+      mlirPassManagerRunOnOp(pm, mlirModuleGetOperation(module));
   if (mlirLogicalResultIsFailure(status)) {
     fprintf(stderr, "Unexpected failure running pass pipeline\n");
     exit(2);
   }
   mlirPassManagerDestroy(pm);
 }
 
 // CHECK-LABEL: Running test 'testSimpleExecution'

mlir/test/CAPI/pass.c

static void registerAllUpstreamDialects(MlirContext ctx) {
   mlirContextAppendDialectRegistry(ctx, registry);
   mlirDialectRegistryDestroy(registry);
 }
 
 void testRunPassOnModule(void) {
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
 
-  MlirModule module = mlirModuleCreateParse(
-      ctx,
-      // clang-format off
-                            mlirStringRefCreateFromCString(
-"func.func @foo(%arg0 : i32) -> i32 {                                   \n"
-"  %res = arith.addi %arg0, %arg0 : i32                                     \n"
-"  return %res : i32                                                        \n"
-"}"));
-  // clang-format on
-  if (mlirModuleIsNull(module)) {
-    fprintf(stderr, "Unexpected failure parsing module.\n");
+  const char *funcAsm = //
+      "func.func @foo(%arg0 : i32) -> i32 {   \n"
+      "  %res = arith.addi %arg0, %arg0 : i32 \n"
+      "  return %res : i32                    \n"
+      "}                                      \n";
+  MlirOperation func =
+      mlirOperationCreateParse(ctx, mlirStringRefCreateFromCString(funcAsm),
+                               mlirStringRefCreateFromCString("funcAsm"));
+  if (mlirOperationIsNull(func)) {
+    fprintf(stderr, "Unexpected failure parsing asm.\n");
     exit(EXIT_FAILURE);
   }
 
   // Run the print-op-stats pass on the top-level module:
   // CHECK-LABEL: Operations encountered:
   // CHECK: arith.addi        , 1
   // CHECK: func.func      , 1
   // CHECK: func.return        , 1
   {
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
     mlirPassManagerAddOwnedPass(pm, printOpStatPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, func);
     if (mlirLogicalResultIsFailure(success)) {
       fprintf(stderr, "Unexpected failure running pass manager.\n");
       exit(EXIT_FAILURE);
     }
     mlirPassManagerDestroy(pm);
   }
-  mlirModuleDestroy(module);
+  mlirOperationDestroy(func);
   mlirContextDestroy(ctx);
 }
 
 void testRunPassOnNestedModule(void) {
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
 
-  MlirModule module = mlirModuleCreateParse(
-      ctx,
-      // clang-format off
-                            mlirStringRefCreateFromCString(
-"func.func @foo(%arg0 : i32) -> i32 {                                   \n"
-"  %res = arith.addi %arg0, %arg0 : i32                                     \n"
-"  return %res : i32                                                        \n"
-"}                                                                          \n"
-"module {                                                                   \n"
-"  func.func @bar(%arg0 : f32) -> f32 {                                     \n"
-"    %res = arith.addf %arg0, %arg0 : f32                                   \n"
-"    return %res : f32                                                      \n"
-"  }                                                                        \n"
-"}"));
-  // clang-format on
-  if (mlirModuleIsNull(module))
+  const char *moduleAsm = //
+      "module {                                   \n"
+      "  func.func @foo(%arg0 : i32) -> i32 {     \n"
+      "    %res = arith.addi %arg0, %arg0 : i32   \n"
+      "    return %res : i32                      \n"
+      "  }                                        \n"
+      "  module {                                 \n"
+      "    func.func @bar(%arg0 : f32) -> f32 {   \n"
+      "      %res = arith.addf %arg0, %arg0 : f32 \n"
+      "      return %res : f32                    \n"
+      "    }                                      \n"
+      "  }                                        \n"
+      "}                                          \n";
+  MlirOperation module =
+      mlirOperationCreateParse(ctx, mlirStringRefCreateFromCString(moduleAsm),
+                               mlirStringRefCreateFromCString("moduleAsm"));
+  if (mlirOperationIsNull(module))
     exit(1);
 
   // Run the print-op-stats pass on functions under the top-level module:
   // CHECK-LABEL: Operations encountered:
   // CHECK: arith.addi        , 1
   // CHECK: func.func      , 1
   // CHECK: func.return        , 1
   {
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     MlirOpPassManager nestedFuncPm = mlirPassManagerGetNestedUnder(
         pm, mlirStringRefCreateFromCString("func.func"));
     MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
     mlirOpPassManagerAddOwnedPass(nestedFuncPm, printOpStatPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsFailure(success))
       exit(2);
     mlirPassManagerDestroy(pm);
   }
   // Run the print-op-stats pass on functions under the nested module:
   // CHECK-LABEL: Operations encountered:
   // CHECK: arith.addf        , 1
   // CHECK: func.func      , 1
   // CHECK: func.return        , 1
   {
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     MlirOpPassManager nestedModulePm = mlirPassManagerGetNestedUnder(
         pm, mlirStringRefCreateFromCString("builtin.module"));
     MlirOpPassManager nestedFuncPm = mlirOpPassManagerGetNestedUnder(
         nestedModulePm, mlirStringRefCreateFromCString("func.func"));
     MlirPass printOpStatPass = mlirCreateTransformsPrintOpStats();
     mlirOpPassManagerAddOwnedPass(nestedFuncPm, printOpStatPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsFailure(success))
       exit(2);
     mlirPassManagerDestroy(pm);
   }
 
-  mlirModuleDestroy(module);
+  mlirOperationDestroy(module);
   mlirContextDestroy(ctx);
 }
 
 static void printToStderr(MlirStringRef str, void *userData) {
   (void)userData;
   fwrite(str.data, 1, str.length, stderr);
 }
 
return (MlirExternalPassCallbacks){testConstructExternalPass,
                                      testDestructExternalPass, initializePass,
                                      testCloneExternalPass, runPass};
 }
 
 void testExternalPass(void) {
   MlirContext ctx = mlirContextCreate();
   registerAllUpstreamDialects(ctx);
 
-  MlirModule module = mlirModuleCreateParse(
-      ctx,
-      // clang-format off
-      mlirStringRefCreateFromCString(
-"func.func @foo(%arg0 : i32) -> i32 {                                   \n"
-"  %res = arith.addi %arg0, %arg0 : i32                                     \n"
-"  return %res : i32                                                        \n"
-"}"));
-  // clang-format on
-  if (mlirModuleIsNull(module)) {
+  const char *moduleAsm = //
+      "module {                                 \n"
+      "  func.func @foo(%arg0 : i32) -> i32 {   \n"
+      "    %res = arith.addi %arg0, %arg0 : i32 \n"
+      "    return %res : i32                    \n"
+      "  }                                      \n"
+      "}";
+  MlirOperation module =
+      mlirOperationCreateParse(ctx, mlirStringRefCreateFromCString(moduleAsm),
+                               mlirStringRefCreateFromCString("moduleAsm"));
+  if (mlirOperationIsNull(module)) {
     fprintf(stderr, "Unexpected failure parsing module.\n");
     exit(EXIT_FAILURE);
   }
 
   MlirStringRef description = mlirStringRefCreateFromCString("");
   MlirStringRef emptyOpName = mlirStringRefCreateFromCString("");
 
   MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
// Run a generic pass
 
     if (userData.constructCallCount != 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     mlirPassManagerAddOwnedPass(pm, externalPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsFailure(success)) {
       fprintf(stderr, "Unexpected failure running external pass.\n");
       exit(EXIT_FAILURE);
     }
 
     if (userData.runCallCount != 1) {
       fprintf(stderr, "Expected runCallCount to be 1\n");
       exit(EXIT_FAILURE);
if (userData.constructCallCount != 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     MlirOpPassManager nestedFuncPm =
         mlirPassManagerGetNestedUnder(pm, funcOpName);
     mlirOpPassManagerAddOwnedPass(nestedFuncPm, externalPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsFailure(success)) {
       fprintf(stderr, "Unexpected failure running external operation pass.\n");
       exit(EXIT_FAILURE);
     }
 
     // Since this is a nested pass, it can be cloned and run in parallel
     if (userData.cloneCallCount != userData.constructCallCount - 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
// Run a pass with `initialize` set
 
     if (userData.constructCallCount != 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     mlirPassManagerAddOwnedPass(pm, externalPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsFailure(success)) {
       fprintf(stderr, "Unexpected failure running external pass.\n");
       exit(EXIT_FAILURE);
     }
 
     if (userData.initializeCallCount != 1) {
       fprintf(stderr, "Expected initializeCallCount to be 1\n");
       exit(EXIT_FAILURE);
// Run a pass that fails during `initialize`
 
     if (userData.constructCallCount != 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     mlirPassManagerAddOwnedPass(pm, externalPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsSuccess(success)) {
       fprintf(
           stderr,
           "Expected failure running pass manager on failing external pass.\n");
       exit(EXIT_FAILURE);
     }
 
     if (userData.initializeCallCount != 1) {
// Run a pass that fails during `run`
 
     if (userData.constructCallCount != 1) {
       fprintf(stderr, "Expected constructCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     MlirPassManager pm = mlirPassManagerCreate(ctx);
     mlirPassManagerAddOwnedPass(pm, externalPass);
-    MlirLogicalResult success = mlirPassManagerRun(pm, module);
+    MlirLogicalResult success = mlirPassManagerRunOnOp(pm, module);
     if (mlirLogicalResultIsSuccess(success)) {
       fprintf(
           stderr,
           "Expected failure running pass manager on failing external pass.\n");
       exit(EXIT_FAILURE);
     }
 
     if (userData.runCallCount != 1) {
       fprintf(stderr, "Expected runCallCount to be 1\n");
       exit(EXIT_FAILURE);
     }
 
     mlirPassManagerDestroy(pm);
 
     if (userData.destructCallCount != userData.constructCallCount) {
       fprintf(stderr, "Expected destructCallCount to be equal to "
                       "constructCallCount\n");
       exit(EXIT_FAILURE);
     }
   }
 
   mlirTypeIDAllocatorDestroy(typeIDAllocator);
-  mlirModuleDestroy(module);
+  mlirOperationDestroy(module);
   mlirContextDestroy(ctx);
 }
 
 int main(void) {
   testRunPassOnModule();
   testRunPassOnNestedModule();
   testPrintPassPipeline();
   testParsePassPipeline();
   testParseErrorCapture();
   testExternalPass();
   return 0;
 }