Next steps on python bindings

MLIR
#1

I’m forking comments from the patch D85481 by @zhanghb97 to this thread to have a higher level discussion about next steps.

@ftynse and I were discussing over chat this morning/evening and were thinking that we could reduce ambiguity on the next step by focusing it a bit more tightly. The patch covers a number of core IR constructs that have a fair bit of subtlety to them, and we would like to take a more incremental/design driven approach to them.

We were wondering if it might make sense to scope the first patch down to a few things:

  • Binding for MlirContext
  • Binding for MlirModule without a constructor (avoids having to bind location)
  • MlirContext.parse() method that returns an MlirModule
  • MlirModule.dump() method

That would be enough for a first patch, and it defers a lot of the more complex issues. A good second and third patch could be (note that there may be some other opinions on naming and details):

  • Add a C-API for mlirModulePrint(...) that allows to get the printed asm form.
  • Add a python binding for MlirModule.print(file) which prints ASM to an open file handle and an MlirModule.print_to_string() that is a shortcut for returning a string directly. As a very first step, you could just provide the print_to_string() method.

I would recommend to have a simple mlirModulePrint() C-API that does not take a mapping of OpPrintingFlags as a first step, and then provide a second mlirModulePrintWithOptions(MlirOpPrintingOptions) as a followup. Leaving out a C-API mapping of the printing flags will simplify things, and there is ambiguity in getting the OpPrintingOptions mapped well in a way that maximizes version compatibility.

Further, I would recommend modeling the mlirModulePrint() function to take a callback and a cookie for appending to the string. Example:

// Accumulates string contents represented by data and len with a caller supplied cookie.
// Returns 0 (false) on IO failure and non-zero on success.
typedef void (*MlirAppendStringCallback)(void *cookie, const char *data, size_t len));

// Prints the ASM form of the module to the given callback with a user-supplied cookie.
void mlirModulePrint(MlirAppendStringCallback callback, void *cookie);

Note that there are multiple ways to model a print function like this, and this is the form that I would choose/advocate for as the most suitable for this kind of API.

I think that getting the python bindings bootstrapped with the ability to create a context and parse/dump/print the ASM form could be a nice thing to build further on. From a work sequencing perspective and given the state we have on it, either Alex or I would like to drive some of the subtler issues of mapping the Operation hierarchy ourselves within the next week or so and were hoping that scoping down your patch a bit to just the basics would open us up to be able to work better in parallel (i.e. after you submit the scaled down first patch above).

[RFC] First step Python bindings for C API
#2

I’m scoping down my patch to meet the first requirement, but I find that as for the MlirModule:

/*============================================================================*/
/* Module API.                                                                */
/*============================================================================*/

/** Creates a new, empty module and transfers ownership to the caller. */
MlirModule mlirModuleCreateEmpty(MlirLocation location);

/** Parses a module from the string and transfers ownership to the caller. */
MlirModule mlirModuleCreateParse(MlirContext context, const char *module);

/** Takes a module owned by the caller and deletes it. */
void mlirModuleDestroy(MlirModule module);

/** Views the module as a generic operation. */
MlirOperation mlirModuleGetOperation(MlirModule module);

there is no dump related function in the C API. So how to implement the MlirModule.dump() method here?

#3

I’d start with this, and then implement the python binding for module.dump()

#4

There is a dump method for operation.

MlirOperation mlirModuleGetOperation(MlirModule module);
...
void mlirOperationDump(MlirOperation op);
#5

I think it’s not make sense to bind MlirModule.dump() with mlirOperationDump. I tried to add a C API:

void mlirModuleDump(MlirModule module) {
  unwrap(module).dump();
}

And bind the mlirModuleDump as the MlirModule.dump():

py::class_<MlirModule>(m, "MlirModule")
      .def("dump", &mlirModuleDump);
#6

Actually mlirOperationDump is perfect: a Module is an Operation, and it can be converted using mlirModuleGetOperation which gives access to every Operation-related C APIs.

