[MLIR][python bindings] Add support for DenseElementsAttr of IndexType

[rkayaith]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

based on https://github.com/llvm/llvm-project/blob/a4bdb27538c3bc5b757976e47d663e63880451e3/mlir/lib/IR/AttributeDetail.h#L33-L41 I think this should always be 64 bits wide.

[makslevental]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

I guess this is really what you're pointing at? https://github.com/llvm/llvm-project/blob/a4bdb27538c3bc5b757976e47d663e63880451e3/mlir/include/mlir/IR/BuiltinTypes.td#L307

static constexpr unsigned kInternalStorageBitWidth = 64;

Seems like that one should be kInternalStorageBitWidth = sizeof(uintptr_t) too? But maybe I'm missing something.

[makslevental]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

You can't edit inline comments? Weird...

Anyway I was going to edit/add that this is all about host arch rather than target so maybe being mindful of sizeof(uintptr_t) > 64 is moot (are there 128bit architectures on the horizon?).

[rkayaith]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

this is all about host arch rather than target so maybe being mindful of sizeof(uintptr_t) > 64 is moot (are there 128bit architectures on the horizon?).

yea that'd be my guess as well. but also.. CHERI kind of :)

[makslevental]

use uint64_t instead of uintptr_t

[makslevental]

since dense<-1> : tensor<index> is allowed, use int64_t

[makslevental]

remove accidentally included cstdint

[makslevental]

update python test

[rkayaith]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

A static_assert that IndexType::kInternalStorageBitWidth == 64 would be nice to explain where int64_t comes from

mlir/test/python/ir/array_attributes.py

374

np.int64 seems like the right type here (though I'm not actually sure what happens when the wrong type gets passed in, maybe it works somehow)

379

I know the other tests here don't, but you could avoid the escaping by using CHECK{LITERAL}: https://llvm.org/docs/CommandGuide/FileCheck.html#directive-modifiers

[makslevental]

add comment and using np.int64 on input array

[makslevental]

mlir/lib/Bindings/Python/IRAttributes.cpp

714

Can't do that without adding something like mlirIndexTypeGetStorageWidth to C API; comment will have to do.

mlir/test/python/ir/array_attributes.py

379

Leaving as is for this test but definitely will use in the future.

mlir/lib/Bindings/Python/IRAttributes.cpp

py::buffer_info accessBuffer() {
     if (mlirTypeIsAF64(elementType)) {
       // f64
       return bufferInfo<double>(shapedType);
     }
     if (mlirTypeIsAF16(elementType)) {
       // f16
       return bufferInfo<uint16_t>(shapedType, "e");
     }
+    if (mlirTypeIsAIndex(elementType)) {
+      // Same as IndexType::kInternalStorageBitWidth
+      return bufferInfo<int64_t>(shapedType);
+    }
     if (mlirTypeIsAInteger(elementType) &&
         mlirIntegerTypeGetWidth(elementType) == 32) {
       if (mlirIntegerTypeIsSignless(elementType) ||
           mlirIntegerTypeIsSigned(elementType)) {
         // i32
         return bufferInfo<int32_t>(shapedType);
       }
       if (mlirIntegerTypeIsUnsigned(elementType)) {

mlir/test/python/ir/array_attributes.py

with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)
     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(np.array(attr))
 
+
+# CHECK-LABEL: TEST: testGetDenseElementsIndex
+@run
+def testGetDenseElementsIndex():
+  with Context(), Location.unknown():
+    idx_type = IndexType.get()
+    array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
+    attr = DenseElementsAttr.get(array, type=idx_type)
+    # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xindex>
+    print(attr)
+    arr = np.array(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(arr)
+    # CHECK: True
+    print(arr.dtype == np.int64)
+