187 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import numpy as np4from mlir.ir import *5from mlir.dialects import quant6 7 8def run(f):9 print("\nTEST:", f.__name__)10 f()11 return f12 13 14# CHECK-LABEL: TEST: test_type_hierarchy15@run16def test_type_hierarchy():17 with Context():18 i8 = IntegerType.get_signless(8)19 any = Type.parse("!quant.any<i8<-8:7>:f32>")20 uniform = Type.parse("!quant.uniform<i8<-8:7>:f32, 0.99872:127>")21 per_axis = Type.parse("!quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>")22 sub_channel = Type.parse(23 "!quant.uniform<i8:f32:{0:1, 1:2}, {{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}>"24 )25 calibrated = Type.parse("!quant.calibrated<f32<-0.998:1.2321>>")26 27 assert not quant.QuantizedType.isinstance(i8)28 assert quant.QuantizedType.isinstance(any)29 assert quant.QuantizedType.isinstance(uniform)30 assert quant.QuantizedType.isinstance(per_axis)31 assert quant.QuantizedType.isinstance(sub_channel)32 assert quant.QuantizedType.isinstance(calibrated)33 34 assert quant.AnyQuantizedType.isinstance(any)35 assert quant.UniformQuantizedType.isinstance(uniform)36 assert quant.UniformQuantizedPerAxisType.isinstance(per_axis)37 assert quant.UniformQuantizedSubChannelType.isinstance(sub_channel)38 assert quant.CalibratedQuantizedType.isinstance(calibrated)39 40 assert not quant.AnyQuantizedType.isinstance(uniform)41 assert not quant.UniformQuantizedType.isinstance(per_axis)42 assert not quant.UniformQuantizedType.isinstance(sub_channel)43 assert not quant.UniformQuantizedPerAxisType.isinstance(sub_channel)44 45 46# CHECK-LABEL: TEST: test_any_quantized_type47@run48def test_any_quantized_type():49 with Context():50 i8 = IntegerType.get_signless(8)51 f32 = F32Type.get()52 any = quant.AnyQuantizedType.get(53 quant.QuantizedType.FLAG_SIGNED, i8, f32, -8, 754 )55 56 # CHECK: flags: 157 print(f"flags: {any.flags}")58 # CHECK: signed: True59 print(f"signed: {any.is_signed}")60 # CHECK: storage type: i861 print(f"storage type: {any.storage_type}")62 # CHECK: expressed type: f3263 print(f"expressed type: {any.expressed_type}")64 # CHECK: storage min: -865 print(f"storage min: {any.storage_type_min}")66 # CHECK: storage max: 767 print(f"storage max: {any.storage_type_max}")68 # CHECK: storage width: 869 print(f"storage width: {any.storage_type_integral_width}")70 # CHECK: quantized element type: !quant.any<i8<-8:7>:f32>71 print(f"quantized element type: {any.quantized_element_type}")72 # CHECK: !quant.any<i8<-8:7>:f32>73 print(any)74 assert any == Type.parse("!quant.any<i8<-8:7>:f32>")75 76 77# CHECK-LABEL: TEST: test_uniform_type78@run79def test_uniform_type():80 with Context():81 i8 = IntegerType.get_signless(8)82 f32 = F32Type.get()83 uniform = quant.UniformQuantizedType.get(84 quant.UniformQuantizedType.FLAG_SIGNED, i8, f32, 0.99872, 127, -8, 785 )86 87 # CHECK: scale: 0.9987288 print(f"scale: {uniform.scale}")89 # CHECK: zero point: 12790 print(f"zero point: {uniform.zero_point}")91 # CHECK: fixed point: False92 print(f"fixed point: {uniform.is_fixed_point}")93 # CHECK: !quant.uniform<i8<-8:7>:f32, 9.987200e-01:127>94 print(uniform)95 assert uniform == Type.parse("!quant.uniform<i8<-8:7>:f32, 0.99872:127>")96 97 98# CHECK-LABEL: TEST: test_uniform_per_axis_type99@run100def test_uniform_per_axis_type():101 with Context():102 i8 = IntegerType.get_signless(8)103 f32 = F32Type.get()104 per_axis = quant.UniformQuantizedPerAxisType.get(105 quant.QuantizedType.FLAG_SIGNED,106 i8,107 f32,108 [200, 0.99872],109 [0, 120],110 quantized_dimension=1,111 storage_type_min=quant.QuantizedType.default_minimum_for_integer(112 is_signed=True, integral_width=8113 ),114 storage_type_max=quant.QuantizedType.default_maximum_for_integer(115 is_signed=True, integral_width=8116 ),117 )118 119 # CHECK: scales: [200.0, 0.99872]120 print(f"scales: {per_axis.scales}")121 # CHECK: zero_points: [0, 120]122 print(f"zero_points: {per_axis.zero_points}")123 # CHECK: quantized dim: 1124 print(f"quantized dim: {per_axis.quantized_dimension}")125 # CHECK: fixed point: False126 print(f"fixed point: {per_axis.is_fixed_point}")127 # CHECK: !quant.uniform<i8:f32:1, {2.000000e+02,9.987200e-01:120}>128 print(per_axis)129 assert per_axis == Type.parse("!quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>")130 131 132# CHECK-LABEL: TEST: test_uniform_sub_channel_type133@run134def test_uniform_sub_channel_type():135 with Context():136 i8 = IntegerType.get_signless(8)137 f32 = F32Type.get()138 sub_channel = quant.UniformQuantizedSubChannelType.get(139 quant.QuantizedType.FLAG_SIGNED,140 i8,141 f32,142 DenseElementsAttr.get(143 np.asarray([2.0, 3.0, 4.0, 5.0], np.float32).reshape(2, 2)144 ),145 DenseElementsAttr.get(np.asarray([10, 20, 30, 40], np.int8).reshape(2, 2)),146 [0, 1],147 [1, 2],148 storage_type_min=quant.QuantizedType.default_minimum_for_integer(149 is_signed=True, integral_width=8150 ),151 storage_type_max=quant.QuantizedType.default_maximum_for_integer(152 is_signed=True, integral_width=8153 ),154 )155 156 # CHECK: quantized dimensions: [0, 1]157 print(f"quantized dimensions: {sub_channel.quantized_dimensions}")158 # CHECK: block sizes: [1, 2]159 print(f"block sizes: {sub_channel.block_sizes}")160 # CHECK: scales: {{\[}}[2. 3.]161 # CHECK: [4. 5.]]162 print(f"scales: {np.asarray(sub_channel.scales)}")163 # CHECK: zero-points: {{\[}}[10 20]164 # CHECK: [30 40]]165 print(f"zero-points: {np.asarray(sub_channel.zero_points)}")166 # CHECK: !quant.uniform<i8:f32:{0:1, 1:2}, {{\{}}{2.000000e+00:10, 3.000000e+00:20}, {4.000000e+00:30, 5.000000e+00:40}}>167 print(sub_channel)168 assert sub_channel == Type.parse(169 "!quant.uniform<i8:f32:{0:1, 1:2},{{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}>"170 )171 172 173# CHECK-LABEL: TEST: test_calibrated_type174@run175def test_calibrated_type():176 with Context():177 f32 = F32Type.get()178 calibrated = quant.CalibratedQuantizedType.get(f32, -0.998, 1.2321)179 180 # CHECK: min: -0.998181 print(f"min: {calibrated.min}")182 # CHECK: max: 1.2321183 print(f"max: {calibrated.max}")184 # CHECK: !quant.calibrated<f32<-0.998:1.232100e+00>>185 print(calibrated)186 assert calibrated == Type.parse("!quant.calibrated<f32<-0.998:1.2321>>")187