Skip to content

Commit

Permalink
Add wmem flag to tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@lukamac
lukamac committed Dec 21, 2024
1 parent 7db0519 commit daa5def
Show file tree
Hide file tree
Showing 7 changed files with 75 additions and 50 deletions.
12 changes: 5 additions & 7 deletions test/Ne16Weight.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,15 +20,14 @@
import numpy.typing as npt

from HeaderWriter import HeaderWriter
from NnxTestClasses import NnxWeight, WmemLiteral
from NnxTestClasses import NnxWeight, NnxWmem


class Ne16Weight(NnxWeight):
_CIN_SUBTILE = 16

@staticmethod
def encode(
weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
self, weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
) -> npt.NDArray[np.uint8]:
"""Unroll weight into expected memory format
Expand Down Expand Up @@ -76,8 +75,8 @@ def encode(
# (-1, )
return weight.flatten()

@staticmethod
def decode(
self,
weight: npt.NDArray[np.uint8],
bits: int,
cout: int,
Expand All @@ -99,11 +98,10 @@ def decode(

return weight

@staticmethod
def source_generate(
wmem: WmemLiteral, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
self, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
) -> None:
assert wmem == "tcdm", f"Invalid wmem source provided: {wmem}"
assert self.wmem == NnxWmem.tcdm, f"Unsupported weight memory destination {self.wmem}"
section = "PI_L1"

header_writer.generate_vector_files(
Expand Down
18 changes: 9 additions & 9 deletions test/NeurekaV2Weight.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,15 @@
import numpy.typing as npt

from HeaderWriter import HeaderWriter
from NnxTestClasses import NnxWeight, WmemLiteral
from NnxTestClasses import NnxWeight, NnxWmem


class NeurekaV2Weight(NnxWeight):
_WEIGHT_BANDWIDTH = 288
_CIN_SUBTILE = 32

@staticmethod
def encode(
weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
self, weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
) -> npt.NDArray[np.uint8]:
"""Unroll weight into expected memory format
Expand Down Expand Up @@ -83,8 +82,8 @@ def encode(
# (-1, )
return weight.flatten()

@staticmethod
def decode(
self,
weight: npt.NDArray[np.uint8],
bits: int,
cout: int,
Expand Down Expand Up @@ -118,16 +117,17 @@ def decode(

return weight

@staticmethod
def source_generate(
wmem: WmemLiteral, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
self, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
) -> None:
if wmem == "sram":
if self.wmem == NnxWmem.sram:
section = '__attribute__((section(".weightmem_sram")))'
elif wmem == "mram":
elif self.wmem == NnxWmem.mram:
section = '__attribute__((section(".weightmem_mram")))'
else:
elif self.wmem == NnxWmem.tcdm:
section = "PI_L1"
else:
assert False, f"Unsupported weight memory destination {self.wmem}"

header_writer.generate_vector_files(
"weight",
Expand Down
18 changes: 9 additions & 9 deletions test/NeurekaWeight.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,17 +21,16 @@
import numpy.typing as npt

from HeaderWriter import HeaderWriter
from NnxTestClasses import NnxWeight, WmemLiteral
from NnxTestClasses import NnxWeight, NnxWmem


class NeurekaWeight(NnxWeight):
_WEIGHT_BANDWIDTH = 256
_CIN_SUBTILE_1x1 = 32
_CIN_SUBTILE_3x3 = 28

@staticmethod
def encode(
weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
self, weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
) -> npt.NDArray[np.uint8]:
"""Unroll weight into expected memory format
Expand Down Expand Up @@ -118,8 +117,8 @@ def encode(
# (-1, )
return weight.flatten()

@staticmethod
def decode(
self,
weight: npt.NDArray[np.uint8],
bits: int,
cout: int,
Expand Down Expand Up @@ -157,16 +156,17 @@ def decode(

return weight

@staticmethod
def source_generate(
wmem: WmemLiteral, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
self, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
) -> None:
if wmem == "sram":
if self.wmem == NnxWmem.sram:
section = '__attribute__((section(".weightmem_sram")))'
elif wmem == "mram":
elif self.wmem == NnxWmem.mram:
section = '__attribute__((section(".weightmem_mram")))'
else:
elif self.wmem == NnxWmem.tcdm:
section = "PI_L1"
else:
assert False, f"Unsupported weight memory destination {self.wmem}"

header_writer.generate_vector_files(
"weight",
Expand Down
31 changes: 19 additions & 12 deletions test/NnxTestClasses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,14 @@
from NeuralEngineFunctionalModel import NeuralEngineFunctionalModel
from TestClasses import IntegerType, KernelShape, Padding, Stride, implies

WmemLiteral = Literal["tcdm", "sram", "mram"]

class NnxWmem(Enum):
tcdm = "tcdm"
sram = "sram"
mram = "mram"

def __str__(self) -> str:
return self.value


class NnxTestConf(BaseModel):
Expand All @@ -52,7 +59,6 @@ class NnxTestConf(BaseModel):
has_norm_quant: bool
has_bias: bool
has_relu: bool
wmem: WmemLiteral

@model_validator(mode="after") # type: ignore
def check_valid_depthwise_channels(self) -> NnxTestConf:
Expand Down Expand Up @@ -352,20 +358,22 @@ def regenerate(

class NnxWeight(ABC):

@staticmethod
def __init__(self, wmem: NnxWmem) -> None:
self.wmem = wmem

@abstractmethod
def encode(
weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
self, weight: npt.NDArray[np.uint8], bits: int, depthwise: bool = False
) -> npt.NDArray[np.uint8]:
"""Unroll weight into expected memory format
Expected input weight shape is (cout, cin, height, width).
"""
...

@staticmethod
@abstractmethod
def decode(
self,
weight: npt.NDArray[np.uint8],
bits: int,
cout: int,
Expand All @@ -376,10 +384,9 @@ def decode(
"""Reverse of encode"""
...

@staticmethod
@abstractmethod
def source_generate(
wmem: WmemLiteral, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
self, init: npt.NDArray[np.uint8], header_writer: HeaderWriter
) -> None:
"""Function implementing generation of weight's sources"""
...
Expand All @@ -390,15 +397,15 @@ class NnxTestHeaderGenerator:

def __init__(
self,
nnxWeightCls: Type[NnxWeight],
nnxWeight: NnxWeight,
headers_dir: Optional[Union[str, os.PathLike]] = None,
):
if headers_dir is None:
headers_dir = NnxTestHeaderGenerator.DEFAULT_HEADERS_DIR
self.header_writer = HeaderWriter(headers_dir)
# function that takes the weights in CoutCinK format, bitwidth, and a depthwise flag,
# and returns a numpy array of dtype=np.uint8 of data in a layout correct for the accelerator
self.nnxWeightCls = nnxWeightCls
self.nnxWeight = nnxWeight

def generate(self, test_name: str, test: NnxTest):
assert test.input is not None and test.output is not None
Expand Down Expand Up @@ -432,14 +439,14 @@ def generate(self, test_name: str, test: NnxTest):
weight_offset = -(2 ** (weight_bits - 1))
weight_out_ch, weight_in_ch, weight_ks_h, weight_ks_w = test.weight.shape
weight_data: np.ndarray = test.weight.numpy() - weight_offset
weight_init = self.nnxWeightCls.encode(
weight_init = self.nnxWeight.encode(
weight_data.astype(np.uint8),
weight_type._bits,
test.conf.depthwise,
)

self.nnxWeightCls.source_generate(
test.conf.wmem, weight_init, self.header_writer
self.nnxWeight.source_generate(
weight_init, self.header_writer
)

# Render scale
Expand Down
16 changes: 14 additions & 2 deletions test/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,14 @@
# SPDX-License-Identifier: Apache-2.0

import os
import subprocess
from typing import Union

import pydantic
import pytest

from NnxBuildFlow import CmakeBuildFlow, NnxBuildFlowName
from NnxMapping import NnxMapping, NnxName
from NnxTestClasses import NnxTest, NnxTestGenerator
from NnxTestClasses import NnxTest, NnxTestGenerator, NnxWmem
from TestClasses import implies


Expand Down Expand Up @@ -69,6 +68,14 @@ def pytest_addoption(parser):
default=NnxBuildFlowName.make,
help="Choose the build flow. Default: make",
)
parser.addoption(
"--wmem",
dest="wmem",
type=NnxWmem,
choices=list(NnxWmem),
default=NnxWmem.tcdm,
help="Choose the weight memory destination. Default: tcdm",
)


@pytest.fixture
Expand All @@ -91,6 +98,11 @@ def buildFlowName(request) -> NnxBuildFlowName:
return buildFlowName


@pytest.fixture
def wmem(request) -> NnxWmem:
return request.config.getoption("wmem")


def _find_test_dirs(path: Union[str, os.PathLike]):
return [dirpath for dirpath, _, _ in os.walk(path) if NnxTest.is_test_dir(dirpath)]

Expand Down
5 changes: 3 additions & 2 deletions test/test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

from NnxBuildFlow import NnxBuildFlowClsMapping, NnxBuildFlowName
from NnxMapping import NnxMapping, NnxName
from NnxTestClasses import NnxTest, NnxTestHeaderGenerator
from NnxTestClasses import NnxTest, NnxTestHeaderGenerator, NnxWmem

HORIZONTAL_LINE = "\n" + "-" * 100 + "\n"

Expand All @@ -33,13 +33,14 @@ def test(
nnxName: NnxName,
buildFlowName: NnxBuildFlowName,
nnxTestName: str,
wmem: NnxWmem,
):
testConfCls, weightCls = NnxMapping[nnxName]

# conftest.py makes sure the test is valid and generated
nnxTest = NnxTest.load(testConfCls, nnxTestName)

NnxTestHeaderGenerator(weightCls).generate(nnxTestName, nnxTest)
NnxTestHeaderGenerator(weightCls(wmem)).generate(nnxTestName, nnxTest)

buildFlow = NnxBuildFlowClsMapping[buildFlowName](nnxName)
buildFlow.build()
Expand Down
Loading

0 comments on commit daa5def

Please sign in to comment.