[mlir][Python] Add an Operation.result property.

[stellaraccident]

Summary

• If ODS redefines this, it is fine, but I have found this accessor to be universally useful in the old npcomp bindings and I'm closing gaps that will let me switch.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

1050

spurious semicolon

2776

Can you augment the message with the operation name and the number of results it has?

2781

[stellaraccident]

Address comments.

Also add a using for StringRef and Twine (and replaces fully qualified references).

mlir/lib/Bindings/Python/IRModules.h

public:
 
   /// 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());
 
   /// Gets the backing operation.
-  MlirOperation get() {
+  operator MlirOperation() const { return get(); }
+  MlirOperation get() const {
     checkValid();
     return operation;
   }
 
   PyOperationRef getRef() {
     return PyOperationRef(
         this, pybind11::reinterpret_borrow<pybind11::object>(handle));
   }
 
   bool isAttached() { return attached; }
   void setAttached() {
     assert(!attached && "operation already attached");
     attached = true;
   }
-  void checkValid();
+  void checkValid() const;
 
   /// Gets the owning block or raises an exception if the operation has no
   /// owning block.
   PyBlock getBlock();
 
   /// Gets the parent operation or raises an exception if the operation has
   /// no parent.
   PyOperationRef getParentOperation();

mlir/lib/Bindings/Python/IRModules.cpp

 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
+using llvm::StringRef;
+using llvm::Twine;
 
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
 static const char kContextParseTypeDocstring[] =
     R"(Parses the assembly form of a type.
 
MlirDialect PyDialects::getDialectForKey(const std::string &key,
                                          bool attrError) {
   // If the "std" dialect was asked for, substitute the empty namespace :(
   static const std::string emptyKey;
   const std::string *canonKey = key == "std" ? &emptyKey : &key;
   MlirDialect dialect = mlirContextGetOrLoadDialect(
       getContext()->get(), {canonKey->data(), canonKey->size()});
   if (mlirDialectIsNull(dialect)) {
     throw SetPyError(attrError ? PyExc_AttributeError : PyExc_IndexError,
-                     llvm::Twine("Dialect '") + key + "' not found");
+                     Twine("Dialect '") + key + "' not found");
   }
   return dialect;
 }
 
 //------------------------------------------------------------------------------
 // PyLocation
 //------------------------------------------------------------------------------
 
PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
   (void)liveOperations;
 
   PyOperationRef created = createInstance(std::move(contextRef), operation,
                                           std::move(parentKeepAlive));
   created->attached = false;
   return created;
 }
 
