Packs data into format necessary for bitserial computation Parameters ---------- data : tvm.te.Tensor The input tvm tensor bits : int Number of bits to use for packing pack_axis : int index of the axis to pack in data bit_axis : int index of axi
(data, bits, pack_axis, bit_axis, pack_type, name="QuantizeInput")
| 27 | |
| 28 | |
| 29 | def bitpack(data, bits, pack_axis, bit_axis, pack_type, name="QuantizeInput"): |
| 30 | """Packs data into format necessary for bitserial computation |
| 31 | |
| 32 | Parameters |
| 33 | ---------- |
| 34 | data : tvm.te.Tensor |
| 35 | The input tvm tensor |
| 36 | bits : int |
| 37 | Number of bits to use for packing |
| 38 | pack_axis : int |
| 39 | index of the axis to pack in data |
| 40 | bit_axis : int |
| 41 | index of axis to place bit axis in resulting packed data |
| 42 | pack_type : str |
| 43 | Data type for packing, must be one of: ['uint8', 'uint16', 'uint32', 'uint64'] |
| 44 | name : Optional[str] = "QuantizeInput" |
| 45 | Name for the operation |
| 46 | """ |
| 47 | ishape = data.shape |
| 48 | n = len(ishape) |
| 49 | if pack_type == "uint8": |
| 50 | data_width = 8 |
| 51 | elif pack_type == "uint16": |
| 52 | data_width = 16 |
| 53 | elif pack_type == "uint32": |
| 54 | data_width = 32 |
| 55 | elif pack_type == "uint64": |
| 56 | data_width = 64 |
| 57 | |
| 58 | # Data must be in multiples of the data_width |
| 59 | assert get_const_int(ishape[pack_axis]) % data_width == 0, "Not a multiple of word size" |
| 60 | |
| 61 | shape_vec = list(ishape) |
| 62 | shape_vec[pack_axis] = shape_vec[pack_axis] // data_width |
| 63 | shape_vec.insert(bit_axis, 1) |
| 64 | bitserial_oshape = tuple(shape_vec) |
| 65 | masks = np.array([0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80]) |
| 66 | |
| 67 | # pack axis shifts if bit axis comes before |
| 68 | if bit_axis <= pack_axis: |
| 69 | pack_axis += 1 |
| 70 | |
| 71 | def _bitpack(*indices): |
| 72 | packed_data = [tvm.tirx.const(0, pack_type)] * bits |
| 73 | for k in range(data_width): |
| 74 | # Translate indices for packed data back to original |
| 75 | idx = [0] * n |
| 76 | j = 0 |
| 77 | for i in range(n + 1): |
| 78 | if i == bit_axis: |
| 79 | continue |
| 80 | if i == pack_axis: |
| 81 | idx[j] = indices[i] * data_width + k |
| 82 | else: |
| 83 | idx[j] = indices[i] |
| 84 | j += 1 |
| 85 | |
| 86 | element = data(*idx) |
no test coverage detected
searching dependent graphs…