Skip to content

[DO NOT REVIEW] PoC for U3 tight layout #37430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olpipi wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[DO NOT REVIEW] PoC for U3 tight layout #37430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 22 additions & 149 deletions src/bindings/python/src/openvino/helpers/packing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,8 @@
def pack_data(array: np.ndarray, type: Type) -> np.ndarray:
"""Represent array values as u1, u2, u3, u4, u6 or i4 openvino element type and pack them into uint8 numpy array.

For u1, u4, i4: Standard bit packing where 8 % bitwidth == 0
For u3: Transposed packing - 8 values in 3 bytes
For u6: Transposed packing - 4 values in 3 bytes
For u1, u2, u4, i4: Standard bit packing where 8 % bitwidth == 0
For u3, u6: Linear LSB-first bit-stream packing, values may straddle a byte boundary

If the number of elements in array is odd we pad them with zero value to be able to fit the bit
sequence into the uint8 array.
Expand All @@ -27,11 +26,11 @@ def pack_data(array: np.ndarray, type: Type) -> np.ndarray:
:param type: Type to interpret the array values. Type must be u1, u2, u3, u4, u6, i4, nf4 or f4e2m1.
:type type: openvino.Type
"""
# Handle u3 and u6 with special transposed packing
# Handle u3 and u6 with linear LSB-first bit-stream packing
if type == Type.u3:
return _pack_u3(array)
return _pack_linear(array, 3)
elif type == Type.u6:
return _pack_u6(array)
return _pack_linear(array, 6)

assert type in [Type.u1, Type.u2, Type.u4, Type.i4, Type.nf4, Type.f4e2m1], "Packing algorithm for the" "data types stored in 1, 2 or 4 bits"

Expand Down Expand Up @@ -61,9 +60,8 @@ def pack_data(array: np.ndarray, type: Type) -> np.ndarray:
def unpack_data(array: np.ndarray, type: Type, shape: Union[list, Shape]) -> np.ndarray:
"""Extract openvino element type values from array into new uint8/int8 array given shape.

For u1, u4, i4: Standard bit unpacking where 8 % bitwidth == 0
For u3: Transposed unpacking - 8 values from 3 bytes
For u6: Transposed unpacking - 4 values from 3 bytes
For u1, u2, u4, i4: Standard bit unpacking where 8 % bitwidth == 0
For u3, u6: Linear LSB-first bit-stream unpacking, values may straddle a byte boundary

Example: uint8 value [120] can be represented as two u4 values and be unpacked into [7, 8]
because [120] bit representation is [01111000] will be viewed as [0111, 1000],
Expand All @@ -76,11 +74,11 @@ def unpack_data(array: np.ndarray, type: Type, shape: Union[list, Shape]) -> np.
:param shape: the new shape for the unpacked array.
:type shape: Union[list, openvino.Shape]
"""
# Handle u3 and u6 with special transposed unpacking
# Handle u3 and u6 with linear LSB-first bit-stream unpacking
if type == Type.u3:
return _unpack_u3(array, shape)
return _unpack_linear(array, 3, shape)
elif type == Type.u6:
return _unpack_u6(array, shape)
return _unpack_linear(array, 6, shape)

assert type in [Type.u1, Type.u2, Type.u4, Type.i4, Type.nf4, Type.f4e2m1], "Unpacking algorithm for the" "data types stored in 1, 2 or 4 bits"
unpacked = np.unpackbits(array.view(np.uint8))
Expand All @@ -107,148 +105,23 @@ def unpack_data(array: np.ndarray, type: Type, shape: Union[list, Shape]) -> np.
return np.resize(packed, shape)