-void PyOperation::checkValid() {
+void PyOperation::checkValid() const {
   if (!valid) {
     throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
   }
 }
 
 void PyOperationBase::print(py::object fileObject, bool binary,
                             llvm::Optional<int64_t> largeElementsLimit,
                             bool enableDebugInfo, bool prettyDebugInfo,
                             bool printGenericOpForm, bool useLocalScope) {
   PyOperation &operation = getOperation();
   operation.checkValid();
   if (fileObject.is_none())
     fileObject = py::module::import("sys").attr("stdout");
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   if (largeElementsLimit)
     mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
   if (enableDebugInfo)
     mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/prettyDebugInfo);
   if (printGenericOpForm)
     mlirOpPrintingFlagsPrintGenericOpForm(flags);
 
   PyFileAccumulator accum(fileObject, binary);
   py::gil_scoped_release();
-  mlirOperationPrintWithFlags(operation.get(), flags, accum.getCallback(),
+  mlirOperationPrintWithFlags(operation, flags, accum.getCallback(),
                               accum.getUserData());
   mlirOpPrintingFlagsDestroy(flags);
 }
 
 py::object PyOperationBase::getAsm(bool binary,
                                    llvm::Optional<int64_t> largeElementsLimit,
                                    bool enableDebugInfo, bool prettyDebugInfo,
                                    bool printGenericOpForm,
py::object PyOperation::create(
 
   return created->createOpView();
 }
 
 py::object PyOperation::createOpView() {
   MlirIdentifier ident = mlirOperationGetName(get());
   MlirStringRef identStr = mlirIdentifierStr(ident);
   auto opViewClass = PyGlobals::get().lookupRawOpViewClass(
-      llvm::StringRef(identStr.data, identStr.length));
+      StringRef(identStr.data, identStr.length));
   if (opViewClass)
     return (*opViewClass)(getRef().getObject());
   return py::cast(PyOpView(getRef().getObject()));
 }
 
 PyOpView::PyOpView(py::object operationObject)
     // Casting through the PyOperationBase base-class and then back to the
     // Operation lets us accept any PyOperationBase subclass.
throw SetPyError(PyExc_ValueError,
                      "Attempt to insert operation that is already attached");
   block.getParentOperation()->checkValid();
   MlirOperation beforeOp = {nullptr};
   if (refOperation) {
     // Insert before operation.
     (*refOperation)->checkValid();
     beforeOp = (*refOperation)->get();
   }
-  mlirBlockInsertOwnedOperationBefore(block.get(), beforeOp, operation.get());
+  mlirBlockInsertOwnedOperationBefore(block.get(), beforeOp, operation);
   operation.setAttached();
 }
 
 PyInsertionPoint PyInsertionPoint::atBlockBegin(PyBlock &block) {
   MlirOperation firstOp = mlirBlockGetFirstOperation(block.get());
   if (mlirOperationIsNull(firstOp)) {
     // Just insert at end.
     return PyInsertionPoint(block);
   }
public:
   PyConcreteValue(PyValue &orig)
       : PyConcreteValue(orig.getParentOperation(), castFrom(orig)) {}
 
   /// Attempts to cast the original value to the derived type and throws on
   /// type mismatches.
   static MlirValue castFrom(PyValue &orig) {
     if (!DerivedTy::isaFunction(orig.get())) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
-      throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast value to ") +
+      throw SetPyError(PyExc_ValueError, Twine("Cannot cast value to ") +
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
     return orig.get();
   }
 
   /// Binds the Python module objects to functions of this class.
   static void bind(py::module &m) {
public:
   PyConcreteAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
       : BaseTy(std::move(contextRef), attr) {}
   PyConcreteAttribute(PyAttribute &orig)
       : PyConcreteAttribute(orig.getContext(), castFrom(orig)) {}
 
   static MlirAttribute castFrom(PyAttribute &orig) {
     if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
-      throw SetPyError(PyExc_ValueError,
-                       llvm::Twine("Cannot cast attribute to ") +
-                           DerivedTy::pyClassName + " (from " + origRepr + ")");
+      throw SetPyError(PyExc_ValueError, Twine("Cannot cast attribute to ") +
+                                             DerivedTy::pyClassName +
+                                             " (from " + origRepr + ")");
     }
     return orig;
   }
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::buffer_protocol());
     cls.def(py::init<PyAttribute &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &type, double value, DefaultingPyLocation loc) {
           MlirAttribute attr = mlirFloatAttrDoubleGetChecked(type, value, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirAttributeIsNull(attr)) {
             throw SetPyError(PyExc_ValueError,
-                             llvm::Twine("invalid '") +
+                             Twine("invalid '") +
                                  py::repr(py::cast(type)).cast<std::string>() +
                                  "' and expected floating point type.");
           }
           return PyFloatAttribute(type.getContext(), attr);
         },
         py::arg("type"), py::arg("value"), py::arg("loc") = py::none(),
         "Gets an uniqued float point attribute associated to a type");
     c.def_static(
public:
   PyConcreteType(PyMlirContextRef contextRef, MlirType t)
       : BaseTy(std::move(contextRef), t) {}
   PyConcreteType(PyType &orig)
       : PyConcreteType(orig.getContext(), castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
     if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
-      throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
+      throw SetPyError(PyExc_ValueError, Twine("Cannot cast type to ") +
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
     return orig;
   }
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName);
c.def_static(
         [](PyType &elementType) {
           // The element must be a floating point or integer scalar type.
           if (mlirTypeIsAIntegerOrFloat(elementType)) {
             MlirType t = mlirComplexTypeGet(elementType);
             return PyComplexType(elementType.getContext(), t);
           }
           throw SetPyError(
               PyExc_ValueError,
-              llvm::Twine("invalid '") +
+              Twine("invalid '") +
                   py::repr(py::cast(elementType)).cast<std::string>() +
                   "' and expected floating point or integer type.");
         },
         "Create a complex type");
     c.def_property_readonly(
         "element_type",
         [](PyComplexType &self) -> PyType {
           MlirType t = mlirComplexTypeGetElementType(self);
c.def_static(
            DefaultingPyLocation loc) {
           MlirType t = mlirVectorTypeGetChecked(shape.size(), shape.data(),
                                                 elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
-                llvm::Twine("invalid '") +
+                Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point or integer type.");
           }
           return PyVectorType(elementType.getContext(), t);
         },
         py::arg("shape"), py::arg("elementType"), py::arg("loc") = py::none(),
         "Create a vector type");
   }
c.def_static(
            DefaultingPyLocation loc) {
           MlirType t = mlirRankedTensorTypeGetChecked(
               shape.size(), shape.data(), elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
-                llvm::Twine("invalid '") +
+                Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "
                     "complex "
                     "type.");
           }
           return PyRankedTensorType(elementType.getContext(), t);
         },
         py::arg("shape"), py::arg("element_type"), py::arg("loc") = py::none(),
