[mlir][capi] Add dynamic pass creation to MLIR C-API

[trilorez]

Summary

Adds the ability to create dynamic passes using the C-API. This allows passes
to be written in C or languages that use the C-bindings.

[trilorez]

Rebase main, rename builtin.func to func.func

[rriddle]

mlir/include/mlir-c/Pass.h

127

Please use complete sentences for comments.

134

I would try to model more closely to the C++ API, which doesn't return a result.

mlir/lib/CAPI/IR/Pass.cpp

89

nit: Drop the newline here.

109

nit: Drop trivial braces here and below

[mehdi_amini]

LG overall, I suggested another name for this inline.

mlir/include/mlir-c/Pass.h

123

We probably should expand the doc here, this isn't a "core context" for which we can just point to the C++ doc.

I would rather name this differently than "Dynamic Pass" because it's too close to the dynamic pipeline we have for passes. What about "External Passes" or "Pluggable Pass"?
Something like this:

//===----------------------------------------------------------------------===//
// External Pass API.
// 
// This API allows to define passes outside of MLIR, not necessarily in
// C++, and register them with the MLIR pass management
// infrastructure.
//
//===----------------------------------------------------------------------===//

125

127

I would even point to the C++ method also by adding something like "(see Pass::initialize())"
We have a 1-1 mapping here right?

[trilorez]

mlir/include/mlir-c/Pass.h

123

I like External Pass, it makes the use case pretty clear

127

There is a 1-1 mapping except that initialize and deinitialize map to the C++ constructor and destructor. Maybe construct and destruct would be a more clear mapping?

134

The result is used to signal the pass failure. I suppose I can add an argument to run that is a pointer to a bool or MlirLogicalResult that can be set to signal pass failure.

Otherwise I can add callback argument to run that would directly call signalPassFailure, which would need to be called with a pointer to the C++ pass class that would also be passed to run. That approach seems a bit cumbersome but does model the C++ API more closely.

Which approach would you prefer?

[mehdi_amini]

mlir/include/mlir-c/Pass.h

127

Ah right, that's confusing with the initialize() method that I believe exists on passes (should it be exposed here as well by the way?)

[trilorez]

mlir/include/mlir-c/Pass.h

127

Oh, I missed that Pass::initialize() exists. Seems like that it is worth exposing as an optional callback.

[trilorez]

Renamed DynamicPass to ExternalPass.
Renamed initialize/deinitialize to construct/destruct.
Exposed Pass::initialize.
Addressed nits.

[trilorez]

Rebase main

[trilorez]

mlir/include/mlir-c/Pass.h

134

@rriddle Any thoughts on these approaches?

[rriddle]

mlir/include/mlir-c/Pass.h

134

I think that passing in the pointer of the MLIR pass itself would likely be better in the long run. This API will/may likely be expanded to expose access to things like analyses, which would also require access to the pass pointer itself (same with running nested pass managers, or various other types of Pass features). What do you think of that approach?

[rriddle]

Sorry for the late review, I just came back from OOO this week.

mlir/include/mlir-c/Pass.h

157

?

mlir/lib/CAPI/IR/Pass.cpp

120

nit: Drop the trivial braces, also can you spell out auto here?

mlir/test/CAPI/pass.c

284

Technically TypeIDs are supposed to be aligned by 8. I don't think anything would break right now, but it could.

[trilorez]

Added MlirExternalPass to external pass run callback, which can be used to
signal pass failure.

Added alignment check to mlirTypeIDCreate.

Addressed nits.

mlir/include/mlir-c/Pass.h

134

That makes sense for future flexibility. I ended up exposing the pass in the callback as a separate MlirExternalPass type so it can be safely casted to the c++ ExternalPass type and call signalPassFailure (which was otherwise protected on Pass). I kept the return type of mlirCreateExternalPass as MlirPass since its only valid to directly use it as a generic pass but that might be kind of confusing.

