Skip to content
This repository was archived by the owner on Oct 11, 2025. It is now read-only.

Commit 59d5333

Browse files
authored
Reapply "[mlir][python] allow DenseIntElementsAttr for index type (#118947)" (#124804)
This reapplies #118947 and adapts to nanobind.
1 parent d4e4c7a commit 59d5333

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

IRAttributes.cpp

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1372,13 +1372,19 @@ class PyDenseIntElementsAttribute
13721372

13731373
MlirType type = mlirAttributeGetType(*this);
13741374
type = mlirShapedTypeGetElementType(type);
1375-
assert(mlirTypeIsAInteger(type) &&
1376-
"expected integer element type in dense int elements attribute");
1375+
// Index type can also appear as a DenseIntElementsAttr and therefore can be
1376+
// casted to integer.
1377+
assert(mlirTypeIsAInteger(type) ||
1378+
mlirTypeIsAIndex(type) && "expected integer/index element type in "
1379+
"dense int elements attribute");
13771380
// Dispatch element extraction to an appropriate C function based on the
13781381
// elemental type of the attribute. nb::int_ is implicitly constructible
13791382
// from any C++ integral type and handles bitwidth correctly.
13801383
// TODO: consider caching the type properties in the constructor to avoid
13811384
// querying them on each element access.
1385+
if (mlirTypeIsAIndex(type)) {
1386+
return nb::int_(mlirDenseElementsAttrGetIndexValue(*this, pos));
1387+
}
13821388
unsigned width = mlirIntegerTypeGetWidth(type);
13831389
bool isUnsigned = mlirIntegerTypeIsUnsigned(type);
13841390
if (isUnsigned) {

0 commit comments

Comments
 (0)