#7

Ah, I get it. However, for the MlirModule.dump() method, is it a more direct choice to add mlirModuleDump C API that wraps the ModuleOp::dump()? Or is there any reason I didn’t notice to avoid adding the mlirModuleDump C API?

#8

Do we want to connect this to __str__ or __repr__ directly?

Just confirming the intention.

I have been discussing this with @jingpu as well. I’d really like to avoid double allocation for the string here. Currently, IR objects can be printed to llvm::raw_ostream. There’s no subclass that would work with C strings AFAICS, so the naive solution would be to print to llvm::raw_string_ostream and then either return a copy and transfer ownership to the caller, or give the caller a non-owning view letting them copy themselves. We cannot just transfer the ownership of the data owned by the string to the API caller because they won’t be able to properly deallocate it. Thinking about the API you propose, we may implement our own subclass of llvm::raw_ostream that forwards chunks of the string to the user-provided callback. It’s then up to the caller to make allocation efficient (it will likely get called multiple times for long strings), but all memory management will happen on the caller side so no duplication except for buffering.

You don’t need to add anything to the C API. Like we already told you, Pybind lets you call any function from the python API, including a function declared in the bindings source file or a lambda. I expect the following to work.

py::class_<MlirModule>(m, "MlirModule")
      .def("dump", [](MlirModule module) {
    mlirOperationDump(mlirModuleGetOperation(module));
});
#9

+1 to __str__. I’ve had mixed results with connecting asm generation directly to the __repr__ method. That one is used in a lot of debug print contexts that really need a compact representation of the actual object.

We could still retain a dedicated print method for the version with all options. We might want to name it something else though (diverging from the MLIR naming): print typically meant to output to the console in some way in python. Perhaps a to_asm for the generic form. Anyway, that can be bikeshed on another patch.

That is exactly what I had in mind. This falls under my rule of “friends don’t let friends allocate their strings” when crossing a boundary like this. My approach also interops with whatever file handle like thing the caller shows up with.

In past systems, I’ve eventually had to have typedefs for generic read/write callbacks like this, but to be general, they should also return an error indication that can be optionally used by the callee to know whether the stream is in an error state (possibly ending it’s operation prematurely). We probably want to add an int/bool return to the callback to future proof it.

#10

I have scoped down my patch D85481 to meet the first patch requirements:

According to the previous feedback, my patch uses a PyMlirContext class to wrap the MlirContext so that the context object can be created and destroyed correctly:

class PyMlirContext {
public:
  PyMlirContext() { context = mlirContextCreate(); }
  ~PyMlirContext() { mlirContextDestroy(context); }
  // Parses the module from asm.
  MlirModule parse(const std::string &module) {
    return mlirModuleCreateParse(context, module.c_str());
  }

  MlirContext context;
};

As for the module, it can be parsed from the ASM, and the dump method is implemented with the mlirModuleGetOperation and mlirOperationDump.

So far, I have trouble of linking the libraries separately. As the previous feedback showed, multiple files should be linked separately. When I tried like this:

add_library(MLIRIRModulesBindingsPythonExtension ${PYEXT_LINK_MODE}
  IRModules.cpp
)
... ...
target_link_libraries(MLIRIRModulesBindingsPythonExtension
  PRIVATE
  MLIRIR
  MLIRCAPIIR
  MLIRCAPIRegistration
  ${PYEXT_LIBADD}
)

add_library(MLIRBindingsPythonExtension ${PYEXT_LINK_MODE}
  MainModule.cpp
)
... ...
target_link_libraries(MLIRBindingsPythonExtension
  PRIVATE
  MLIRIR
  MLIRIRModulesBindingsPythonExtension
  ${PYEXT_LIBADD}
)

but it reported:

MainModule.cpp:(.text.PyInit__mlir+0x5f0): undefined reference to `populateIRSubmodule(pybind11::module&)'

I don’t know why the MainModule can’t link the IRModules library?