diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 1f8155000525..3e927ceec190 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -868,22 +868,23 @@ py::object PyOperationBase::getAsm(bool binary, return fileObject.attr("getvalue")(); } -PyOperationRef PyOperation::getParentOperation() { +llvm::Optional PyOperation::getParentOperation() { checkValid(); if (!isAttached()) throw SetPyError(PyExc_ValueError, "Detached operations have no parent"); MlirOperation operation = mlirOperationGetParentOperation(get()); if (mlirOperationIsNull(operation)) - throw SetPyError(PyExc_ValueError, "Operation has no parent."); + return {}; return PyOperation::forOperation(getContext(), operation); } PyBlock PyOperation::getBlock() { checkValid(); - PyOperationRef parentOperation = getParentOperation(); + llvm::Optional parentOperation = getParentOperation(); MlirBlock block = mlirOperationGetBlock(get()); assert(!mlirBlockIsNull(block) && "Attached operation has null parent"); - return PyBlock{std::move(parentOperation), block}; + assert(parentOperation && "Operation has no parent"); + return PyBlock{std::move(*parentOperation), block}; } py::object PyOperation::getCapsule() { @@ -2121,8 +2122,11 @@ void mlir::python::populateIRCore(py::module &m) { py::arg("loc") = py::none(), py::arg("ip") = py::none(), kOperationCreateDocstring) .def_property_readonly("parent", - [](PyOperation &self) { - return self.getParentOperation().getObject(); + [](PyOperation &self) -> py::object { + auto parent = self.getParentOperation(); + if (parent) + return parent->getObject(); + return py::none(); }) .def("erase", &PyOperation::erase) .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h index 79c480e9446f..9d217c872191 100644 --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -18,6 +18,7 @@ #include "mlir-c/IR.h" #include "mlir-c/IntegerSet.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/Optional.h" namespace mlir { namespace python { @@ -452,7 +453,7 @@ public: /// Gets the parent operation or raises an exception if the operation has /// no parent. - PyOperationRef getParentOperation(); + llvm::Optional getParentOperation(); /// Gets a capsule wrapping the void* within the MlirOperation. pybind11::object getCapsule();