mlir/test/CAPI/pass.c

284

I updated these to use aligned_alloc and added an assertion in mlirTypeIDCreate

[trilorez]

Fix formatting

[trilorez]

Use _Alignas instead of aligned_alloc

[trilorez]

Add bindings for TypeIDAllocator, use in c api tests.
This works around Windows not having c type alignment apis available.

[rriddle]

mlir/include/mlir-c/Pass.h

124

mlir/include/mlir/CAPI/IR.h

31

Should this rather be in Support.h? (Not sure what the layering here is supposed to be given that TypeID is also in this file)

mlir/lib/CAPI/IR/IR.cpp

795

Will this still hold if some low bits are stolen? (Not sure if that would happen much/ever for this CAPI).

mlir/lib/CAPI/IR/Pass.cpp

87

Missing namespace comment here.

88

Is there a downside to sinking this below the ExternalPass definition?

113

Can you spell out auto here?

127

Can we drop the llvm::?

137

Can we spell out auto here?

[trilorez]

Address comments

mlir/lib/CAPI/IR/IR.cpp

795

This api is for the initial creation of a type id from outside of MLIR, so no bits should have been stolen yet. My understanding is the bit-stealing could only happen to already created type ids.

mlir/lib/CAPI/IR/Pass.cpp

87

Weird how clang-format didn't pick up on that, added.

88

I need it before ExternalPass is defined since ExternalPass::runOnOperation uses the wrap function generated by this macro.

[trilorez]

mlir/include/mlir/CAPI/IR.h

31

We'd have to move both MlirTypeID and MlirTypeIDAllocator to Support.h since the allocate function returns an MlirTypeID as opposed to a pointer to one (so it couldn't be forward declared). I'm not sure if Support.h is a better place considering TypeID is a core MLIR data structure as opposed to auxiliary but I don't have a strong notion about the best place for it.

[rriddle]

mlir/include/mlir/CAPI/IR.h

31

I mentioned it because TypeID is defined in the Support library(same as LogicalResult), not IR.

[trilorez]

mlir/include/mlir/CAPI/IR.h

31

Oh, in that case moving it seems like the right way to go

[trilorez]

Move MlirTypeID and MlirTypeIDAllocator to Support.h

[rriddle]

mlir/test/CAPI/pass.c

515

This function is leaking memory, seems to be missing a mlirModuleDestroy here.

[trilorez]

mlir/test/CAPI/pass.c

515

Whoops, fixed here: https://github.com/llvm/llvm-project/commit/e394c4e5d972861f005379f490f316e0d1b84f21

[rriddle]

mlir/test/CAPI/pass.c

515

Thanks!

mlir/include/mlir-c/IR.h

 DEFINE_C_API_STRUCT(MlirContext, void);
 DEFINE_C_API_STRUCT(MlirDialect, void);
 DEFINE_C_API_STRUCT(MlirDialectRegistry, void);
 DEFINE_C_API_STRUCT(MlirOperation, void);
 DEFINE_C_API_STRUCT(MlirOpPrintingFlags, void);
 DEFINE_C_API_STRUCT(MlirBlock, void);
 DEFINE_C_API_STRUCT(MlirRegion, void);
 DEFINE_C_API_STRUCT(MlirSymbolTable, void);
+DEFINE_C_API_STRUCT(MlirTypeIDAllocator, void);
 
 DEFINE_C_API_STRUCT(MlirAttribute, const void);
 DEFINE_C_API_STRUCT(MlirIdentifier, const void);
 DEFINE_C_API_STRUCT(MlirLocation, const void);
 DEFINE_C_API_STRUCT(MlirModule, const void);
 DEFINE_C_API_STRUCT(MlirType, const void);
 DEFINE_C_API_STRUCT(MlirTypeID, const void);
 DEFINE_C_API_STRUCT(MlirValue, const void);

 
 /// Gets the string value of the identifier.
 MLIR_CAPI_EXPORTED MlirStringRef mlirIdentifierStr(MlirIdentifier ident);
 
 //===----------------------------------------------------------------------===//
 // TypeID API.
 //===----------------------------------------------------------------------===//
 