def _pack_u3(array: np.ndarray) -> np.ndarray:
"""Pack u3 values using transposed packing scheme.

8 values (each 3 bits) are packed into 3 bytes:
- Byte 0: bits [1:0] of values 0-3 (4 values * 2 bits = 8 bits)
- Byte 1: bits [1:0] of values 4-7 (4 values * 2 bits = 8 bits)
- Byte 2: bits [2] of all 8 values (8 values * 1 bit = 8 bits)
def _pack_linear(array: np.ndarray, bits: int) -> np.ndarray:
"""Pack values into a linear, LSB-first bit-stream (used for u3/u6): value i occupies bits
[i * bits, i * bits + bits), possibly straddling a byte boundary.
"""
array = array.astype(np.uint8, casting="unsafe").flatten()
# Pad to multiple of 8
pad = (-len(array)) % 8
if pad:
array = np.concatenate([array, np.zeros(pad, dtype=np.uint8)])

groups = array.reshape(-1, 8)
result = []

for group in groups:
# Extract lower 2 bits and MSB for each value
lower_bits = group & 0x03 # bits [1:0]
msb = (group >> 2) & 0x01 # bit [2]

# Pack into 3 bytes
byte0 = (lower_bits[0] << 6) | (lower_bits[1] << 4) | (lower_bits[2] << 2) | lower_bits[3]
byte1 = (lower_bits[4] << 6) | (lower_bits[5] << 4) | (lower_bits[6] << 2) | lower_bits[7]
byte2 = (msb[0] << 7) | (msb[1] << 6) | (msb[2] << 5) | (msb[3] << 4) | \
(msb[4] << 3) | (msb[5] << 2) | (msb[6] << 1) | msb[7]
bit_order_little = (array[:, None] & (1 << np.arange(bits)) > 0).astype(np.uint8)
return np.packbits(bit_order_little.flatten(), bitorder="little")

result.extend([byte0, byte1, byte2])

return np.array(result, dtype=np.uint8)


def _unpack_u3(array: np.ndarray, shape: Union[list, Shape]) -> np.ndarray:
"""Unpack u3 values using transposed unpacking scheme.

3 bytes are unpacked into 8 values (each 3 bits).
"""
array = array.view(np.uint8)
shape = list(shape)

# Process 3 bytes at a time
result = []
for i in range(0, len(array), 3):
if i + 2 >= len(array):
break
byte0, byte1, byte2 = array[i:i+3]

# Unpack lower 2 bits from first two bytes
val0 = (byte0 >> 6) & 0x03
val1 = (byte0 >> 4) & 0x03
val2 = (byte0 >> 2) & 0x03
val3 = byte0 & 0x03
val4 = (byte1 >> 6) & 0x03
val5 = (byte1 >> 4) & 0x03
val6 = (byte1 >> 2) & 0x03
val7 = byte1 & 0x03

# Unpack MSBs from third byte
msb0 = (byte2 >> 7) & 0x01
msb1 = (byte2 >> 6) & 0x01
msb2 = (byte2 >> 5) & 0x01
msb3 = (byte2 >> 4) & 0x01
msb4 = (byte2 >> 3) & 0x01
msb5 = (byte2 >> 2) & 0x01
msb6 = (byte2 >> 1) & 0x01
msb7 = byte2 & 0x01

# Combine to form 3-bit values
result.extend([
val0 | (msb0 << 2), val1 | (msb1 << 2), val2 | (msb2 << 2), val3 | (msb3 << 2),
val4 | (msb4 << 2), val5 | (msb5 << 2), val6 | (msb6 << 2), val7 | (msb7 << 2)
])

result = np.array(result, dtype=np.uint8)
return np.resize(result, shape)


def _pack_u6(array: np.ndarray) -> np.ndarray:
"""Pack u6 values using transposed packing scheme.

4 values (each 6 bits) are packed into 3 bytes:
- Byte 0: bits [3:0] of values 0-1 (2 values * 4 bits = 8 bits)
- Byte 1: bits [3:0] of values 2-3 (2 values * 4 bits = 8 bits)
- Byte 2: bits [5:4] of all 4 values (4 values * 2 bits = 8 bits)
"""
array = array.astype(np.uint8, casting="unsafe").flatten()
# Pad to multiple of 4
pad = (-len(array)) % 4
if pad:
array = np.concatenate([array, np.zeros(pad, dtype=np.uint8)])

groups = array.reshape(-1, 4)
result = []

for group in groups:
lower_bits = group & 0x0F # bits [3:0]
upper_bits = (group >> 4) & 0x03 # bits [5:4]

# Pack into 3 bytes
byte0 = (lower_bits[0] << 4) | lower_bits[1]
byte1 = (lower_bits[2] << 4) | lower_bits[3]
byte2 = (upper_bits[0] << 6) | (upper_bits[1] << 4) | (upper_bits[2] << 2) | upper_bits[3]

result.extend([byte0, byte1, byte2])

return np.array(result, dtype=np.uint8)


def _unpack_u6(array: np.ndarray, shape: Union[list, Shape]) -> np.ndarray:
"""Unpack u6 values using transposed unpacking scheme.

