Skip to content

Commit

Permalink
Maybe fix 3x3?
Browse files Browse the repository at this point in the history
  • Loading branch information
FrancescoConti committed Apr 18, 2024
1 parent b55e714 commit 1e19f46
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions test/NeurekaMemoryLayout.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@


class NeurekaMemoryLayout:
_WEIGHT_BANDWIDTH = 256
_WEIGHT_BANDWIDTH_1x1 = 256
_WEIGHT_BANDWIDTH_3x3 = 288
_CIN_SUBTILE_1x1 = 32
_CIN_SUBTILE_3x3 = 32

Expand Down Expand Up @@ -77,27 +78,29 @@ def weightEncode(
weight = weight.reshape(-1, height * width * cinSubtile)
# Pad only the last dimension to weight bandwidth size
# (-1, Weight Bandwidth)
print("DEBUG", weight.shape)
weight = np.pad(
weight,
((0, 0), (0, NeurekaMemoryLayout._WEIGHT_BANDWIDTH - weight.shape[-1])),
((0, 0), (0, NeurekaMemoryLayout._WEIGHT_BANDWIDTH_3x3 - weight.shape[-1])),
"constant",
constant_values=0,
)
weightBandwidthBytes = int(np.ceil(NeurekaMemoryLayout._WEIGHT_BANDWIDTH_3x3 / 8))
elif height == 1 and width == 1:
# (cout * cinMajor, Bits * cinSubtile)
weight = weight.reshape(-1, bits * cinSubtile)
# Pad only the last dimension to weight bandwidth size
# (-1, Weight Bandwidth)
weight = np.pad(
weight,
((0, 0), (0, NeurekaMemoryLayout._WEIGHT_BANDWIDTH - weight.shape[-1])),
((0, 0), (0, NeurekaMemoryLayout._WEIGHT_BANDWIDTH_1x1 - weight.shape[-1])),
"constant",
constant_values=0,
)
weightBandwidthBytes = int(np.ceil(NeurekaMemoryLayout._WEIGHT_BANDWIDTH_1x1 / 8))

# Prepare for packing
# (-1, Weight Bandwidth Bytes, 8)
weightBandwidthBytes = int(np.ceil(NeurekaMemoryLayout._WEIGHT_BANDWIDTH / 8))
weight = np.stack(np.split(weight, weightBandwidthBytes, axis=-1), axis=-2)

# Pack bits
Expand Down

0 comments on commit 1e19f46

Please sign in to comment.