Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

SDNQ fuse repeating bitwise ops

Browse Source
This commit is contained in:
Disty0
2025-09-06 13:18:17 +03:00
parent feaa51cb5d
commit 6a954ffcde
1 changed files with 109 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/packed_int.py
+109 -91
View File
@@ -27,28 +27,45 @@ def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, wei
def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.stack(
(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)),
torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 2), 128)),
torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128)),
torch.bitwise_or(packed_tensor[:, 3], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128)),
torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128)),
torch.bitwise_or(packed_tensor[:, 5], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 128)),
torch.bitwise_or(packed_tensor[:, 6], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 7), 128)),
packed_tensor = torch.bitwise_or(
packed_tensor[:, :7],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 7], 1),
torch.bitwise_left_shift(packed_tensor[:, 7], 2),
torch.bitwise_left_shift(packed_tensor[:, 7], 3),
torch.bitwise_left_shift(packed_tensor[:, 7], 4),
torch.bitwise_left_shift(packed_tensor[:, 7], 5),
torch.bitwise_left_shift(packed_tensor[:, 7], 6),
torch.bitwise_left_shift(packed_tensor[:, 7], 7),
),
dim=-1
),
128
),
dim=-1
)
return packed_tensor
def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 4)
packed_tensor = torch.stack(
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)),
torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
torch.bitwise_or(
packed_tensor[:, :2],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 3], 2),
torch.bitwise_left_shift(packed_tensor[:, 3], 4),
),
dim=-1
),
192
)
),
torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)).unsqueeze(-1),
),
dim=-1
)
@@ -57,25 +74,23 @@ def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.stack(
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)),
torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)),
torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)),
torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 5:8], 5)),
torch.bitwise_or(
packed_tensor[:, 3],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128),
),
),
).unsqueeze(-1),
torch.bitwise_or(
packed_tensor[:, 4],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128),
),
),
).unsqueeze(-1),
),
dim=-1
)
@@ -90,25 +105,18 @@ def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor:
def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.stack(
(
torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)),
torch.bitwise_left_shift(packed_tensor[:, 6], 6),
),
torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)),
torch.bitwise_left_shift(packed_tensor[:, 7], 6),
),
torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)),
packed_tensor = torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 3:6], 3)),
torch.cat(
(
torch.bitwise_left_shift(packed_tensor[:, 6:8], 6),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128),
)
).unsqueeze(-1),
),
),
dim=-1
dim=-1
)
)
return packed_tensor
@@ -123,7 +131,7 @@ def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor:
def pack_uint1(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().reshape(-1, 8)
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 1)),
@@ -138,15 +146,9 @@ def pack_uint1(tensor: torch.Tensor) -> torch.Tensor:
def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack(
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, 0], 127),
torch.bitwise_and(packed_tensor[:, 1], 127),
torch.bitwise_and(packed_tensor[:, 2], 127),
torch.bitwise_and(packed_tensor[:, 3], 127),
torch.bitwise_and(packed_tensor[:, 4], 127),
torch.bitwise_and(packed_tensor[:, 5], 127),
torch.bitwise_and(packed_tensor[:, 6], 127),
torch.bitwise_and(packed_tensor[:, :7], 127),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
@@ -165,7 +167,7 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
),
torch.bitwise_right_shift(packed_tensor[:, 6], 7),
),
)
).unsqueeze(-1)
),
dim=-1
).view(shape)
@@ -173,18 +175,16 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack(
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, 0], 63),
torch.bitwise_and(packed_tensor[:, 1], 63),
torch.bitwise_and(packed_tensor[:, 2], 63),
torch.bitwise_and(packed_tensor[:, 0:3], 63),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
),
torch.bitwise_right_shift(packed_tensor[:, 2], 6)
)
).unsqueeze(-1)
),
dim=-1
).view(shape)
@@ -192,28 +192,21 @@ def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack(
result_bitwise_right_shift = torch.bitwise_right_shift(packed_tensor[:, :3], 5)
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, 0], 31),
torch.bitwise_and(packed_tensor[:, 1], 31),
torch.bitwise_and(packed_tensor[:, 2], 31),
torch.bitwise_and(packed_tensor[:, 3], 31),
torch.bitwise_and(packed_tensor[:, 4], 31),
torch.bitwise_and(packed_tensor[:, :5], 31),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, 0], 5),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24),
result_bitwise_right_shift[:, :2],
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:5], 2), 24),
),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, 1], 5),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24),
),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, 2], 5),
result_bitwise_right_shift[:, 2],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8),
),
),
).unsqueeze(-1),
),
dim=-1
).view(shape)
@@ -226,21 +219,30 @@ def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack(
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, 0], 7),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7),
torch.bitwise_and(packed_tensor[:, 1], 7),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7),
torch.bitwise_and(packed_tensor[:, 2], 7),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, 0], 6),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4),
torch.bitwise_and(
torch.cat(
(
packed_tensor[:, :3],
torch.bitwise_right_shift(packed_tensor[:, :3], 3)
),
dim=-1
),
7
),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, 1], 6),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4),
torch.bitwise_right_shift(packed_tensor[:, :2], 6),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_right_shift(packed_tensor[:, 2], 4),
torch.bitwise_right_shift(packed_tensor[:, 2], 5),
),
dim=-1
),
4
),
),
),
dim=-1
@@ -249,12 +251,20 @@ def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack(
result = torch.cat(
(
torch.bitwise_and(packed_tensor, 3),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3),
torch.bitwise_right_shift(packed_tensor, 6),
torch.bitwise_and(
torch.stack(
(
packed_tensor,
torch.bitwise_right_shift(packed_tensor, 2),
torch.bitwise_right_shift(packed_tensor, 4)
),
dim=-1
),
3
),
torch.bitwise_right_shift(packed_tensor, 6).unsqueeze(-1),
),
dim=-1
).view(shape)
@@ -262,19 +272,27 @@ def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By
def unpack_uint1(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.stack(
result = torch.cat(
(
torch.bitwise_and(packed_tensor, 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 1), 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 3), 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 5), 1),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 6), 1),
torch.bitwise_right_shift(packed_tensor, 7),
torch.bitwise_and(
torch.stack(
(
packed_tensor,
torch.bitwise_right_shift(packed_tensor, 1),
torch.bitwise_right_shift(packed_tensor, 2),
torch.bitwise_right_shift(packed_tensor, 3),
torch.bitwise_right_shift(packed_tensor, 4),
torch.bitwise_right_shift(packed_tensor, 5),
torch.bitwise_right_shift(packed_tensor, 6),
),
dim=-1
),
1
),
torch.bitwise_right_shift(packed_tensor, 7).unsqueeze(-1),
),
dim=-1
).reshape(shape)
).view(shape)
return result