diff --git a/mlir/include/mlir-c/StandardAttributes.h b/mlir/include/mlir-c/StandardAttributes.h --- a/mlir/include/mlir-c/StandardAttributes.h +++ b/mlir/include/mlir-c/StandardAttributes.h @@ -404,6 +404,10 @@ MLIR_CAPI_EXPORTED MlirStringRef mlirDenseElementsAttrGetStringValue(MlirAttribute attr, intptr_t pos); +/** Returns the raw data of the given dense elements attribute. */ +MLIR_CAPI_EXPORTED const void * +mlirDenseElementsAttrGetRawData(MlirAttribute attr); + //===----------------------------------------------------------------------===// // Opaque elements attribute. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp --- a/mlir/lib/Bindings/Python/IRModules.cpp +++ b/mlir/lib/Bindings/Python/IRModules.cpp @@ -16,6 +16,7 @@ #include "mlir-c/StandardAttributes.h" #include "mlir-c/StandardTypes.h" #include "llvm/ADT/SmallVector.h" +#include #include namespace py = pybind11; @@ -1621,6 +1622,48 @@ intptr_t dunderLen() { return mlirElementsAttrGetNumElements(attr); } + py::array accessArray() { + MlirType shapedType = mlirAttributeGetType(this->attr); + MlirType elementType = mlirShapedTypeGetElementType(shapedType); + + if (mlirTypeIsAF32(elementType)) { + // f32 + return arrayInfo(py::dtype("f"), shapedType, + mlirDenseElementsAttrGetFloatValue); + } else if (mlirTypeIsAF64(elementType)) { + // f64 + return arrayInfo(py::dtype("d"), shapedType, + mlirDenseElementsAttrGetDoubleValue); + } else if (mlirTypeIsAInteger(elementType) && + mlirIntegerTypeGetWidth(elementType) == 32) { + if (mlirIntegerTypeIsSignless(elementType) || + mlirIntegerTypeIsSigned(elementType)) { + // i32 + return arrayInfo(py::dtype("i"), shapedType, + mlirDenseElementsAttrGetInt32Value); + } else if (mlirIntegerTypeIsUnsigned(elementType)) { + // unsigned i32 + return arrayInfo(py::dtype("I"), shapedType, + mlirDenseElementsAttrGetUInt32Value); + } + } else if (mlirTypeIsAInteger(elementType) && + mlirIntegerTypeGetWidth(elementType) == 64) { + if (mlirIntegerTypeIsSignless(elementType) || + mlirIntegerTypeIsSigned(elementType)) { + // i64 + return arrayInfo(py::dtype("l"), shapedType, + mlirDenseElementsAttrGetInt64Value); + } else if (mlirIntegerTypeIsUnsigned(elementType)) { + // unsigned i64 + return arrayInfo(py::dtype("L"), shapedType, + mlirDenseElementsAttrGetUInt64Value); + } + } + + std::string message = "unimplemented array format."; + throw SetPyError(PyExc_ValueError, message); + } + static void bindDerived(ClassTy &c) { c.def("__len__", &PyDenseElementsAttribute::dunderLen) .def_static("get", PyDenseElementsAttribute::getFromBuffer, @@ -1633,7 +1676,8 @@ .def_property_readonly("is_splat", [](PyDenseElementsAttribute &self) -> bool { return mlirDenseElementsAttrIsSplat(self.attr); - }); + }) + .def_property_readonly("array", &PyDenseElementsAttribute::accessArray); } private: @@ -1650,6 +1694,33 @@ const ElementTy *contents = static_cast(arrayInfo.ptr); return ctor(shapedType, numElements, contents); } + + template + py::array arrayInfo(const py::dtype dtype, MlirType shapedType, + Type (*value)(MlirAttribute, intptr_t)) { + intptr_t rank = mlirShapedTypeGetRank(shapedType); + // Prepare the data for the buffer_info. + Type *data = reinterpret_cast( + const_cast(mlirDenseElementsAttrGetRawData(this->attr))); + // Prepare the shape for the buffer_info. + std::vector shape; + for (intptr_t i = 0; i < rank; ++i) + shape.push_back(mlirShapedTypeGetDimSize(shapedType, i)); + // Prepare the strides for the buffer_info. + std::vector strides; + intptr_t strideFactor = 1; + for (intptr_t i = 1; i < rank; ++i) { + strideFactor = 1; + for (intptr_t j = i; j < rank; ++j) { + strideFactor *= mlirShapedTypeGetDimSize(shapedType, j); + } + strides.push_back(sizeof(Type) * strideFactor); + } + strides.push_back(sizeof(Type)); + // Set the base object to keep the backing data alive. + auto base = py::cast(this, py::return_value_policy::take_ownership); + return py::array(dtype, shape, strides, data, std::move(base)); + } }; /// Refinement of the PyDenseElementsAttribute for attributes containing integer diff --git a/mlir/lib/CAPI/IR/StandardAttributes.cpp b/mlir/lib/CAPI/IR/StandardAttributes.cpp --- a/mlir/lib/CAPI/IR/StandardAttributes.cpp +++ b/mlir/lib/CAPI/IR/StandardAttributes.cpp @@ -516,6 +516,14 @@ pos)); } +//===----------------------------------------------------------------------===// +// Raw data accessors. + +const void *mlirDenseElementsAttrGetRawData(MlirAttribute attr) { + return static_cast( + unwrap(attr).cast().getRawData().data()); +} + //===----------------------------------------------------------------------===// // Opaque elements attribute. //===----------------------------------------------------------------------===// diff --git a/mlir/test/Bindings/Python/ir_array_attributes.py b/mlir/test/Bindings/Python/ir_array_attributes.py --- a/mlir/test/Bindings/Python/ir_array_attributes.py +++ b/mlir/test/Bindings/Python/ir_array_attributes.py @@ -106,6 +106,9 @@ print(attr) # CHECK: is_splat: False print("is_splat:", attr.is_splat) + # CHECK: {{\[}}[1.1 2.2 3.3] + # CHECK: {{\[}}4.4 5.5 6.6]] + print(attr.array) run(testGetDenseElementsF32) @@ -117,6 +120,9 @@ attr = DenseElementsAttr.get(array) # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf64> print(attr) + # CHECK: {{\[}}[1.1 2.2 3.3] + # CHECK: {{\[}}4.4 5.5 6.6]] + print(attr.array) run(testGetDenseElementsF64) @@ -129,6 +135,9 @@ attr = DenseElementsAttr.get(array) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsI32Signless) @@ -140,6 +149,9 @@ attr = DenseElementsAttr.get(array) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsUI32Signless) @@ -150,6 +162,9 @@ attr = DenseElementsAttr.get(array, signless=False) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi32> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsI32) @@ -161,6 +176,9 @@ attr = DenseElementsAttr.get(array, signless=False) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui32> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsUI32) @@ -173,6 +191,9 @@ attr = DenseElementsAttr.get(array) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsI64Signless) @@ -184,6 +205,9 @@ attr = DenseElementsAttr.get(array) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsUI64Signless) @@ -194,6 +218,9 @@ attr = DenseElementsAttr.get(array, signless=False) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi64> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsI64) @@ -205,6 +232,9 @@ attr = DenseElementsAttr.get(array, signless=False) # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui64> print(attr) + # CHECK: {{\[}}[1 2 3] + # CHECK: {{\[}}4 5 6]] + print(attr.array) run(testGetDenseElementsUI64) diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c --- a/mlir/test/CAPI/ir.c +++ b/mlir/test/CAPI/ir.c @@ -903,6 +903,26 @@ fabs(mlirDenseElementsAttrGetDoubleSplatValue(splatDouble) - 1.0) > 1E-6) return 17; + uint32_t *uint32RawData = + (uint32_t *)mlirDenseElementsAttrGetRawData(uint32Elements); + int32_t *int32RawData = + (int32_t *)mlirDenseElementsAttrGetRawData(int32Elements); + uint64_t *uint64RawData = + (uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements); + int64_t *int64RawData = + (int64_t *)mlirDenseElementsAttrGetRawData(int64Elements); + float *floatRawData = + (float *)mlirDenseElementsAttrGetRawData(floatElements); + double *doubleRawData = + (double *)mlirDenseElementsAttrGetRawData(doubleElements); + if (uint32RawData[0] != 0u || uint32RawData[1] != 1u || + int32RawData[0] != 0 || int32RawData[1] != 1 || + uint64RawData[0] != 0u || uint64RawData[1] != 1u || + int64RawData[0] != 0 || int64RawData[1] != 1 || + floatRawData[0] != 0.0f || floatRawData[1] != 1.0f || + doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0) + return 18; + mlirAttributeDump(splatBool); mlirAttributeDump(splatUInt32); mlirAttributeDump(splatInt32);