[Mlir-commits] [mlir] [mlir][python] allow DenseIntElementsAttr for index type (PR #118947)
Matthias Gehre
llvmlistbot at llvm.org
Mon Dec 16 15:12:35 PST 2024
https://github.com/mgehre-amd updated https://github.com/llvm/llvm-project/pull/118947
>From ae6df1b8d3689783062d7adf626c5dd3fff77f3e Mon Sep 17 00:00:00 2001
From: Christopher McGirr <christopher.mcgirr at amd.com>
Date: Mon, 29 Jan 2024 15:02:29 +0100
Subject: [PATCH 1/2] [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 <tiago.trevisanjost at amd.com>
Co-authored-by: Matthias Gehre <matthias.gehre at amd.com>
---
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 7c8c84e55b962f..ece4e7d5e4ac16 100644
--- a/mlir/include/mlir-c/BuiltinAttributes.h
+++ b/mlir/include/mlir-c/BuiltinAttributes.h
@@ -556,6 +556,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 cc9532f4e33b2c..f05f9f02e50faa 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1172,13 +1172,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 11d1ade552f5a2..2a44533d781856 100644
--- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
@@ -758,6 +758,9 @@ int64_t mlirDenseElementsAttrGetInt64Value(MlirAttribute attr, intptr_t pos) {
uint64_t mlirDenseElementsAttrGetUInt64Value(MlirAttribute attr, intptr_t pos) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<uint64_t>()[pos];
}
+size_t mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, intptr_t pos) {
+ return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<size_t>()[pos];
+}
float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr, intptr_t pos) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<float>()[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 256a69a939658d..ef1d835fc64012 100644
--- a/mlir/test/python/ir/array_attributes.py
+++ b/mlir/test/python/ir/array_attributes.py
@@ -572,6 +572,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 00c3e1b4decdb7..2f3c4460d3f590 100644
--- a/mlir/test/python/ir/attributes.py
+++ b/mlir/test/python/ir/attributes.py
@@ -366,6 +366,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():
>From c08b2381891f98c47b4e21cf5a888ccec654f30f Mon Sep 17 00:00:00 2001
From: Matthias Gehre <matthias.gehre at amd.com>
Date: Tue, 17 Dec 2024 00:12:17 +0100
Subject: [PATCH 2/2] size_t -> uint64_t
---
mlir/include/mlir-c/BuiltinAttributes.h | 4 ++--
mlir/lib/CAPI/IR/BuiltinAttributes.cpp | 4 ++--
2 files changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/include/mlir-c/BuiltinAttributes.h b/mlir/include/mlir-c/BuiltinAttributes.h
index ece4e7d5e4ac16..1d0edf9ea809d2 100644
--- a/mlir/include/mlir-c/BuiltinAttributes.h
+++ b/mlir/include/mlir-c/BuiltinAttributes.h
@@ -556,8 +556,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 uint64_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/CAPI/IR/BuiltinAttributes.cpp b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
index 2a44533d781856..8d57ab6b59e79c 100644
--- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
@@ -758,8 +758,8 @@ int64_t mlirDenseElementsAttrGetInt64Value(MlirAttribute attr, intptr_t pos) {
uint64_t mlirDenseElementsAttrGetUInt64Value(MlirAttribute attr, intptr_t pos) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<uint64_t>()[pos];
}
-size_t mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, intptr_t pos) {
- return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<size_t>()[pos];
+uint64_t mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, intptr_t pos) {
+ return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<uint64_t>()[pos];
}
float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr, intptr_t pos) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<float>()[pos];
More information about the Mlir-commits
mailing list