c.def_static(
         "get",
         [](PyType &elementType, DefaultingPyLocation loc) {
           MlirType t = mlirUnrankedTensorTypeGetChecked(elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
-                llvm::Twine("invalid '") +
+                Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "
                     "complex "
                     "type.");
           }
           return PyUnrankedTensorType(elementType.getContext(), t);
         },
         py::arg("element_type"), py::arg("loc") = py::none(),
c.def_static(
             unsigned memorySpace, DefaultingPyLocation loc) {
            MlirType t = mlirMemRefTypeContiguousGetChecked(
                elementType, shape.size(), shape.data(), memorySpace, loc);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // in C API.
            if (mlirTypeIsNull(t)) {
              throw SetPyError(
                  PyExc_ValueError,
-                 llvm::Twine("invalid '") +
+                 Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("shape"), py::arg("memory_space"),
c.def_static(
             DefaultingPyLocation loc) {
            MlirType t =
                mlirUnrankedMemRefTypeGetChecked(elementType, memorySpace, loc);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // in C API.
            if (mlirTypeIsNull(t)) {
              throw SetPyError(
                  PyExc_ValueError,
-                 llvm::Twine("invalid '") +
+                 Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyUnrankedMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("memory_space"),
py::class_<PyMlirContext>(m, "Context")
           "Alias for 'dialect'")
       .def(
           "get_dialect_descriptor",
           [=](PyMlirContext &self, std::string &name) {
             MlirDialect dialect = mlirContextGetOrLoadDialect(
                 self.get(), {name.data(), name.size()});
             if (mlirDialectIsNull(dialect)) {
               throw SetPyError(PyExc_ValueError,
-                               llvm::Twine("Dialect '") + name + "' not found");
+                               Twine("Dialect '") + name + "' not found");
             }
             return PyDialectDescriptor(self.getRef(), dialect);
           },
           "Gets or loads a dialect by name, returning its descriptor object")
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
py::class_<PyOperationBase>(m, "_OperationBase")
                                    self.getOperation().getRef());
                              })
       .def_property_readonly(
           "results",
           [](PyOperationBase &self) {
             return PyOpResultList(self.getOperation().getRef());
           },
           "Returns the list of Operation results.")
+      .def_property_readonly(
+          "result",
+          [](PyOperationBase &self) {
+            auto &operation = self.getOperation();
+            auto numResults = mlirOperationGetNumResults(operation);
+            if (numResults != 1) {
+              auto name = mlirIdentifierStr(mlirOperationGetName(operation));
+              throw SetPyError(
+                  PyExc_ValueError,
+                  Twine("Cannot call .result on operation ") +
+                      StringRef(name.data, name.length) + " which has " +
+                      Twine(numResults) +
+                      " results (it is only valid for operations with a "
+                      "single result)");
+            }
+            return PyOpResult(operation.getRef(),
+                              mlirOperationGetResult(operation, 0));
+          },
+          "Shortcut to get an op result if it has only one (throws an error "
+          "otherwise).")
       .def("__iter__",
            [](PyOperationBase &self) {
              return PyRegionIterator(self.getOperation().getRef());
            })
       .def(
           "__str__",
           [](PyOperationBase &self) {
             return self.getAsm(/*binary=*/false,
py::class_<PyAttribute>(m, "Attribute")
           "parse",
           [](std::string attrSpec, DefaultingPyMlirContext context) {
             MlirAttribute type = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirAttributeIsNull(type)) {
               throw SetPyError(PyExc_ValueError,
-                               llvm::Twine("Unable to parse attribute: '") +
+                               Twine("Unable to parse attribute: '") +
                                    attrSpec + "'");
             }
             return PyAttribute(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           "Parses an attribute from an assembly form")
       .def_property_readonly(
           "context",
py::class_<PyType>(m, "Type")
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
             MlirType type =
                 mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirTypeIsNull(type)) {
               throw SetPyError(PyExc_ValueError,
-                               llvm::Twine("Unable to parse type: '") +
-                                   typeSpec + "'");
+                               Twine("Unable to parse type: '") + typeSpec +
+                                   "'");
             }
             return PyType(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")

mlir/test/Bindings/Python/ir_operation.py

def testOperationPrint():
   # CHECK: value = opaque<"", "0xDEADBEEF"> : tensor<4xi32>
   # CHECK: "std.return"(%arg0) : (i32) -> () -:4:7
   module.operation.print(large_elements_limit=2, enable_debug_info=True,
       pretty_debug_info=True, print_generic_op_form=True, use_local_scope=True)
 
 run(testOperationPrint)
 
 
+# CHECK-LABEL: TEST: testKnownOpView
 def testKnownOpView():
   with Context(), Location.unknown():
     Context.current.allow_unregistered_dialects = True
     module = Module.parse(r"""
       %1 = "custom.f32"() : () -> f32
       %2 = "custom.f32"() : () -> f32
       %3 = addf %1, %2 : f32
     """)
with Context(), Location.unknown():
     print(repr(custom))
 
     # Check again to make sure negative caching works.
     custom = module.body.operations[0]
     # CHECK: <_mlir.ir.OpView object
     print(repr(custom))
 
 run(testKnownOpView)
+
+
+# CHECK-LABEL: TEST: testSingleResultProperty
+def testSingleResultProperty():
+  with Context(), Location.unknown():
+    Context.current.allow_unregistered_dialects = True
+    module = Module.parse(r"""
+      "custom.no_result"() : () -> ()
+      %0:2 = "custom.two_result"() : () -> (f32, f32)
+      %1 = "custom.one_result"() : () -> f32
+    """)
+    print(module)
+
+  try:
+    module.body.operations[0].result
+  except ValueError as e:
+    # CHECK: Cannot call .result on operation custom.no_result which has 0 results
+    print(e)
+  else:
+    assert False, "Expected exception"
+
+  try:
+    module.body.operations[1].result
+  except ValueError as e:
+    # CHECK: Cannot call .result on operation custom.two_result which has 2 results
+    print(e)
+  else:
+    assert False, "Expected exception"
+
+  # CHECK: %1 = "custom.one_result"() : () -> f32
+  print(module.body.operations[2])
+
+run(testSingleResultProperty)