| Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame^] | 1 | # Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved. |
| 2 | # |
| 3 | # SPDX-License-Identifier: Apache-2.0 |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the License); you may |
| 6 | # not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an AS IS BASIS, WITHOUT |
| 13 | # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
| 16 | |
| 17 | |
| 18 | # Description: |
| 19 | # Defines the basic numeric type classes for tensors. |
| 20 | |
| 21 | from .numeric_util import round_up_divide |
| 22 | import enum |
| 23 | |
| 24 | |
| 25 | class BaseType(enum.Flag): |
| 26 | Signed = 1 |
| 27 | Unsigned = 2 |
| 28 | Asymmetric = 4 |
| 29 | Int = 8 |
| 30 | SignedInt = Int | Signed |
| 31 | UnsignedInt = Int | Unsigned |
| 32 | AsymmSInt = Int | Asymmetric | Signed |
| 33 | AsymmUInt = Int | Asymmetric | Unsigned |
| 34 | Float = 16 |
| 35 | BFloat = 32 |
| 36 | Bool = 64 |
| 37 | String = 128 |
| 38 | Resource = 256 |
| 39 | Variant = 512 |
| 40 | |
| 41 | |
| 42 | class DataType: |
| 43 | """Defines a data type. Consists of a base type, and the number of bits used for this type""" |
| 44 | |
| 45 | __slots__ = "type", "bits" |
| 46 | |
| 47 | def __init__(self, type_, bits): |
| 48 | self.type = type_ |
| 49 | self.bits = bits |
| 50 | |
| 51 | def __eq__(self, other): |
| 52 | return self.type == other.type and self.bits == other.bits |
| 53 | |
| 54 | def __hash__(self): |
| 55 | return hash((self.type, self.bits)) |
| 56 | |
| 57 | def size_in_bytes(self): |
| 58 | return round_up_divide(self.bits, 8) |
| 59 | |
| 60 | def size_in_bits(self): |
| 61 | return self.bits |
| 62 | |
| 63 | def __str__(self): |
| 64 | stem, needs_format = DataType.stem_name[self.type] |
| 65 | if not needs_format: |
| 66 | return stem |
| 67 | else: |
| 68 | return stem % (self.bits,) |
| 69 | |
| 70 | __repr__ = __str__ |
| 71 | |
| 72 | stem_name = { |
| 73 | BaseType.UnsignedInt: ("uint%s", True), |
| 74 | BaseType.SignedInt: ("int%s", True), |
| 75 | BaseType.AsymmUInt: ("quint%s", True), |
| 76 | BaseType.AsymmSInt: ("qint%s", True), |
| 77 | BaseType.Float: ("float%s", True), |
| 78 | BaseType.BFloat: ("bfloat%s", True), |
| 79 | BaseType.Bool: ("bool", False), |
| 80 | BaseType.String: ("string", False), |
| 81 | BaseType.Resource: ("resource", False), |
| 82 | BaseType.Variant: ("variant", False), |
| 83 | } |
| 84 | |
| 85 | |
| 86 | # generate the standard set of data types |
| 87 | DataType.int8 = DataType(BaseType.SignedInt, 8) |
| 88 | DataType.int16 = DataType(BaseType.SignedInt, 16) |
| 89 | DataType.int32 = DataType(BaseType.SignedInt, 32) |
| 90 | DataType.int64 = DataType(BaseType.SignedInt, 64) |
| 91 | |
| 92 | DataType.uint8 = DataType(BaseType.UnsignedInt, 8) |
| 93 | DataType.uint16 = DataType(BaseType.UnsignedInt, 16) |
| 94 | DataType.uint32 = DataType(BaseType.UnsignedInt, 32) |
| 95 | DataType.uint64 = DataType(BaseType.UnsignedInt, 64) |
| 96 | |
| 97 | DataType.quint4 = DataType(BaseType.AsymmUInt, 4) |
| 98 | DataType.quint8 = DataType(BaseType.AsymmUInt, 8) |
| 99 | DataType.quint12 = DataType(BaseType.AsymmUInt, 12) |
| 100 | DataType.quint16 = DataType(BaseType.AsymmUInt, 16) |
| 101 | DataType.quint32 = DataType(BaseType.AsymmUInt, 32) |
| 102 | |
| 103 | DataType.qint4 = DataType(BaseType.AsymmSInt, 4) |
| 104 | DataType.qint8 = DataType(BaseType.AsymmSInt, 8) |
| 105 | DataType.qint12 = DataType(BaseType.AsymmSInt, 12) |
| 106 | DataType.qint16 = DataType(BaseType.AsymmSInt, 16) |
| 107 | DataType.qint32 = DataType(BaseType.AsymmSInt, 32) |
| 108 | |
| 109 | DataType.float16 = DataType(BaseType.Float, 16) |
| 110 | DataType.float32 = DataType(BaseType.Float, 32) |
| 111 | DataType.float64 = DataType(BaseType.Float, 64) |
| 112 | |
| 113 | DataType.string = DataType(BaseType.String, 64) |
| 114 | DataType.bool = DataType(BaseType.Bool, 8) |
| 115 | DataType.resource = DataType(BaseType.Resource, 8) |
| 116 | DataType.variant = DataType(BaseType.Variant, 8) |