+/// `ptr` must be 8 byte aligned and unique to a type valid for the duration of
+/// the returned type id's usage
+MLIR_CAPI_EXPORTED MlirTypeID mlirTypeIDCreate(const void *ptr);
+
 /// Checks whether a type id is null.
 static inline bool mlirTypeIDIsNull(MlirTypeID typeID) { return !typeID.ptr; }
 
 /// Checks if two type ids are equal.
 MLIR_CAPI_EXPORTED bool mlirTypeIDEqual(MlirTypeID typeID1, MlirTypeID typeID2);
 
 /// Returns the hash value of the type id.
 MLIR_CAPI_EXPORTED size_t mlirTypeIDHashValue(MlirTypeID typeID);
 
 //===----------------------------------------------------------------------===//
+// TypeIDAllocator API.
+//===----------------------------------------------------------------------===//
+
+/// Creates a type id allocator for dynamic type id creation
+MLIR_CAPI_EXPORTED MlirTypeIDAllocator mlirTypeIDAllocatorCreate();
+
+/// Deallocates the allocator and all allocated type ids
+MLIR_CAPI_EXPORTED void
+mlirTypeIDAllocatorDestroy(MlirTypeIDAllocator allocator);
+
+/// Allocates a type id that is valid for the lifetime of the allocator
+MLIR_CAPI_EXPORTED MlirTypeID
+mlirTypeIDAllocatorAllocateTypeID(MlirTypeIDAllocator allocator);
+
+//===----------------------------------------------------------------------===//
 // Symbol and SymbolTable API.
 //===----------------------------------------------------------------------===//
 
 /// Returns the name of the attribute used to store symbol names compatible with
 /// symbol tables.
 MLIR_CAPI_EXPORTED MlirStringRef mlirSymbolTableGetSymbolAttributeName();
 
 /// Returns the name of the attribute used to store symbol visibility.

mlir/include/mlir-c/Pass.h

 // This header declares the C interface to MLIR pass manager.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_C_PASS_H
 #define MLIR_C_PASS_H
 
 #include "mlir-c/IR.h"
+#include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 //===----------------------------------------------------------------------===//
 // Opaque type declarations.

 
 #define DEFINE_C_API_STRUCT(name, storage)                                     \
   struct name {                                                                \
     storage *ptr;                                                              \
   };                                                                           \
   typedef struct name name
 
 DEFINE_C_API_STRUCT(MlirPass, void);
+DEFINE_C_API_STRUCT(MlirExternalPass, void);
 DEFINE_C_API_STRUCT(MlirPassManager, void);
 DEFINE_C_API_STRUCT(MlirOpPassManager, void);
 
 #undef DEFINE_C_API_STRUCT
 
