Add userData to the diagnostic handler C API

[GeorgeLyon]

Summary

Previously, there was no way to add context to the diagnostic engine via the C API

[mehdi_amini]

mlir/include/mlir-c/Diagnostics.h

83

Please document the userData here and above

[GeorgeLyon]

Added deleteUserData callback

[GeorgeLyon]

@mehdi_amini I added one more bit of functionality here, so you may want to take another look.

[mehdi_amini]

Can you expand a bit on the use-case for deleteUserData? When is the user not in control of the lifetime? This seems all kind of tied to the context in the end?

[GeorgeLyon]

If a host language is registering dialect handlers that may need to be cleaned up (for instance, via a reference count decrement), it theoretically _could_ keep a list of attached handlers and manually release them on detach or when the context is destroyed. deleteUserData is a simple way to eliminate this bit of bookkeeping and make "ownership" a bit more formal in this case, namely the context is responsible for telling you when it is done using the handler (and thus, the userData). This way, if we hypothetically add a setHandlers method which takes an array of handles and _replaces_ the current handler stack, then the host language wouldn't need to account for this, deleteUserData would just be called when the handler is no longer in use.

[mehdi_amini]

@stellaraccident / @ftynse can you double check on the deleteUserData part and the shared_ptr?

[ftynse]

Shared pointer logic looks fine: the lambda takes it by-copy and is stored internally in std::function, so the data remains live.

I am not entirely convinced that we want to burden the C API with additional memory management facilities though. It's C after all, so I don't see a problem with requiring the caller to manage their data. MLIR will not remove the handler by itself. So the user of C API can be reasonably expected to clean their data if they remove the handler, or after calling mlirContextFree.

mlir/lib/CAPI/IR/Diagnostics.cpp

60

Nit: elide trivial braces

mlir/test/CAPI/ir.c

424–425

This change is irrelevant and weakens the test (CHECK only looks at the current line of the comment).

915–922

This change is irrelevant. Consider using git clang-format that would only touch the diff.

[GeorgeLyon]

Address feedback

[GeorgeLyon]

I have worked with many C APIs in a number of languages and the attitude of "just do the memory management bookkeeping in the host language" has created some very unfortunate complexity in the bindings and more than a few bugs. This is a very low-cost change for the C bindings, and doesn't add much complexity for language bindings that don't care about this functionality (they can just pass NULL to the extra arguments). I think this is also a win from a self-documentation standpoint: before, I would need to go and figure out what are the different ways diagnostic handlers can be unregistered to understand the semantics. I would also need to worry about concurrency, specifically whether or not my handler is guaranteed to not be called after detach returns. I could read the code to figure this out (it turns out, the answer is "no") but I think it is better to minimize this sort of spelunking and just say "we will call your handler and when we are done calling your handler we will call your callback".

Also, if I remember correctly the // CHECK lines used to be surrounded by // clang-format off/on. Is there a reason this changed? I used to run clang-format on save in my editor and it doesn't look like git clang-format has support for this type of flow.

[mehdi_amini]

LGTM

[ftynse]

In D91738#2408526, @GeorgeLyon wrote:

I have worked with many C APIs in a number of languages and the attitude of "just do the memory management bookkeeping in the host language" has created some very unfortunate complexity in the bindings and more than a few bugs. This is a very low-cost change for the C bindings, and doesn't add much complexity for language bindings that don't care about this functionality (they can just pass NULL to the extra arguments). I think this is also a win from a self-documentation standpoint: before, I would need to go and figure out what are the different ways diagnostic handlers can be unregistered to understand the semantics. I would also need to worry about concurrency, specifically whether or not my handler is guaranteed to not be called after detach returns. I could read the code to figure this out (it turns out, the answer is "no") but I think it is better to minimize this sort of spelunking and just say "we will call your handler and when we are done calling your handler we will call your callback".

Works for me as long as you put some version of this argument in the commit message.

In D91738#2408526, @GeorgeLyon wrote:

Also, if I remember correctly the // CHECK lines used to be surrounded by // clang-format off/on. Is there a reason this changed? I used to run clang-format on save in my editor and it doesn't look like git clang-format has support for this type of flow.

clang-format used to mess up CHECK comments, but does not anymore, at least for me. So I assumed it has been taught about FileCheck lines.

mlir/include/mlir-c/Diagnostics.h

enum MlirDiagnosticSeverity {
   MlirDiagnosticNote,
   MlirDiagnosticRemark
 };
 typedef enum MlirDiagnosticSeverity MlirDiagnosticSeverity;
 
 /// Opaque identifier of a diagnostic handler, useful to detach a handler.
 typedef uint64_t MlirDiagnosticHandlerID;
 