3 bytes are unpacked into 4 values (each 6 bits).
"""
def _unpack_linear(array: np.ndarray, bits: int, shape: Union[list, Shape]) -> np.ndarray:
"""Unpack values from a linear, LSB-first bit-stream (used for u3/u6), inverse of _pack_linear."""
array = array.view(np.uint8)
shape = list(shape)
num_values = int(np.prod(shape))

# Process 3 bytes at a time
result = []
for i in range(0, len(array), 3):
if i + 2 >= len(array):
break
byte0, byte1, byte2 = array[i:i+3]

# Unpack lower 4 bits from first two bytes
val0 = (byte0 >> 4) & 0x0F
val1 = byte0 & 0x0F
val2 = (byte1 >> 4) & 0x0F
val3 = byte1 & 0x0F

# Unpack upper 2 bits from third byte
upper0 = (byte2 >> 6) & 0x03
upper1 = (byte2 >> 4) & 0x03
upper2 = (byte2 >> 2) & 0x03
upper3 = byte2 & 0x03

# Combine to form 6-bit values
result.extend([
val0 | (upper0 << 4), val1 | (upper1 << 4),
val2 | (upper2 << 4), val3 | (upper3 << 4)
])

result = np.array(result, dtype=np.uint8)

bits_unpacked = np.unpackbits(array, bitorder="little")[: num_values * bits]
values = bits_unpacked.reshape(-1, bits)
weights = (1 << np.arange(bits)).astype(np.uint8)
result = (values * weights).sum(axis=1).astype(np.uint8)
return np.resize(result, shape)
44 changes: 9 additions & 35 deletions src/bindings/python/src/openvino/helpers/packing.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,47 +6,22 @@ import numpy
import numpy as np
import openvino._pyopenvino
__all__: list[str] = ['Shape', 'Type', 'np', 'pack_data', 'unpack_data']
def _pack_u3(array: numpy.ndarray) -> numpy.ndarray:
def _pack_linear(array: numpy.ndarray, bits: int) -> numpy.ndarray:
"""
Pack u3 values using transposed packing scheme.

8 values (each 3 bits) are packed into 3 bytes:
- Byte 0: bits [1:0] of values 0-3 (4 values * 2 bits = 8 bits)
- Byte 1: bits [1:0] of values 4-7 (4 values * 2 bits = 8 bits)
- Byte 2: bits [2] of all 8 values (8 values * 1 bit = 8 bits)

"""
def _pack_u6(array: numpy.ndarray) -> numpy.ndarray:
"""
Pack u6 values using transposed packing scheme.

4 values (each 6 bits) are packed into 3 bytes:
- Byte 0: bits [3:0] of values 0-1 (2 values * 4 bits = 8 bits)
- Byte 1: bits [3:0] of values 2-3 (2 values * 4 bits = 8 bits)
- Byte 2: bits [5:4] of all 4 values (4 values * 2 bits = 8 bits)
Pack values into a linear, LSB-first bit-stream (used for u3/u6): value i occupies bits
[i * bits, i * bits + bits), possibly straddling a byte boundary.

"""
def _unpack_u3(array: numpy.ndarray, shape: typing.Union[list, openvino._pyopenvino.Shape]) -> numpy.ndarray:
def _unpack_linear(array: numpy.ndarray, bits: int, shape: typing.Union[list, openvino._pyopenvino.Shape]) -> numpy.ndarray:
"""
Unpack u3 values using transposed unpacking scheme.

3 bytes are unpacked into 8 values (each 3 bits).

"""
def _unpack_u6(array: numpy.ndarray, shape: typing.Union[list, openvino._pyopenvino.Shape]) -> numpy.ndarray:
"""
Unpack u6 values using transposed unpacking scheme.

3 bytes are unpacked into 4 values (each 6 bits).