+//===----------------------------------------------------------------------===//
+// PassManager/OpPassManager APIs.
+//===----------------------------------------------------------------------===//
+
 /// Create a new top-level PassManager.
 MLIR_CAPI_EXPORTED MlirPassManager mlirPassManagerCreate(MlirContext ctx);
 
 /// Destroy the provided PassManager.
 MLIR_CAPI_EXPORTED void mlirPassManagerDestroy(MlirPassManager passManager);
 
 /// Checks if a PassManager is null.
 static inline bool mlirPassManagerIsNull(MlirPassManager passManager) {
MLIR_CAPI_EXPORTED void mlirPrintPassPipeline(MlirOpPassManager passManager,
                                               MlirStringCallback callback,
                                               void *userData);
 
 /// Parse a textual MLIR pass pipeline and add it to the provided OpPassManager.
 
 MLIR_CAPI_EXPORTED MlirLogicalResult
 mlirParsePassPipeline(MlirOpPassManager passManager, MlirStringRef pipeline);
 
+//===----------------------------------------------------------------------===//
+// External Pass API.
+//
+// This API allows to define passes outside of MLIR, not necessarily in
+// C++, and register them with the MLIR pass management infrastructure.
+//
+//===----------------------------------------------------------------------===//
+
+/// Structure of external `MlirPass` callbacks.
+/// All callbacks are required to be set unless otherwise specified.
+struct MlirExternalPassCallbacks {
+  /// This callback is called from the pass is created.
+  /// This is analogous to a C++ pass constructor.
+  void (*construct)(void *userData);
+
+  /// This callback is called when the pass is destroyed
+  /// This is analogous to a C++ pass destructor.
+  void (*destruct)(void *userData);
+
+  /// This callback is optional.
+  /// The callback is called before the pass is run, allowing a chance to
+  /// initialize any complex state necessary for running the pass.
+  /// See Pass::initialize(MLIRContext *).
+  MlirLogicalResult (*initialize)(MlirContext ctx, void *userData);
+
+  /// This callback is called when the pass is cloned.
+  /// See Pass::clonePass().
+  void *(*clone)(void *userData);
+
+  /// This callback is called when the pass is run.
+  /// See Pass::runOnOperation().
+  void (*run)(MlirOperation op, MlirExternalPass pass, void *userData);
+};
+typedef struct MlirExternalPassCallbacks MlirExternalPassCallbacks;
+
+/// Creates an external `MlirPass` that calls the supplied `callbacks` using the
+/// supplied `userData`. If `opName` is empty, the pass is a generic operation
+/// pass. Otherwise it is an operation pass specific to the specified pass name.
+MLIR_CAPI_EXPORTED MlirPass mlirCreateExternalPass(
+    MlirTypeID passID, MlirStringRef name, MlirStringRef argument,
+    MlirStringRef description, MlirStringRef opName,
+    intptr_t nDependentDialects, MlirDialectHandle *dependentDialects,
+    MlirExternalPassCallbacks callbacks, void *userData);
+
+/// This signals that the pass has failed. This is only valid to call during
+/// the `run` callback of `MlirExternalPassCallbacks`.
+/// See Pass::signalPassFailure().
+MLIR_CAPI_EXPORTED void mlirExternalPassSignalFailure(MlirExternalPass pass);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_PASS_H

mlir/include/mlir/CAPI/IR.h

 DEFINE_C_API_PTR_METHODS(MlirContext, mlir::MLIRContext)
 DEFINE_C_API_PTR_METHODS(MlirDialect, mlir::Dialect)
 DEFINE_C_API_PTR_METHODS(MlirDialectRegistry, mlir::DialectRegistry)
 DEFINE_C_API_PTR_METHODS(MlirOperation, mlir::Operation)
 DEFINE_C_API_PTR_METHODS(MlirBlock, mlir::Block)
 DEFINE_C_API_PTR_METHODS(MlirOpPrintingFlags, mlir::OpPrintingFlags)
 DEFINE_C_API_PTR_METHODS(MlirRegion, mlir::Region)
 DEFINE_C_API_PTR_METHODS(MlirSymbolTable, mlir::SymbolTable)
+DEFINE_C_API_PTR_METHODS(MlirTypeIDAllocator, mlir::TypeIDAllocator)
 
 DEFINE_C_API_METHODS(MlirAttribute, mlir::Attribute)
 DEFINE_C_API_METHODS(MlirIdentifier, mlir::StringAttr)
 DEFINE_C_API_METHODS(MlirLocation, mlir::Location)
 DEFINE_C_API_METHODS(MlirModule, mlir::ModuleOp)
 DEFINE_C_API_METHODS(MlirType, mlir::Type)
 DEFINE_C_API_METHODS(MlirTypeID, mlir::TypeID)
 DEFINE_C_API_METHODS(MlirValue, mlir::Value)
 
 #endif // MLIR_CAPI_IR_H

mlir/lib/CAPI/IR/IR.cpp

 MlirStringRef mlirIdentifierStr(MlirIdentifier ident) {
   return wrap(unwrap(ident).strref());
 }
 
 //===----------------------------------------------------------------------===//
 // TypeID API.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirTypeIDCreate(const void *ptr) {
+  assert(reinterpret_cast<uintptr_t>(ptr) % 8 == 0 &&
+         "ptr must be 8 byte aligned");
+  // This is essentially a no-op that returns back `ptr`, but by going through
+  // the `TypeID` functions we can get compiler errors in case the `TypeID`
+  // api/representation changes
+  return wrap(TypeID::getFromOpaquePointer(ptr));
+}
+
 bool mlirTypeIDEqual(MlirTypeID typeID1, MlirTypeID typeID2) {
   return unwrap(typeID1) == unwrap(typeID2);
 }
 
 size_t mlirTypeIDHashValue(MlirTypeID typeID) {
   return hash_value(unwrap(typeID));
 }
 
 //===----------------------------------------------------------------------===//
