From 308d8b8c6618f570166bcc7dbb87f97c04bba1b2 Mon Sep 17 00:00:00 2001 From: Rahul Kayaith Date: Fri, 21 Jan 2022 05:21:00 +0000 Subject: [PATCH] [mlir][python] 8b/16b DenseIntElements access This extends dense attribute element access to support 8b and 16b ints. Also extends the corresponding parts of the C api. Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D117731 --- mlir/include/mlir-c/BuiltinAttributes.h | 8 +++++ mlir/lib/Bindings/Python/IRAttributes.cpp | 12 +++++++ mlir/lib/CAPI/IR/BuiltinAttributes.cpp | 16 +++++++++ mlir/test/CAPI/ir.c | 11 ++++++ mlir/test/python/ir/attributes.py | 44 +++++++++++++++++++++++ 5 files changed, 91 insertions(+) diff --git a/mlir/include/mlir-c/BuiltinAttributes.h b/mlir/include/mlir-c/BuiltinAttributes.h index 5839cd3d2408..973b7e99469c 100644 --- a/mlir/include/mlir-c/BuiltinAttributes.h +++ b/mlir/include/mlir-c/BuiltinAttributes.h @@ -355,6 +355,10 @@ MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrUInt8Get( MlirType shapedType, intptr_t numElements, const uint8_t *elements); MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrInt8Get( MlirType shapedType, intptr_t numElements, const int8_t *elements); +MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrUInt16Get( + MlirType shapedType, intptr_t numElements, const uint16_t *elements); +MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrInt16Get( + MlirType shapedType, intptr_t numElements, const int16_t *elements); MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrUInt32Get( MlirType shapedType, intptr_t numElements, const uint32_t *elements); MLIR_CAPI_EXPORTED MlirAttribute mlirDenseElementsAttrInt32Get( @@ -416,6 +420,10 @@ MLIR_CAPI_EXPORTED int8_t mlirDenseElementsAttrGetInt8Value(MlirAttribute attr, intptr_t pos); MLIR_CAPI_EXPORTED uint8_t mlirDenseElementsAttrGetUInt8Value(MlirAttribute attr, intptr_t pos); +MLIR_CAPI_EXPORTED int16_t +mlirDenseElementsAttrGetInt16Value(MlirAttribute attr, intptr_t pos); +MLIR_CAPI_EXPORTED uint16_t +mlirDenseElementsAttrGetUInt16Value(MlirAttribute attr, intptr_t pos); MLIR_CAPI_EXPORTED int32_t mlirDenseElementsAttrGetInt32Value(MlirAttribute attr, intptr_t pos); MLIR_CAPI_EXPORTED uint32_t diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp index fd44ffe6ba5f..5d87641c379d 100644 --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -673,6 +673,12 @@ public: if (width == 1) { return mlirDenseElementsAttrGetBoolValue(*this, pos); } + if (width == 8) { + return mlirDenseElementsAttrGetUInt8Value(*this, pos); + } + if (width == 16) { + return mlirDenseElementsAttrGetUInt16Value(*this, pos); + } if (width == 32) { return mlirDenseElementsAttrGetUInt32Value(*this, pos); } @@ -683,6 +689,12 @@ public: if (width == 1) { return mlirDenseElementsAttrGetBoolValue(*this, pos); } + if (width == 8) { + return mlirDenseElementsAttrGetInt8Value(*this, pos); + } + if (width == 16) { + return mlirDenseElementsAttrGetInt16Value(*this, pos); + } if (width == 32) { return mlirDenseElementsAttrGetInt32Value(*this, pos); } diff --git a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp index c20548bd4759..7b718da88cee 100644 --- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp +++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp @@ -426,6 +426,16 @@ MlirAttribute mlirDenseElementsAttrInt8Get(MlirType shapedType, const int8_t *elements) { return getDenseAttribute(shapedType, numElements, elements); } +MlirAttribute mlirDenseElementsAttrUInt16Get(MlirType shapedType, + intptr_t numElements, + const uint16_t *elements) { + return getDenseAttribute(shapedType, numElements, elements); +} +MlirAttribute mlirDenseElementsAttrInt16Get(MlirType shapedType, + intptr_t numElements, + const int16_t *elements) { + return getDenseAttribute(shapedType, numElements, elements); +} MlirAttribute mlirDenseElementsAttrUInt32Get(MlirType shapedType, intptr_t numElements, const uint32_t *elements) { @@ -530,6 +540,12 @@ int8_t mlirDenseElementsAttrGetInt8Value(MlirAttribute attr, intptr_t pos) { uint8_t mlirDenseElementsAttrGetUInt8Value(MlirAttribute attr, intptr_t pos) { return unwrap(attr).cast().getValues()[pos]; } +int16_t mlirDenseElementsAttrGetInt16Value(MlirAttribute attr, intptr_t pos) { + return unwrap(attr).cast().getValues()[pos]; +} +uint16_t mlirDenseElementsAttrGetUInt16Value(MlirAttribute attr, intptr_t pos) { + return unwrap(attr).cast().getValues()[pos]; +} int32_t mlirDenseElementsAttrGetInt32Value(MlirAttribute attr, intptr_t pos) { return unwrap(attr).cast().getValues()[pos]; } diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c index d01ccaeb0e93..257d5e9b8683 100644 --- a/mlir/test/CAPI/ir.c +++ b/mlir/test/CAPI/ir.c @@ -904,6 +904,8 @@ int printBuiltinAttributes(MlirContext ctx) { int bools[] = {0, 1}; uint8_t uints8[] = {0u, 1u}; int8_t ints8[] = {0, 1}; + uint16_t uints16[] = {0u, 1u}; + int16_t ints16[] = {0, 1}; uint32_t uints32[] = {0u, 1u}; int32_t ints32[] = {0, 1}; uint64_t uints64[] = {0u, 1u}; @@ -921,6 +923,13 @@ int printBuiltinAttributes(MlirContext ctx) { MlirAttribute int8Elements = mlirDenseElementsAttrInt8Get( mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 8), encoding), 2, ints8); + MlirAttribute uint16Elements = mlirDenseElementsAttrUInt16Get( + mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 16), + encoding), + 2, uints16); + MlirAttribute int16Elements = mlirDenseElementsAttrInt16Get( + mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 16), encoding), + 2, ints16); MlirAttribute uint32Elements = mlirDenseElementsAttrUInt32Get( mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 32), encoding), @@ -956,6 +965,8 @@ int printBuiltinAttributes(MlirContext ctx) { if (mlirDenseElementsAttrGetBoolValue(boolElements, 1) != 1 || mlirDenseElementsAttrGetUInt8Value(uint8Elements, 1) != 1 || mlirDenseElementsAttrGetInt8Value(int8Elements, 1) != 1 || + mlirDenseElementsAttrGetUInt16Value(uint16Elements, 1) != 1 || + mlirDenseElementsAttrGetInt16Value(int16Elements, 1) != 1 || mlirDenseElementsAttrGetUInt32Value(uint32Elements, 1) != 1 || mlirDenseElementsAttrGetInt32Value(int32Elements, 1) != 1 || mlirDenseElementsAttrGetUInt64Value(uint64Elements, 1) != 1 || diff --git a/mlir/test/python/ir/attributes.py b/mlir/test/python/ir/attributes.py index 5f8dd0ad1183..48f2d4b3df06 100644 --- a/mlir/test/python/ir/attributes.py +++ b/mlir/test/python/ir/attributes.py @@ -292,6 +292,50 @@ def testDenseIntAttr(): print(ShapedType(a.type).element_type) +# CHECK-LABEL: TEST: testDenseIntAttrGetItem +@run +def testDenseIntAttrGetItem(): + def print_item(attr_asm): + attr = DenseIntElementsAttr(Attribute.parse(attr_asm)) + dtype = ShapedType(attr.type).element_type + try: + item = attr[0] + print(f"{dtype}:", item) + except TypeError as e: + print(f"{dtype}:", e) + + with Context(): + # CHECK: i1: 1 + print_item("dense : tensor") + # CHECK: i8: 123 + print_item("dense<123> : tensor") + # CHECK: i16: 123 + print_item("dense<123> : tensor") + # CHECK: i32: 123 + print_item("dense<123> : tensor") + # CHECK: i64: 123 + print_item("dense<123> : tensor") + # CHECK: ui8: 123 + print_item("dense<123> : tensor") + # CHECK: ui16: 123 + print_item("dense<123> : tensor") + # CHECK: ui32: 123 + print_item("dense<123> : tensor") + # CHECK: ui64: 123 + print_item("dense<123> : tensor") + # CHECK: si8: -123 + print_item("dense<-123> : tensor") + # CHECK: si16: -123 + print_item("dense<-123> : tensor") + # CHECK: si32: -123 + print_item("dense<-123> : tensor") + # CHECK: si64: -123 + print_item("dense<-123> : tensor") + + # CHECK: i7: Unsupported integer type + print_item("dense<123> : tensor") + + # CHECK-LABEL: TEST: testDenseFPAttr @run def testDenseFPAttr():