Switch the CallbackOstream wrapper in the MLIR C API to an Unbuffered stream

[mehdi_amini]

Summary

This delegate the control of the buffering to the user of the API. This
seems like a safer option as messages are immediately propagated to the
user, which may lead to less surprising behavior during debugging for
instance.
In terms of performance, a user can add a buffered stream on the other
side of the callback.

[rriddle]

mlir/include/mlir/CAPI/Utils.h

31

nit: Seems like there is an extra space in the comment here.

[stellaraccident]

Good catch.

mlir/include/mlir/CAPI/Utils.h

 namespace mlir {
 namespace detail {
 /// A simple raw ostream subclass that forwards write_impl calls to the
 /// user-supplied callback together with opaque user-supplied data.
 class CallbackOstream : public llvm::raw_ostream {
 public:
   CallbackOstream(std::function<void(const char *, intptr_t, void *)> callback,
                   void *opaqueData)
-      : callback(callback), opaqueData(opaqueData), pos(0u) {}
+      : raw_ostream(/*unbuffered =*/true), callback(callback),
+        opaqueData(opaqueData), pos(0u) {}
 
   void write_impl(const char *ptr, size_t size) override {
     callback(ptr, size, opaqueData);
     pos += size;
   }
 
   uint64_t current_pos() const override { return pos; }
 

mlir/lib/CAPI/IR/AffineExpr.cpp

 MlirContext mlirAffineExprGetContext(MlirAffineExpr affineExpr) {
   return wrap(unwrap(affineExpr).getContext());
 }
 
 void mlirAffineExprPrint(MlirAffineExpr affineExpr, MlirStringCallback callback,
                          void *userData) {
   mlir::detail::CallbackOstream stream(callback, userData);
   unwrap(affineExpr).print(stream);
-  stream.flush();
 }
 
 void mlirAffineExprDump(MlirAffineExpr affineExpr) {
   unwrap(affineExpr).dump();
 }
 
 int mlirAffineExprIsSymbolicOrConstant(MlirAffineExpr affineExpr) {
   return unwrap(affineExpr).isSymbolicOrConstant();

mlir/lib/CAPI/IR/AffineMap.cpp

 int mlirAffineMapEqual(MlirAffineMap a1, MlirAffineMap a2) {
   return unwrap(a1) == unwrap(a2);
 }
 
 void mlirAffineMapPrint(MlirAffineMap affineMap, MlirStringCallback callback,
                         void *userData) {
   mlir::detail::CallbackOstream stream(callback, userData);
   unwrap(affineMap).print(stream);
-  stream.flush();
 }
 
 void mlirAffineMapDump(MlirAffineMap affineMap) { unwrap(affineMap).dump(); }
 
 MlirAffineMap mlirAffineMapEmptyGet(MlirContext ctx) {
   return wrap(AffineMap::get(unwrap(ctx)));
 }
 

mlir/lib/CAPI/IR/Diagnostics.cpp

 #include "mlir/IR/Diagnostics.h"
 
 using namespace mlir;
 
 void mlirDiagnosticPrint(MlirDiagnostic diagnostic, MlirStringCallback callback,
                          void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(diagnostic).print(stream);
-  stream.flush();
 }
 
 MlirLocation mlirDiagnosticGetLocation(MlirDiagnostic diagnostic) {
   return wrap(unwrap(diagnostic).getLocation());
 }
 
 MlirDiagnosticSeverity mlirDiagnosticGetSeverity(MlirDiagnostic diagnostic) {
   switch (unwrap(diagnostic).getSeverity()) {

mlir/lib/CAPI/IR/IR.cpp

 MlirContext mlirLocationGetContext(MlirLocation location) {
   return wrap(unwrap(location).getContext());
 }
 
 void mlirLocationPrint(MlirLocation location, MlirStringCallback callback,
                        void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(location).print(stream);
-  stream.flush();
 }
 
 /* ========================================================================== */
 /* Module API.                                                                */
 /* ========================================================================== */
 
 MlirModule mlirModuleCreateEmpty(MlirLocation location) {
   return wrap(ModuleOp::create(unwrap(location)));
int mlirOperationRemoveAttributeByName(MlirOperation op, const char *name) {
   auto removeResult = unwrap(op)->removeAttr(name);
   return removeResult == MutableDictionaryAttr::RemoveResult::Removed;
 }
 
 void mlirOperationPrint(MlirOperation op, MlirStringCallback callback,
                         void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(op)->print(stream);
-  stream.flush();
 }
 
 void mlirOperationPrintWithFlags(MlirOperation op, MlirOpPrintingFlags flags,
                                  MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(op)->print(stream, *unwrap(flags));
-  stream.flush();
 }
 
 void mlirOperationDump(MlirOperation op) { return unwrap(op)->dump(); }
 
 /* ========================================================================== */
 /* Region API.                                                                */
 /* ========================================================================== */
 

 MlirValue mlirBlockGetArgument(MlirBlock block, intptr_t pos) {
   return wrap(unwrap(block)->getArgument(static_cast<unsigned>(pos)));
 }
 
 void mlirBlockPrint(MlirBlock block, MlirStringCallback callback,
                     void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(block)->print(stream);
-  stream.flush();
 }
 
 /* ========================================================================== */
 /* Value API.                                                                 */
 /* ========================================================================== */
 
 int mlirValueIsABlockArgument(MlirValue value) {
   return unwrap(value).isa<BlockArgument>();

 }
 
 void mlirValueDump(MlirValue value) { unwrap(value).dump(); }
 
 void mlirValuePrint(MlirValue value, MlirStringCallback callback,
                     void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(value).print(stream);
-  stream.flush();
 }
 
 /* ========================================================================== */
 /* Type API.                                                                  */
 /* ========================================================================== */
 
 MlirType mlirTypeParseGet(MlirContext context, const char *type) {
   return wrap(mlir::parseType(type, unwrap(context)));
 }
 
 MlirContext mlirTypeGetContext(MlirType type) {
   return wrap(unwrap(type).getContext());
 }
 
 int mlirTypeEqual(MlirType t1, MlirType t2) { return unwrap(t1) == unwrap(t2); }
 
 void mlirTypePrint(MlirType type, MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(type).print(stream);
-  stream.flush();
 }
 
 void mlirTypeDump(MlirType type) { unwrap(type).dump(); }
 
 /* ========================================================================== */
 /* Attribute API.                                                             */
 /* ========================================================================== */
 

 int mlirAttributeEqual(MlirAttribute a1, MlirAttribute a2) {
   return unwrap(a1) == unwrap(a2);
 }
 
 void mlirAttributePrint(MlirAttribute attr, MlirStringCallback callback,
                         void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(attr).print(stream);
-  stream.flush();
 }
 
 void mlirAttributeDump(MlirAttribute attr) { unwrap(attr).dump(); }
 
 MlirNamedAttribute mlirNamedAttributeGet(const char *name, MlirAttribute attr) {
   return MlirNamedAttribute{name, attr};
 }