+// TypeIDAllocator API.
+//===----------------------------------------------------------------------===//
+
+MlirTypeIDAllocator mlirTypeIDAllocatorCreate() {
+  return wrap(new TypeIDAllocator());
+}
+
+void mlirTypeIDAllocatorDestroy(MlirTypeIDAllocator allocator) {
+  delete unwrap(allocator);
+}
+
+MlirTypeID mlirTypeIDAllocatorAllocateTypeID(MlirTypeIDAllocator allocator) {
+  return wrap(unwrap(allocator)->allocate());
+}
+
+//===----------------------------------------------------------------------===//
 // Symbol and SymbolTable API.
 //===----------------------------------------------------------------------===//
 
 MlirStringRef mlirSymbolTableGetSymbolAttributeName() {
   return wrap(SymbolTable::getSymbolAttrName());
 }
 
 MlirStringRef mlirSymbolTableGetVisibilityAttributeName() {

mlir/lib/CAPI/IR/Pass.cpp

 }
 
 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)));
 }
+
+//===----------------------------------------------------------------------===//
+// External Pass API.
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+class ExternalPass;
+} // namespace mlir
+DEFINE_C_API_PTR_METHODS(MlirExternalPass, mlir::ExternalPass)
+
+namespace mlir {
+/// This pass class wraps external passes defined in other languages using the
+/// MLIR C-interface
+class ExternalPass : public Pass {
+public:
+  ExternalPass(TypeID passID, StringRef name, StringRef argument,
+               StringRef description, Optional<StringRef> opName,
+               ArrayRef<MlirDialectHandle> dependentDialects,
+               MlirExternalPassCallbacks callbacks, void *userData)
+      : Pass(passID, opName), id(passID), name(name), argument(argument),
+        description(description), dependentDialects(dependentDialects),
+        callbacks(callbacks), userData(userData) {
+    callbacks.construct(userData);
+  }
+
+  ~ExternalPass() override { callbacks.destruct(userData); }
+
+  StringRef getName() const override { return name; }
+  StringRef getArgument() const override { return argument; }
+  StringRef getDescription() const override { return description; }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    MlirDialectRegistry cRegistry = wrap(&registry);
+    for (MlirDialectHandle dialect : dependentDialects)
+      mlirDialectHandleInsertDialect(dialect, cRegistry);
+  }
+
+  void signalPassFailure() { Pass::signalPassFailure(); }
+
+protected:
+  LogicalResult initialize(MLIRContext *ctx) override {
+    if (callbacks.initialize)
+      return unwrap(callbacks.initialize(wrap(ctx), userData));
+    return success();
+  }
+
+  bool canScheduleOn(RegisteredOperationName opName) const override {
+    if (Optional<StringRef> specifiedOpName = getOpName())
+      return opName.getStringRef() == specifiedOpName;
+    return true;
+  }
+
+  void runOnOperation() override {
+    callbacks.run(wrap(getOperation()), wrap(this), userData);
+  }
+
+  std::unique_ptr<Pass> clonePass() const override {
+    void *clonedUserData = callbacks.clone(userData);
+    return std::make_unique<ExternalPass>(id, name, argument, description,
+                                          getOpName(), dependentDialects,
+                                          callbacks, clonedUserData);
+  }
+
+private:
+  TypeID id;
+  std::string name;
+  std::string argument;
+  std::string description;
+  std::vector<MlirDialectHandle> dependentDialects;
+  MlirExternalPassCallbacks callbacks;
+  void *userData;
+};
+} // namespace mlir
+
+MlirPass mlirCreateExternalPass(MlirTypeID passID, MlirStringRef name,
+                                MlirStringRef argument,
+                                MlirStringRef description, MlirStringRef opName,
+                                intptr_t nDependentDialects,
+                                MlirDialectHandle *dependentDialects,
+                                MlirExternalPassCallbacks callbacks,
+                                void *userData) {
+  return wrap(static_cast<mlir::Pass *>(new mlir::ExternalPass(
+      unwrap(passID), unwrap(name), unwrap(argument), unwrap(description),
+      opName.length > 0 ? Optional<StringRef>(unwrap(opName)) : None,
+      {dependentDialects, static_cast<size_t>(nDependentDialects)}, callbacks,
+      userData)));
+}
+
+void mlirExternalPassSignalFailure(MlirExternalPass pass) {
+  unwrap(pass)->signalPassFailure();
+}

