Skip to content

Commit 33df617

Browse files
stellaraccidentstellaraccident
authored
[mlir] Quality of life improvements to python API types. (#66723)
* Moves several orphaned methods from Operation/OpView -> _OperationBase so that both hierarchies share them (whether unknown or known to ODS). * Adds typing information for missing `MLIRError` exception. * Adds `DiagnosticInfo` typing. * Adds `DenseResourceElementsAttr` typing that was missing.
1 parent acfb99d commit 33df617

File tree

3 files changed

+67
-39
lines changed

3 files changed

+67
-39
lines changed

‎mlir/lib/Bindings/Python/IRCore.cpp

+36-31
Original file line numberDiff line numberDiff line change
@@ -2768,6 +2768,24 @@ void mlir::python::populateIRCore(py::module &m) {
27682768
return PyOpAttributeMap(
27692769
self.getOperation().getRef());
27702770
})
2771+
.def_property_readonly(
2772+
"context",
2773+
[](PyOperationBase &self) {
2774+
PyOperation &concreteOperation = self.getOperation();
2775+
concreteOperation.checkValid();
2776+
return concreteOperation.getContext().getObject();
2777+
},
2778+
"Context that owns the Operation")
2779+
.def_property_readonly("name",
2780+
[](PyOperationBase &self) {
2781+
auto &concreteOperation = self.getOperation();
2782+
concreteOperation.checkValid();
2783+
MlirOperation operation =
2784+
concreteOperation.get();
2785+
MlirStringRef name = mlirIdentifierStr(
2786+
mlirOperationGetName(operation));
2787+
return py::str(name.data, name.length);
2788+
})
27712789
.def_property_readonly("operands",
27722790
[](PyOperationBase &self) {
27732791
return PyOpOperandList(
@@ -2813,6 +2831,14 @@ void mlir::python::populateIRCore(py::module &m) {
28132831
},
28142832
"Returns the source location the operation was defined or derived "
28152833
"from.")
2834+
.def_property_readonly("parent",
2835+
[](PyOperationBase &self) -> py::object {
2836+
auto parent =
2837+
self.getOperation().getParentOperation();
2838+
if (parent)
2839+
return parent->getObject();
2840+
return py::none();
2841+
})
28162842
.def(
28172843
"__str__",
28182844
[](PyOperationBase &self) {
@@ -2855,6 +2881,12 @@ void mlir::python::populateIRCore(py::module &m) {
28552881
.def("move_before", &PyOperationBase::moveBefore, py::arg("other"),
28562882
"Puts self immediately before the other operation in its parent "
28572883
"block.")
2884+
.def(
2885+
"clone",
2886+
[](PyOperationBase &self, py::object ip) {
2887+
return self.getOperation().clone(ip);
2888+
},
2889+
py::arg("ip") = py::none())
28582890
.def(
28592891
"detach_from_parent",
28602892
[](PyOperationBase &self) {
@@ -2866,7 +2898,8 @@ void mlir::python::populateIRCore(py::module &m) {
28662898
operation.detachFromParent();
28672899
return operation.createOpView();
28682900
},
2869-
"Detaches the operation from its parent block.");
2901+
"Detaches the operation from its parent block.")
2902+
.def("erase", [](PyOperationBase &self) { self.getOperation().erase(); });
28702903

28712904
py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
28722905
.def_static("create", &PyOperation::create, py::arg("name"),
@@ -2887,45 +2920,17 @@ void mlir::python::populateIRCore(py::module &m) {
28872920
py::arg("context") = py::none(),
28882921
"Parses an operation. Supports both text assembly format and binary "
28892922
"bytecode format.")
2890-
.def_property_readonly("parent",
2891-
[](PyOperation &self) -> py::object {
2892-
auto parent = self.getParentOperation();
2893-
if (parent)
2894-
return parent->getObject();
2895-
return py::none();
2896-
})
2897-
.def("erase", &PyOperation::erase)
2898-
.def("clone", &PyOperation::clone, py::arg("ip") = py::none())
28992923
.def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
29002924
&PyOperation::getCapsule)
29012925
.def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyOperation::createFromCapsule)
2902-
.def_property_readonly("name",
2903-
[](PyOperation &self) {
2904-
self.checkValid();
2905-
MlirOperation operation = self.get();
2906-
MlirStringRef name = mlirIdentifierStr(
2907-
mlirOperationGetName(operation));
2908-
return py::str(name.data, name.length);
2909-
})
2910-
.def_property_readonly(
2911-
"context",
2912-
[](PyOperation &self) {
2913-
self.checkValid();
2914-
return self.getContext().getObject();
2915-
},
2916-
"Context that owns the Operation")
2926+
.def_property_readonly("operation", [](py::object self) { return self; })
29172927
.def_property_readonly("opview", &PyOperation::createOpView);
29182928

29192929
auto opViewClass =
29202930
py::class_<PyOpView, PyOperationBase>(m, "OpView", py::module_local())
29212931
.def(py::init<py::object>(), py::arg("operation"))
29222932
.def_property_readonly("operation", &PyOpView::getOperationObject)
2923-
.def_property_readonly(
2924-
"context",
2925-
[](PyOpView &self) {
2926-
return self.getOperation().getContext().getObject();
2927-
},
2928-
"Context that owns the Operation")
2933+
.def_property_readonly("opview", [](py::object self) { return self; })
29292934
.def("__str__", [](PyOpView &self) {
29302935
return py::str(self.getOperationObject());
29312936
});

‎mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

+30-7
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,13 @@ __all__ = [
4343
"DenseElementsAttr",
4444
"DenseFPElementsAttr",
4545
"DenseIntElementsAttr",
46+
"DenseResourceElementsAttr",
4647
"Dialect",
4748
"DialectDescriptor",
4849
"Dialects",
4950
"Diagnostic",
5051
"DiagnosticHandler",
52+
"DiagnosticInfo",
5153
"DiagnosticSeverity",
5254
"DictAttr",
5355
"Float8E4M3FNType",
@@ -74,6 +76,7 @@ __all__ = [
7476
"Location",
7577
"MemRefType",
7678
"Module",
79+
"MLIRError",
7780
"NamedAttribute",
7881
"NoneType",
7982
"OpaqueType",
@@ -123,10 +126,16 @@ class _OperationBase:
123126
@property
124127
def attributes(self) -> OpAttributeMap: ...
125128
@property
129+
def context(self) -> Context: ...
130+
@property
126131
def location(self) -> Location: ...
127132
@property
133+
def name(self) -> str: ...
134+
@property
128135
def operands(self) -> OpOperandList: ...
129136
@property
137+
@property
138+
def parent(self) -> Optional[_OperationBase]: ...
130139
def regions(self) -> RegionSequence: ...
131140
@property
132141
def result(self) -> OpResult: ...
@@ -530,6 +539,10 @@ class DenseIntElementsAttr(DenseElementsAttr):
530539
@property
531540
def type(self) -> Type: ...
532541

542+
class DenseResourceElementsAttr(Attribute):
543+
@staticmethod
544+
def get_from_buffer(array: Any, name: str, type: Type, alignment: Optional[int] = None, is_mutable: bool = False, context: Optional[Context] = None) -> None: ...
545+
533546
class Dialect:
534547
def __init__(self, descriptor: DialectDescriptor) -> None: ...
535548
@property
@@ -563,6 +576,17 @@ class DiagnosticHandler:
563576
def __enter__(self) -> DiagnosticHandler: ...
564577
def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
565578

579+
class DiagnosticInfo:
580+
def __init__(self, diag: Diagnostic) -> None: ...
581+
@property
582+
def severity(self) -> "DiagnosticSeverity": ...
583+
@property
584+
def location(self) -> "Location": ...
585+
@property
586+
def message(self) -> str: ...
587+
@property
588+
def notes(self) -> Sequence["DiagnosticInfo"]: ...
589+
566590
class DiagnosticSeverity:
567591
ERROR: DiagnosticSeverity
568592
WARNING: DiagnosticSeverity
@@ -871,6 +895,9 @@ class Module:
871895
@property
872896
def operation(self) -> Operation: ...
873897

898+
class MLIRError(Exception):
899+
def __init__(self, message: str, error_diagnostics: List[DiagnosticInfo]) -> None: ...
900+
874901
class NamedAttribute:
875902
@property
876903
def attr(self) -> Attribute: ...
@@ -950,9 +977,9 @@ class OpView(_OperationBase):
950977
loc: Optional[Location] = None,
951978
ip: Optional[InsertionPoint] = None) -> _TOperation: ...
952979
@property
953-
def context(self) -> Context: ...
954-
@property
955980
def operation(self) -> Operation: ...
981+
@property
982+
def opview(self) -> "OpView": ...
956983

957984
class Operation(_OperationBase):
958985
def _CAPICreate(self) -> object: ...
@@ -968,13 +995,9 @@ class Operation(_OperationBase):
968995
@property
969996
def _CAPIPtr(self) -> object: ...
970997
@property
971-
def context(self) -> Context: ...
972-
@property
973-
def name(self) -> str: ...
998+
def operation(self) -> "Operation": ...
974999
@property
9751000
def opview(self) -> OpView: ...
976-
@property
977-
def parent(self) -> Optional[_OperationBase]: ...
9781001

9791002
class OperationIterator:
9801003
def __iter__(self) -> OperationIterator: ...

‎mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi

+1-1
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,6 @@ class PassManager:
2020
def enable_verifier(self, enable: bool) -> None: ...
2121
@staticmethod
2222
def parse(pipeline: str, context: Optional[_ir.Context] = None) -> PassManager: ...
23-
def run(self, module: _ir.Module) -> None: ...
23+
def run(self, module: _ir._OperationBase) -> None: ...
2424
@property
2525
def _CAPIPtr(self) -> object: ...

0 commit comments

Comments
 (0)