[mlir][python] Add generic operation parse APIs

[rkayaith]

Summary

Currently the bindings only allow for parsing IR with a top-level
builtin.module op, since the parse APIs insert an implicit module op.
This change adds Operation.parse, which returns whatever top-level op
is actually in the source.

To simplify parsing of specific operations, OpView.parse is also
added, which handles the error checking for OpView subclasses.

[rkayaith]

fix comment

[rkayaith]

rebase

[mehdi_amini]

mlir/lib/Bindings/Python/IRCore.cpp

1069

What's missing from mlir/include/mlir-c/Diagnostics.h?

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

1069

I'm not sure exactly what this comment refers to, but it's on all the other parse methods. Though, it would be nice if by default these methods collected errors and included them in the exception message, right now the user has to set up their own diagnostic handler.

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

1069

It should be possible now, just nobody updated those methods after the diagnostic handling was exposed to C API. It would be great not to cargo-cult the comment and actually try implementing this.

1071

Nit: it has been possible to do throw py::value_error for quite a while now, https://pybind11.readthedocs.io/en/stable/advanced/exceptions.html, let's not fallback to C APIs.

[rkayaith]

use py::value_error, and make extra arguments keyword-only in case more need to be added

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

1069

I'll look into fixing these when I have some time, I had an initial question that I left here: https://github.com/llvm/llvm-project/issues/60595

[rkayaith]

update bitcode -> bytecode in comments

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

1069

Any concerns with leaving this for a separate patch? I'd want to update all the other cases at the same time anyways, so the behaviour is consistent.

[rkayaith]

friendly ping

[stellaraccident]

Python changes lgtm. The parser changed make sense to but I usually defer those reviews to river.

[rriddle]

mlir/include/mlir/Parser/Parser.h

143–144

The doc for sourceName doesn't read well. This looks like it is treated as the filename of the source buffer, can we better document this?

[rkayaith]

improve docs, put sourceName arg before sourceFileLoc since it's common to
both overloads

mlir/include/mlir/Parser/Parser.h

143–144

I added a bit more info here

mlir/include/mlir-c/IR.h

 /// not be further used. Particularly, this applies to any regions added to
 /// the state (the implementation may invalidate any such pointers).
 ///
 /// This call can fail under the following conditions, in which case, it will
 /// return a null operation and emit diagnostics:
 ///   - Result type inference is enabled and cannot be performed.
 MLIR_CAPI_EXPORTED MlirOperation mlirOperationCreate(MlirOperationState *state);
 
+/// Parses an operation, giving ownership to the caller. If parsing fails a null
+/// operation will be returned, and an error diagnostic emitted.
+///
+/// `sourceStr` may be either the text assembly format, or binary bytecode
+/// format. `sourceName` is used when filling in missing location information.
+MLIR_CAPI_EXPORTED MlirOperation mlirOperationCreateParse(
+    MlirContext context, MlirStringRef sourceStr, MlirStringRef sourceName);
+
 /// Creates a deep copy of an operation. The operation is not inserted and
 /// ownership is transferred to the caller.
 MLIR_CAPI_EXPORTED MlirOperation mlirOperationClone(MlirOperation op);
 
 /// Takes an operation owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirOperationDestroy(MlirOperation op);
 
 /// Removes the given operation from its parent block. The operation is not

mlir/include/mlir/Parser/Parser.h

LogicalResult parseSourceFile(llvm::StringRef filename,
                               LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the IR string and appends parsed operations to the given block.
 /// If the block is non-empty, the operations are placed before the current
 /// terminator. If parsing is successful, success is returned. Otherwise, an
 /// error message is emitted through the error handler registered in the
 /// context, and failure is returned. If `sourceFileLoc` is non-null, it is
 /// populated with a file location representing the start of the source file
-/// that is being parsed.
+/// that is being parsed. `sourceName` is used when filling in missing location
+/// information.
 LogicalResult parseSourceString(llvm::StringRef sourceStr, Block *block,
                                 const ParserConfig &config,
-                                LocationAttr *sourceFileLoc = nullptr);
+                                LocationAttr *sourceFileLoc = nullptr,
+                                StringRef sourceName = "");
 
 namespace detail {
 /// The internal implementation of the templated `parseSourceFile` methods
 /// below, that simply forwards to the non-templated version.
 template <typename ContainerOpT, typename... ParserArgs>
 inline OwningOpRef<ContainerOpT> parseSourceFile(const ParserConfig &config,
                                                  ParserArgs &&...args) {
   LocationAttr sourceFileLoc;

 
 /// This parses the provided string containing MLIR. If the source IR contained
 /// a single instance of `ContainerOpT`, it is returned. Otherwise, a new
 /// instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the
-/// `SingleBlockImplicitTerminator` trait.
+/// `SingleBlockImplicitTerminator` trait. `sourceName` is used when filling in
+/// missing location information.
 template <typename ContainerOpT = Operation *>
 inline OwningOpRef<ContainerOpT> parseSourceString(llvm::StringRef sourceStr,
-                                                   const ParserConfig &config) {
+                                                   const ParserConfig &config,
+                                                   StringRef sourceName = "") {
   LocationAttr sourceFileLoc;
   Block block;
-  if (failed(parseSourceString(sourceStr, &block, config, &sourceFileLoc)))
+  if (failed(parseSourceString(sourceStr, &block, config, &sourceFileLoc,
+                               sourceName)))
     return OwningOpRef<ContainerOpT>();
   return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
       &block, config.getContext(), sourceFileLoc);
 }
 
 } // namespace mlir
 
 #endif // MLIR_PARSER_PARSER_H