mlir/test/CAPI/CMakeLists.txt

LINK_LIBS PRIVATE
     MLIRCAPIIR
     MLIRCAPILLVM
     MLIRCAPIRegistration
 )
 
 _add_capi_test_executable(mlir-capi-pass-test
   pass.c
   LINK_LIBS PRIVATE
+    MLIRCAPIFunc
     MLIRCAPIIR
     MLIRCAPIRegistration
     MLIRCAPITransforms
 )
 
 _add_capi_test_executable(mlir-capi-sparse-tensor-test
   sparse_tensor.c
   LINK_LIBS PRIVATE

mlir/test/CAPI/pass.c

 //===- pass.c - Simple test of C APIs -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM
 // Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 /* RUN: mlir-capi-pass-test 2>&1 | FileCheck %s
  */
 
 #include "mlir-c/Pass.h"
+#include "mlir-c/Dialect/Func.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Registration.h"
 #include "mlir-c/Transforms.h"
 
 #include <assert.h>
 #include <math.h>
 #include <stdio.h>
 #include <stdlib.h>
void testParsePassPipeline() {
   MlirPassManager pm = mlirPassManagerCreate(ctx);
   // Try parse a pipeline.
   MlirLogicalResult status = mlirParsePassPipeline(
       mlirPassManagerGetAsOpPassManager(pm),
       mlirStringRefCreateFromCString("builtin.module(func.func(print-op-stats),"
                                      " func.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");
+    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("builtin.module(func.func(print-op-stats),"
                                      " func.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");
+    fprintf(stderr,
+            "Unexpected failure parsing pipeline after registering the pass\n");
     exit(EXIT_FAILURE);
   }
 
   // CHECK: Round-trip: builtin.module(func.func(print-op-stats),
   // func.func(print-op-stats))
   fprintf(stderr, "Round-trip: ");
   mlirPrintPassPipeline(mlirPassManagerGetAsOpPassManager(pm), printToStderr,
                         NULL);
   fprintf(stderr, "\n");
   mlirPassManagerDestroy(pm);
   mlirContextDestroy(ctx);
 }
 
+struct TestExternalPassUserData {
+  int constructCallCount;
+  int destructCallCount;
+  int initializeCallCount;
+  int cloneCallCount;
+  int runCallCount;
+};
+typedef struct TestExternalPassUserData TestExternalPassUserData;
+
+void testConstructExternalPass(void *userData) {
+  ++((TestExternalPassUserData *)userData)->constructCallCount;
+}
+
+void testDestructExternalPass(void *userData) {
+  ++((TestExternalPassUserData *)userData)->destructCallCount;
+}
+
+MlirLogicalResult testInitializeExternalPass(MlirContext ctx, void *userData) {
+  ++((TestExternalPassUserData *)userData)->initializeCallCount;
+  return mlirLogicalResultSuccess();
+}
+
+MlirLogicalResult testInitializeFailingExternalPass(MlirContext ctx,
+                                                    void *userData) {
+  ++((TestExternalPassUserData *)userData)->initializeCallCount;
+  return mlirLogicalResultFailure();
+}
+
+void *testCloneExternalPass(void *userData) {
+  ++((TestExternalPassUserData *)userData)->cloneCallCount;
+  return userData;
+}
+
+void testRunExternalPass(MlirOperation op, MlirExternalPass pass,
+                         void *userData) {
+  ++((TestExternalPassUserData *)userData)->runCallCount;
+}
+
+void testRunExternalFuncPass(MlirOperation op, MlirExternalPass pass,
+                             void *userData) {
+  ++((TestExternalPassUserData *)userData)->runCallCount;
+  MlirStringRef opName = mlirIdentifierStr(mlirOperationGetName(op));
+  if (!mlirStringRefEqual(opName,
+                          mlirStringRefCreateFromCString("func.func"))) {
+    mlirExternalPassSignalFailure(pass);
+  }
+}
+
+void testRunFailingExternalPass(MlirOperation op, MlirExternalPass pass,
+                                void *userData) {
+  ++((TestExternalPassUserData *)userData)->runCallCount;
+  mlirExternalPassSignalFailure(pass);
+}
+
+MlirExternalPassCallbacks makeTestExternalPassCallbacks(
+    MlirLogicalResult (*initializePass)(MlirContext ctx, void *userData),
+    void (*runPass)(MlirOperation op, MlirExternalPass, void *userData)) {
+  return (MlirExternalPassCallbacks){testConstructExternalPass,
+                                     testDestructExternalPass, initializePass,
+                                     testCloneExternalPass, runPass};
+}
+
+void testExternalPass() {
+  MlirContext ctx = mlirContextCreate();
+  mlirRegisterAllDialects(ctx);
+
+  MlirModule module = mlirModuleCreateParse(
+      ctx,
+      // clang-format off
+      mlirStringRefCreateFromCString(
+"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");
+    exit(EXIT_FAILURE);
+  }
+
+  MlirStringRef description = mlirStringRefCreateFromCString("");
+  MlirStringRef emptyOpName = mlirStringRefCreateFromCString("");
+
+  MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
+
+  // Run a generic pass
+  {
+    MlirTypeID passID = mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+    MlirStringRef name = mlirStringRefCreateFromCString("TestExternalPass");
+    MlirStringRef argument =
+        mlirStringRefCreateFromCString("test-external-pass");
+    TestExternalPassUserData userData = {0};
+
+    MlirPass externalPass = mlirCreateExternalPass(
+        passID, name, argument, description, emptyOpName, 0, NULL,
+        makeTestExternalPassCallbacks(NULL, testRunExternalPass), &userData);
+
+    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);
+    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);
+    }
+
+    mlirPassManagerDestroy(pm);
+
+    if (userData.destructCallCount != userData.constructCallCount) {
+      fprintf(stderr, "Expected destructCallCount to be equal to "
+                      "constructCallCount\n");
+      exit(EXIT_FAILURE);
+    }
+  }
+
+  // Run a func operation pass
+  {
+    MlirTypeID passID = mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+    MlirStringRef name = mlirStringRefCreateFromCString("TestExternalFuncPass");
+    MlirStringRef argument =
+        mlirStringRefCreateFromCString("test-external-func-pass");
+    TestExternalPassUserData userData = {0};
+    MlirDialectHandle funcHandle = mlirGetDialectHandle__func__();
+    MlirStringRef funcOpName = mlirStringRefCreateFromCString("func.func");
+
+    MlirPass externalPass = mlirCreateExternalPass(
+        passID, name, argument, description, funcOpName, 1, &funcHandle,
+        makeTestExternalPassCallbacks(NULL, testRunExternalFuncPass),
+        &userData);
+
+    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);
+    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");
+      exit(EXIT_FAILURE);
+    }
+
+    // The pass should only be run once this there is only one func op
+    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);
+    }
+  }
+
+  // Run a pass with `initialize` set
+  {
+    MlirTypeID passID = mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+    MlirStringRef name = mlirStringRefCreateFromCString("TestExternalPass");
+    MlirStringRef argument =
+        mlirStringRefCreateFromCString("test-external-pass");
+    TestExternalPassUserData userData = {0};
+
+    MlirPass externalPass = mlirCreateExternalPass(
+        passID, name, argument, description, emptyOpName, 0, NULL,
+        makeTestExternalPassCallbacks(testInitializeExternalPass,
+                                      testRunExternalPass),
+        &userData);
+
+    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);
+    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);
+    }
+
+    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);
+    }
+  }
+
+  // Run a pass that fails during `initialize`
+  {
+    MlirTypeID passID = mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+    MlirStringRef name =
+        mlirStringRefCreateFromCString("TestExternalFailingPass");
+    MlirStringRef argument =
+        mlirStringRefCreateFromCString("test-external-failing-pass");
+    TestExternalPassUserData userData = {0};
+
+    MlirPass externalPass = mlirCreateExternalPass(
+        passID, name, argument, description, emptyOpName, 0, NULL,
+        makeTestExternalPassCallbacks(testInitializeFailingExternalPass,
+                                      testRunExternalPass),
+        &userData);
+
+    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);
+    if (mlirLogicalResultIsSuccess(success)) {
+      fprintf(
+          stderr,
+          "Expected failure running pass manager on failing external pass.\n");
+      exit(EXIT_FAILURE);
+    }
+
+    if (userData.initializeCallCount != 1) {
+      fprintf(stderr, "Expected initializeCallCount to be 1\n");
+      exit(EXIT_FAILURE);
+    }
+
+    if (userData.runCallCount != 0) {
+      fprintf(stderr, "Expected runCallCount to be 0\n");
+      exit(EXIT_FAILURE);
+    }
+
+    mlirPassManagerDestroy(pm);
+
+    if (userData.destructCallCount != userData.constructCallCount) {
+      fprintf(stderr, "Expected destructCallCount to be equal to "
+                      "constructCallCount\n");
+      exit(EXIT_FAILURE);
+    }
+  }
+
+  // Run a pass that fails during `run`
+  {
+    MlirTypeID passID = mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+    MlirStringRef name =
+        mlirStringRefCreateFromCString("TestExternalFailingPass");
+    MlirStringRef argument =
+        mlirStringRefCreateFromCString("test-external-failing-pass");
+    TestExternalPassUserData userData = {0};
+
+    MlirPass externalPass = mlirCreateExternalPass(
+        passID, name, argument, description, emptyOpName, 0, NULL,
+        makeTestExternalPassCallbacks(NULL, testRunFailingExternalPass),
+        &userData);
+
+    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);
+    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);
+  mlirContextDestroy(ctx);
+}
+
 int main() {
   testRunPassOnModule();
   testRunPassOnNestedModule();
   testPrintPassPipeline();
   testParsePassPipeline();
+  testExternalPass();
   return 0;
 }