-/** Diagnostic handler type. Acceps a reference to a diagnostic, which is only
- * guaranteed to be live during the call. If the handler processed the
- * diagnostic completely, it is expected to return success. Otherwise, it is
- * expected to return failure to indicate that other handlers should attempt to
- * process the diagnostic. */
-typedef MlirLogicalResult (*MlirDiagnosticHandler)(MlirDiagnostic);
+/** Diagnostic handler type. Accepts a reference to a diagnostic, which is only
+ * guaranteed to be live during the call. The handler is passed the `userData`
+ * that was provided when the handler was attached to a context. If the handler
+ * processed the diagnostic completely, it is expected to return success.
+ * Otherwise, it is expected to return failure to indicate that other handlers
+ * should attempt to process the diagnostic. */
+typedef MlirLogicalResult (*MlirDiagnosticHandler)(MlirDiagnostic,
+                                                   void *userData);
 
 /// Prints a diagnostic using the provided callback.
 MLIR_CAPI_EXPORTED void mlirDiagnosticPrint(MlirDiagnostic diagnostic,
                                             MlirStringCallback callback,
                                             void *userData);
 
 /// Returns the location at which the diagnostic is reported.
 MLIR_CAPI_EXPORTED MlirLocation

 
 /** Returns `pos`-th note attached to the diagnostic. Expects `pos` to be a
  * valid zero-based index into the list of notes. */
 MLIR_CAPI_EXPORTED MlirDiagnostic
 mlirDiagnosticGetNote(MlirDiagnostic diagnostic, intptr_t pos);
 
 /** Attaches the diagnostic handler to the context. Handlers are invoked in the
  * reverse order of attachment until one of them processes the diagnostic
- * completely. Returns an identifier that can be used to detach the handler. */
+ * completely. When a handler is invoked it is passed the `userData` that was
+ * provided when it was attached. If non-NULL, `deleteUserData` is called once
+ * the system no longer needs to call the handler (for instance after the
+ * handler is detached or the context is destroyed). Returns an identifier that
+ * can be used to detach the handler.
+ */
 MLIR_CAPI_EXPORTED MlirDiagnosticHandlerID mlirContextAttachDiagnosticHandler(
-    MlirContext context, MlirDiagnosticHandler handler);
+    MlirContext context, MlirDiagnosticHandler handler, void *userData,
+    void (*deleteUserData)(void *));
 
 /** Detaches an attached diagnostic handler from the context given its
  * identifier. */
 MLIR_CAPI_EXPORTED void
 mlirContextDetachDiagnosticHandler(MlirContext context,
                                    MlirDiagnosticHandlerID id);
 
 /** Emits an error at the given location through the diagnostics engine. Used

mlir/lib/CAPI/IR/Diagnostics.cpp

 }
 
 // Notes are stored in a vector, so the iterator is a random access iterator,
 // cheap to advance multiple steps at a time.
 MlirDiagnostic mlirDiagnosticGetNote(MlirDiagnostic diagnostic, intptr_t pos) {
   return wrap(*std::next(unwrap(diagnostic).getNotes().begin(), pos));
 }
 
-MlirDiagnosticHandlerID
-mlirContextAttachDiagnosticHandler(MlirContext context,
-                                   MlirDiagnosticHandler handler) {
+static void deleteUserDataNoop(void *userData) {}
+
+MlirDiagnosticHandlerID mlirContextAttachDiagnosticHandler(
+    MlirContext context, MlirDiagnosticHandler handler, void *userData,
+    void (*deleteUserData)(void *)) {
   assert(handler && "unexpected null diagnostic handler");
+  if (deleteUserData == NULL)
+    deleteUserData = deleteUserDataNoop;
+  std::shared_ptr<void> sharedUserData(userData, deleteUserData);
   DiagnosticEngine::HandlerID id =
       unwrap(context)->getDiagEngine().registerHandler(
-          [handler](Diagnostic &diagnostic) {
-            return unwrap(handler(wrap(diagnostic)));
+          [handler, sharedUserData](Diagnostic &diagnostic) {
+            return unwrap(handler(wrap(diagnostic), sharedUserData.get()));
           });
   return static_cast<MlirDiagnosticHandlerID>(id);
 }
 
 void mlirContextDetachDiagnosticHandler(MlirContext context,
                                         MlirDiagnosticHandlerID id) {
   unwrap(context)->getDiagEngine().eraseHandler(
       static_cast<DiagnosticEngine::HandlerID>(id));
 }
 
 void mlirEmitError(MlirLocation location, const char *message) {
   emitError(unwrap(location)) << message;
 }

mlir/test/CAPI/ir.c

static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   mlirOpPrintingFlagsElideLargeElementsAttrs(flags, 2);
   mlirOpPrintingFlagsPrintGenericOpForm(flags);
   mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/0);
   mlirOpPrintingFlagsUseLocalScope(flags);
   fprintf(stderr, "Op print with all flags: ");
   mlirOperationPrintWithFlags(operation, flags, printToStderr, NULL);
   fprintf(stderr, "\n");
   // CHECK: Op print with all flags: %{{.*}} = "std.constant"() {elts = opaque<"", "0xDEADBEEF"> : tensor<4xi32>, value = 0 : index} : () -> index loc(unknown)
 
   mlirOpPrintingFlagsDestroy(flags);
 }
 
 static int constructAndTraverseIr(MlirContext ctx) {
   MlirLocation location = mlirLocationUnknownGet(ctx);
 
   MlirModule moduleOp = makeAndDumpAdd(ctx, location);
   MlirOperation module = mlirModuleGetOperation(moduleOp);
uint32_t *uint32RawData =
       (uint32_t *)mlirDenseElementsAttrGetRawData(uint32Elements);
   int32_t *int32RawData =
       (int32_t *)mlirDenseElementsAttrGetRawData(int32Elements);
   uint64_t *uint64RawData =
       (uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements);
   int64_t *int64RawData =
       (int64_t *)mlirDenseElementsAttrGetRawData(int64Elements);
   float *floatRawData =
       (float *)mlirDenseElementsAttrGetRawData(floatElements);
   double *doubleRawData =
       (double *)mlirDenseElementsAttrGetRawData(doubleElements);
   if (uint32RawData[0] != 0u || uint32RawData[1] != 1u ||
       int32RawData[0] != 0 || int32RawData[1] != 1 ||
       uint64RawData[0] != 0u || uint64RawData[1] != 1u ||
       int64RawData[0] != 0 || int64RawData[1] != 1 ||
       floatRawData[0] != 0.0f || floatRawData[1] != 1.0f ||
       doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0)
     return 18;
 
   mlirAttributeDump(splatBool);
   mlirAttributeDump(splatUInt32);
   mlirAttributeDump(splatInt32);
   mlirAttributeDump(splatUInt64);
   mlirAttributeDump(splatInt64);
int registerOnlyStd() {
 
   fprintf(stderr, "@registration\n");
   // CHECK-LABEL: @registration
 
   return 0;
 }
 
 // Wraps a diagnostic into additional text we can match against.
-MlirLogicalResult errorHandler(MlirDiagnostic diagnostic) {
-  fprintf(stderr, "processing diagnostic <<\n");
+MlirLogicalResult errorHandler(MlirDiagnostic diagnostic, void *userData) {
+  fprintf(stderr, "processing diagnostic (userData: %d) <<\n", (int)userData);
   mlirDiagnosticPrint(diagnostic, printToStderr, NULL);
   fprintf(stderr, "\n");
   MlirLocation loc = mlirDiagnosticGetLocation(diagnostic);
   mlirLocationPrint(loc, printToStderr, NULL);
   assert(mlirDiagnosticGetNumNotes(diagnostic) == 0);
-  fprintf(stderr, ">> end of diagnostic\n");
+  fprintf(stderr, ">> end of diagnostic (userData: %d)\n", (int)userData);
   return mlirLogicalResultSuccess();
 }
 
+// Logs when the delete user data callback is called
+static void deleteUserData(void *userData) {
+  fprintf(stderr, "deleting user data (userData: %d)\n", (int)userData);
+}
+
 void testDiagnostics() {
   MlirContext ctx = mlirContextCreate();
-  MlirDiagnosticHandlerID id =
-      mlirContextAttachDiagnosticHandler(ctx, errorHandler);
+  MlirDiagnosticHandlerID id = mlirContextAttachDiagnosticHandler(
+      ctx, errorHandler, (void *)42, deleteUserData);
   MlirLocation loc = mlirLocationUnknownGet(ctx);
   fprintf(stderr, "@test_diagnostics\n");
   mlirEmitError(loc, "test diagnostics");
   mlirContextDetachDiagnosticHandler(ctx, id);
   mlirEmitError(loc, "more test diagnostics");
   // CHECK-LABEL: @test_diagnostics
-  // CHECK: processing diagnostic <<
+  // CHECK: processing diagnostic (userData: 42) <<
   // CHECK:   test diagnostics
   // CHECK:   loc(unknown)
-  // CHECK: >> end of diagnostic
+  // CHECK: >> end of diagnostic (userData: 42)
+  // CHECK: deleting user data (userData: 42)
   // CHECK-NOT: processing diagnostic
   // CHECK:     more test diagnostics
 }
 
 int main() {
   MlirContext ctx = mlirContextCreate();
   mlirRegisterAllDialects(ctx);
   if (constructAndTraverseIr(ctx))