Unpack values from a linear, LSB-first bit-stream (used for u3/u6), inverse of _pack_linear.
"""
def pack_data(array: numpy.ndarray, type: openvino._pyopenvino.Type) -> numpy.ndarray:
"""
Represent array values as u1, u2, u3, u4, u6 or i4 openvino element type and pack them into uint8 numpy array.

For u1, u4, i4: Standard bit packing where 8 % bitwidth == 0
For u3: Transposed packing - 8 values in 3 bytes
For u6: Transposed packing - 4 values in 3 bytes
For u1, u2, u4, i4: Standard bit packing where 8 % bitwidth == 0
For u3, u6: Linear LSB-first bit-stream packing, values may straddle a byte boundary

If the number of elements in array is odd we pad them with zero value to be able to fit the bit
sequence into the uint8 array.
Expand All @@ -65,9 +40,8 @@ def unpack_data(array: numpy.ndarray, type: openvino._pyopenvino.Type, shape: ty
"""
Extract openvino element type values from array into new uint8/int8 array given shape.

For u1, u4, i4: Standard bit unpacking where 8 % bitwidth == 0
For u3: Transposed unpacking - 8 values from 3 bytes
For u6: Transposed unpacking - 4 values from 3 bytes
For u1, u2, u4, i4: Standard bit unpacking where 8 % bitwidth == 0
For u3, u6: Linear LSB-first bit-stream unpacking, values may straddle a byte boundary

Example: uint8 value [120] can be represented as two u4 values and be unpacked into [7, 8]
because [120] bit representation is [01111000] will be viewed as [0111, 1000],
Expand Down
16 changes: 4 additions & 12 deletions src/bindings/python/tests/test_runtime/test_tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -683,18 +683,14 @@ def test_tensor_from_pillow(numpy_dtype, shape):

def test_pack_unpack_u3_numerical():
"""Test u3 packing/unpacking with known numerical values."""
# 8 values packed into 3 bytes
# 8 values packed into 3 bytes, linear LSB-first bit-stream (value i at bits [3i, 3i+3))
# Decimal values: [0, 1, 2, 3, 4, 5, 6, 7]
# In binary: [000, 001, 010, 011, 100, 101, 110, 111]
data = np.array([0, 1, 2, 3, 4, 5, 6, 7], dtype=np.uint8)

packed = pack_data(data, ov.Type.u3)

# Expected packing:
# Byte 0: bits[1:0] of values 0-3 = [00, 01, 10, 11] = 0x1B
# Byte 1: bits[1:0] of values 4-7 = [00, 01, 10, 11] = 0x1B
# Byte 2: bits[2] of all 8 values = [0,0,0,0,1,1,1,1] = 0x0F
expected = np.array([0x1B, 0x1B, 0x0F], dtype=np.uint8)
expected = np.array([0x88, 0xC6, 0xFA], dtype=np.uint8)

assert len(packed) == 3
assert np.array_equal(packed, expected), f"Expected {expected.tolist()}, got {packed.tolist()}"
Expand All @@ -721,17 +717,13 @@ def test_pack_unpack_u3_edge_cases():


def test_pack_unpack_u6_numerical():
# 4 values packed into 3 bytes
# 4 values packed into 3 bytes, linear LSB-first bit-stream (value i at bits [6i, 6i+6))
# Decimal values: [0, 15, 48, 63]
# In binary: [000000, 001111, 110000, 111111]
data = np.array([0, 15, 48, 63], dtype=np.uint8)
packed = pack_data(data, ov.Type.u6)

# Expected packing:
# Byte 0: bits[3:0] of values 0-1 = [0000, 1111] = 0x0F
# Byte 1: bits[3:0] of values 2-3 = [0000, 1111] = 0x0F
# Byte 2: bits[5:4] of all 4 values = [00, 00, 11, 11] = 0x0F
expected = np.array([0x0F, 0x0F, 0x0F], dtype=np.uint8)
expected = np.array([0xC0, 0x03, 0xFF], dtype=np.uint8)

assert len(packed) == 3
assert np.array_equal(packed, expected), f"Expected {expected.tolist()}, got {packed.tolist()}"
Expand Down
Loading
Loading