mlir/lib/Bindings/Python/IRCore.cpp

 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/IR.h"
 //#include "mlir-c/Registration.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
-#include <utility>
 #include <optional>
+#include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
   (void)liveOperations;
 
   PyOperationRef created = createInstance(std::move(contextRef), operation,
                                           std::move(parentKeepAlive));
   created->attached = false;
   return created;
 }
 
+PyOperationRef PyOperation::parse(PyMlirContextRef contextRef,
+                                  const std::string &sourceStr,
+                                  const std::string &sourceName) {
+  MlirOperation op =
+      mlirOperationCreateParse(contextRef->get(), toMlirStringRef(sourceStr),
+                               toMlirStringRef(sourceName));
+  // TODO: Include error diagnostic messages in the exception message
+  if (mlirOperationIsNull(op))
+    throw py::value_error(
+        "Unable to parse operation assembly (see diagnostics)");
+  return PyOperation::createDetached(std::move(contextRef), op);
+}
+
 void PyOperation::checkValid() const {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
 
 void PyOperationBase::print(py::object fileObject, bool binary,
                             std::optional<int64_t> largeElementsLimit,
void mlir::python::populateIRCore(py::module &m) {
   py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
                   kOperationCreateDocstring)
+      .def_static(
+          "parse",
+          [](const std::string &sourceStr, const std::string &sourceName,
+             DefaultingPyMlirContext context) {
+            return PyOperation::parse(context->getRef(), sourceStr, sourceName)
+                ->createOpView();
+          },
+          py::arg("source"), py::kw_only(), py::arg("source_name") = "",
+          py::arg("context") = py::none(),
+          "Parses an operation. Supports both text assembly format and binary "
+          "bytecode format.")
       .def_property_readonly("parent",
                              [](PyOperation &self) -> py::object {
                                auto parent = self.getParentOperation();
                                if (parent)
                                  return parent->getObject();
                                return py::none();
                              })
       .def("erase", &PyOperation::erase)
void mlir::python::populateIRCore(py::module &m) {
   opViewClass.attr("_ODS_OPERAND_SEGMENTS") = py::none();
   opViewClass.attr("_ODS_RESULT_SEGMENTS") = py::none();
   opViewClass.attr("build_generic") = classmethod(
       &PyOpView::buildGeneric, py::arg("cls"), py::arg("results") = py::none(),
       py::arg("operands") = py::none(), py::arg("attributes") = py::none(),
       py::arg("successors") = py::none(), py::arg("regions") = py::none(),
       py::arg("loc") = py::none(), py::arg("ip") = py::none(),
       "Builds a specific, generated OpView based on class level attributes.");
+  opViewClass.attr("parse") = classmethod(
+      [](const py::object &cls, const std::string &sourceStr,
+         const std::string &sourceName, DefaultingPyMlirContext context) {
+        PyOperationRef parsed =
+            PyOperation::parse(context->getRef(), sourceStr, sourceName);
+
+        // Check if the expected operation was parsed, and cast to to the
+        // appropriate `OpView` subclass if successful.
+        // NOTE: This accesses attributes that have been automatically added to
+        // `OpView` subclasses, and is not intended to be used on `OpView`
+        // directly.
+        std::string clsOpName =
+            py::cast<std::string>(cls.attr("OPERATION_NAME"));
+        MlirStringRef parsedOpName =
+            mlirIdentifierStr(mlirOperationGetName(*parsed.get()));
+        if (!mlirStringRefEqual(parsedOpName, toMlirStringRef(clsOpName)))
+          throw py::value_error(
+              "Expected a '" + clsOpName + "' op, got: '" +
+              std::string(parsedOpName.data, parsedOpName.length) + "'");
+        return cls.attr("_Raw")(parsed.getObject());
+      },
+      py::arg("cls"), py::arg("source"), py::kw_only(),
+      py::arg("source_name") = "", py::arg("context") = py::none(),
+      "Parses a specific, generated OpView based on class level attributes");
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.
   //----------------------------------------------------------------------------
   py::class_<PyRegion>(m, "Region", py::module_local())
       .def_property_readonly(
           "blocks",
           [](PyRegion &self) {

mlir/lib/Bindings/Python/IRModule.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
+#include <optional>
 #include <utility>
 #include <vector>
-#include <optional>
 
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/IntegerSet.h"
forOperation(PyMlirContextRef contextRef, MlirOperation operation,
                pybind11::object parentKeepAlive = pybind11::object());
 
   /// Creates a detached operation. The operation must not be associated with
   /// any existing live operation.
   static PyOperationRef
   createDetached(PyMlirContextRef contextRef, MlirOperation operation,
                  pybind11::object parentKeepAlive = pybind11::object());
 
+  /// Parses a source string (either text assembly or bytecode), creating a
+  /// detached operation.
+  static PyOperationRef parse(PyMlirContextRef contextRef,
+                              const std::string &sourceStr,
+                              const std::string &sourceName);
+
   /// Detaches the operation from its parent block and updates its state
   /// accordingly.
   void detachFromParent() {
     mlirOperationRemoveFromParent(getOperation());
     setDetached();
     parentKeepAlive = pybind11::object();
   }
 

mlir/lib/CAPI/IR/IR.cpp

if (state->enableResultTypeInference) {
     if (failed(inferOperationTypes(cppState)))
       return {nullptr};
   }
 
   MlirOperation result = wrap(Operation::create(cppState));
   return result;
 }
 
+MlirOperation mlirOperationCreateParse(MlirContext context,
+                                       MlirStringRef sourceStr,
+                                       MlirStringRef sourceName) {
+
+  return wrap(
+      parseSourceString(unwrap(sourceStr), unwrap(context), unwrap(sourceName))
+          .release());
+}
+
 MlirOperation mlirOperationClone(MlirOperation op) {
   return wrap(unwrap(op)->clone());
 }
 
 void mlirOperationDestroy(MlirOperation op) { unwrap(op)->erase(); }
 
 void mlirOperationRemoveFromParent(MlirOperation op) { unwrap(op)->remove(); }
 

mlir/lib/Parser/Parser.cpp

LogicalResult mlir::parseSourceFile(
     Block *block, const ParserConfig &config, LocationAttr *sourceFileLoc) {
   if (failed(loadSourceFileBuffer(filename, *sourceMgr, config.getContext())))
     return failure();
   return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
 }
 
 LogicalResult mlir::parseSourceString(llvm::StringRef sourceStr, Block *block,
                                       const ParserConfig &config,
-                                      LocationAttr *sourceFileLoc) {
-  auto memBuffer = llvm::MemoryBuffer::getMemBuffer(sourceStr);
+                                      LocationAttr *sourceFileLoc,
+                                      StringRef sourceName) {
+  auto memBuffer = llvm::MemoryBuffer::getMemBuffer(sourceStr, sourceName);
   if (!memBuffer)
     return failure();
 
   llvm::SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
 }

mlir/test/python/ir/operation.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import io
 import itertools
 from mlir.ir import *
+from mlir.dialects.builtin import ModuleOp
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f

 # CHECK-LABEL: TEST: testOperationHash
 @run
 def testOperationHash():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx, Location.unknown():
     op = Operation.create("custom.op1")
     assert hash(op) == hash(op.operation)
+
+
+# CHECK-LABEL: TEST: testOperationParse
+@run
+def testOperationParse():
+  with Context() as ctx:
+    ctx.allow_unregistered_dialects = True
+
+    # Generic operation parsing.
+    m = Operation.parse('module {}')
+    o = Operation.parse('"test.foo"() : () -> ()')
+    assert isinstance(m, ModuleOp)
+    assert type(o) is OpView
+
+    # Parsing specific operation.
+    m = ModuleOp.parse('module {}')
+    assert isinstance(m, ModuleOp)
+    try:
+      ModuleOp.parse('"test.foo"() : () -> ()')
+    except ValueError as e:
+      # CHECK: error: Expected a 'builtin.module' op, got: 'test.foo'
+      print(f"error: {e}")
+    else:
+      assert False, "expected error"
+
+    o = Operation.parse('"test.foo"() : () -> ()', source_name="my-source-string")
+    # CHECK: op_with_source_name: "test.foo"() : () -> () loc("my-source-string":1:1)
+    print(f"op_with_source_name: {o.get_asm(enable_debug_info=True, use_local_scope=True)}")