shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

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

Browse Source 
Differential Revision: https://reviews.llvm.org/D149690
This commit is contained in:
max 2023-05-02 15:52:04 -05:00
parent 218b50a606
commit ef1b735dfb
2 changed files with 21 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlir/lib/Bindings/Python/IRAttributes.cpp
View File

@ -710,6 +710,10 @@ public:
// 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) ||

17
mlir/test/python/ir/array_attributes.py
View File

@ -365,3 +365,20 @@ def testGetDenseElementsUI64():
# 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)