Author: Mehdi Amini Date: 2020-12-10T20:51:34Z New Revision: 285c0aa262c9255e6ea4efbce1418e5f5f17e9c1
URL: https://github.com/llvm/llvm-project/commit/285c0aa262c9255e6ea4efbce1418e5f5f17e9c1 DIFF: https://github.com/llvm/llvm-project/commit/285c0aa262c9255e6ea4efbce1418e5f5f17e9c1.diff LOG: Add MLIR Python binding for Array Attribute Differential Revision: https://reviews.llvm.org/D92948 Added: Modified: mlir/lib/Bindings/Python/IRModules.cpp mlir/test/Bindings/Python/ir_attributes.py Removed: ################################################################################ diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp index 5519c66ee1ab..5ebb2e4ccee3 100644 --- a/mlir/lib/Bindings/Python/IRModules.cpp +++ b/mlir/lib/Bindings/Python/IRModules.cpp @@ -1461,6 +1461,83 @@ class PyConcreteAttribute : public BaseTy { static void bindDerived(ClassTy &m) {} }; +class PyArrayAttribute : public PyConcreteAttribute<PyArrayAttribute> { +public: + static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAArray; + static constexpr const char *pyClassName = "ArrayAttr"; + using PyConcreteAttribute::PyConcreteAttribute; + + class PyArrayAttributeIterator { + public: + PyArrayAttributeIterator(PyAttribute attr) : attr(attr) {} + + PyArrayAttributeIterator &dunderIter() { return *this; } + + PyAttribute dunderNext() { + if (nextIndex >= mlirArrayAttrGetNumElements(attr.get())) { + throw py::stop_iteration(); + } + return PyAttribute(attr.getContext(), + mlirArrayAttrGetElement(attr.get(), nextIndex++)); + } + + static void bind(py::module &m) { + py::class_<PyArrayAttributeIterator>(m, "ArrayAttributeIterator") + .def("__iter__", &PyArrayAttributeIterator::dunderIter) + .def("__next__", &PyArrayAttributeIterator::dunderNext); + } + + private: + PyAttribute attr; + int nextIndex = 0; + }; + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](py::list attributes, DefaultingPyMlirContext context) { + SmallVector<MlirAttribute> mlirAttributes; + mlirAttributes.reserve(py::len(attributes)); + for (auto attribute : attributes) { + try { + mlirAttributes.push_back(attribute.cast<PyAttribute>()); + } catch (py::cast_error &err) { + std::string msg = std::string("Invalid attribute when attempting " + "to create an ArrayAttribute (") + + err.what() + ")"; + throw py::cast_error(msg); + } catch (py::reference_cast_error &err) { + // This exception seems thrown when the value is "None". + std::string msg = + std::string("Invalid attribute (None?) when attempting to " + "create an ArrayAttribute (") + + err.what() + ")"; + throw py::cast_error(msg); + } + } + MlirAttribute attr = mlirArrayAttrGet( + context->get(), mlirAttributes.size(), mlirAttributes.data()); + return PyArrayAttribute(context->getRef(), attr); + }, + py::arg("attributes"), py::arg("context") = py::none(), + "Gets a uniqued Array attribute"); + c.def("__getitem__", + [](PyArrayAttribute &arr, intptr_t i) { + if (i >= mlirArrayAttrGetNumElements(arr)) + throw py::index_error("ArrayAttribute index out of range"); + return PyAttribute(arr.getContext(), + mlirArrayAttrGetElement(arr, i)); + }) + .def("__len__", + [](const PyArrayAttribute &arr) { + return mlirArrayAttrGetNumElements(arr); + }) + .def("__iter__", [](const PyArrayAttribute &arr) { + return PyArrayAttributeIterator(arr); + }); + } +}; + /// Float Point Attribute subclass - FloatAttr. class PyFloatAttribute : public PyConcreteAttribute<PyFloatAttribute> { public: @@ -3089,6 +3166,8 @@ void mlir::python::populateIRSubmodule(py::module &m) { // Builtin attribute bindings. PyFloatAttribute::bind(m); + PyArrayAttribute::bind(m); + PyArrayAttribute::PyArrayAttributeIterator::bind(m); PyIntegerAttribute::bind(m); PyBoolAttribute::bind(m); PyStringAttribute::bind(m); diff --git a/mlir/test/Bindings/Python/ir_attributes.py b/mlir/test/Bindings/Python/ir_attributes.py index 4ad180bb1b37..642c1f6a836c 100644 --- a/mlir/test/Bindings/Python/ir_attributes.py +++ b/mlir/test/Bindings/Python/ir_attributes.py @@ -269,3 +269,54 @@ def testTypeAttr(): run(testTypeAttr) + + +# CHECK-LABEL: TEST: testArrayAttr +def testArrayAttr(): + with Context(): + raw = Attribute.parse("[42, true, vector<4xf32>]") + # CHECK: attr: [42, true, vector<4xf32>] + print("raw attr:", raw) + # CHECK: - 42 + # CHECK: - true + # CHECK: - vector<4xf32> + for attr in ArrayAttr(raw): + print("- ", attr) + + with Context(): + intAttr = Attribute.parse("42") + vecAttr = Attribute.parse("vector<4xf32>") + boolAttr = BoolAttr.get(True) + raw = ArrayAttr.get([vecAttr, boolAttr, intAttr]) + # CHECK: attr: [vector<4xf32>, true, 42] + print("raw attr:", raw) + # CHECK: - vector<4xf32> + # CHECK: - true + # CHECK: - 42 + arr = ArrayAttr(raw) + for attr in arr: + print("- ", attr) + # CHECK: attr[0]: vector<4xf32> + print("attr[0]:", arr[0]) + # CHECK: attr[1]: true + print("attr[1]:", arr[1]) + # CHECK: attr[2]: 42 + print("attr[2]:", arr[2]) + try: + print("attr[3]:", arr[3]) + except IndexError as e: + # CHECK: Error: ArrayAttribute index out of range + print("Error: ", e) + with Context(): + try: + ArrayAttr.get([None]) + except RuntimeError as e: + # CHECK: Error: Invalid attribute (None?) when attempting to create an ArrayAttribute + print("Error: ", e) + try: + ArrayAttr.get([42]) + except RuntimeError as e: + # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute (Unable to cast Python instance of type <class 'int'> to C++ type 'mlir::python::PyAttribute') + print("Error: ", e) +run(testArrayAttr) + _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits