Reapply "[mlir][python] allow DenseIntElementsAttr for index type (#118947)" (#124804)

This reapplies #118947 and adapts to nanobind.

Author: mgehre-amd

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 uint64_t
+mlirDenseElementsAttrGetIndexValue(MlirAttribute attr, intptr_t pos);
 MLIR_CAPI_EXPORTED float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr,
                                                             intptr_t pos);
 MLIR_CAPI_EXPORTED double

mlir/lib/Bindings/Python/IRAttributes.cpp

@@ -1372,13 +1372,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. nb::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 nb::int_(mlirDenseElementsAttrGetIndexValue(*this, pos));
+    }
     unsigned width = mlirIntegerTypeGetWidth(type);
     bool isUnsigned = mlirIntegerTypeIsUnsigned(type);
     if (isUnsigned) {

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];
 }
+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];
 }