From 37e4ced0a2869d4ece39519e13d2da0b8c9129dc Mon Sep 17 00:00:00 2001 From: Christopher McGirr Date: Mon, 29 Jan 2024 15:02:29 +0100 Subject: [PATCH] [mlir][python] allow DenseIntElementsAttr for index type Model the IndexType as size_t when converting to a python integer. This should follow the definition of the IndexType as system-dependent. This solves the following issue when for example an attribute is defined as: ``` denseattr = dense<[1, 2, 3, 4]> : vector<4xindex> ``` which is a valid DenseElementsAttr. The change will fix an access issue when reading the attribute in python. E.g. ``` DenseIntElementsAttr(op.attributes["denseattr"]) ``` Co-authored-by: Tiago Trevisan Jost Co-authored-by: Matthias Gehre --- mlir/include/mlir-c/BuiltinAttributes.h | 2 ++ mlir/lib/Bindings/Python/IRAttributes.cpp | 10 ++++++++-- mlir/lib/CAPI/IR/BuiltinAttributes.cpp | 3 +++ mlir/test/python/dialects/builtin.py | 4 ++++ mlir/test/python/ir/array_attributes.py | 4 ++++ mlir/test/python/ir/attributes.py | 4 ++++ 6 files changed, 25 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir-c/BuiltinAttributes.h b/mlir/include/mlir-c/BuiltinAttributes.h index 01d1b6008f5e21..5070defddb52d6 100644 --- a/mlir/include/mlir-c/BuiltinAttributes.h +++ b/mlir/include/mlir-c/BuiltinAttributes.h @@ -543,6 +543,8 @@ MLIR_CAPI_EXPORTED int64_t mlirDenseElementsAttrGetInt64Value(MlirAttribute attr, intptr_t pos); MLIR_CAPI_EXPORTED uint64_t mlirDenseElementsAttrGetUInt64Value(MlirAttribute attr, intptr_t pos); +MLIR_CAPI_EXPORTED size_t mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, + intptr_t pos); MLIR_CAPI_EXPORTED float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr, intptr_t pos); MLIR_CAPI_EXPORTED double diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp index dda2003ba0375a..90f774c8289ad2 100644 --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -973,13 +973,19 @@ class PyDenseIntElementsAttribute MlirType type = mlirAttributeGetType(*this); type = mlirShapedTypeGetElementType(type); - assert(mlirTypeIsAInteger(type) && - "expected integer element type in dense int elements attribute"); + // Index type can also appear as a DenseIntElementsAttr and therefore can be + // casted to integer. + assert(mlirTypeIsAInteger(type) || + mlirTypeIsAIndex(type) && "expected integer/index element type in " + "dense int elements attribute"); // Dispatch element extraction to an appropriate C function based on the // elemental type of the attribute. py::int_ is implicitly constructible // from any C++ integral type and handles bitwidth correctly. // TODO: consider caching the type properties in the constructor to avoid // querying them on each element access. + if (mlirTypeIsAIndex(type)) { + return mlirDenseElementsAttrGetIndexValue(*this, pos); + } unsigned width = mlirIntegerTypeGetWidth(type); bool isUnsigned = mlirIntegerTypeIsUnsigned(type); if (isUnsigned) { diff --git a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp index b3066ee0c28bdc..cb6ce16fa317af 100644 --- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp +++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp @@ -745,6 +745,9 @@ int64_t mlirDenseElementsAttrGetInt64Value(MlirAttribute attr, intptr_t pos) { uint64_t mlirDenseElementsAttrGetUInt64Value(MlirAttribute attr, intptr_t pos) { return llvm::cast(unwrap(attr)).getValues()[pos]; } +size_t mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, intptr_t pos) { + return llvm::cast(unwrap(attr)).getValues()[pos]; +} float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr, intptr_t pos) { return llvm::cast(unwrap(attr)).getValues()[pos]; } diff --git a/mlir/test/python/dialects/builtin.py b/mlir/test/python/dialects/builtin.py index 18ebba61e7fea8..973a0eaeca2cdb 100644 --- a/mlir/test/python/dialects/builtin.py +++ b/mlir/test/python/dialects/builtin.py @@ -246,3 +246,7 @@ def testDenseElementsAttr(): # CHECK{LITERAL}: dense<[[0, 1], [2, 3]]> : tensor<2x2xi32> print(DenseElementsAttr.get(values, type=VectorType.get((2, 2), i32))) # CHECK{LITERAL}: dense<[[0, 1], [2, 3]]> : vector<2x2xi32> + idx_values = np.arange(4, dtype=np.int64) + idx_type = IndexType.get() + print(DenseElementsAttr.get(idx_values, type=VectorType.get([4], idx_type))) + # CHECK{LITERAL}: dense<[0, 1, 2, 3]> : vector<4xindex> diff --git a/mlir/test/python/ir/array_attributes.py b/mlir/test/python/ir/array_attributes.py index 9251588a4c48a6..80701a3ee4dac6 100644 --- a/mlir/test/python/ir/array_attributes.py +++ b/mlir/test/python/ir/array_attributes.py @@ -418,6 +418,10 @@ def testGetDenseElementsIndex(): print(arr) # CHECK: True print(arr.dtype == np.int64) + array = np.array([1, 2, 3], dtype=np.int64) + attr = DenseIntElementsAttr.get(array, type=VectorType.get([3], idx_type)) + # CHECK: [1, 2, 3] + print(list(DenseIntElementsAttr(attr))) # CHECK-LABEL: TEST: testGetDenseResourceElementsAttr diff --git a/mlir/test/python/ir/attributes.py b/mlir/test/python/ir/attributes.py index dbd6bad05e01dc..3bad1d4661a004 100644 --- a/mlir/test/python/ir/attributes.py +++ b/mlir/test/python/ir/attributes.py @@ -348,6 +348,10 @@ def testDenseIntAttr(): # CHECK: i1 print(ShapedType(a.type).element_type) + shape = Attribute.parse("dense<[0, 1, 2, 3]> : vector<4xindex>") + # CHECK: attr: dense<[0, 1, 2, 3]> + print("attr:", shape) + @run def testDenseArrayGetItem():