diff options
author | Ayaan Masood <Ayaan.Masood@arm.com> | 2022-04-21 14:28:03 +0100 |
---|---|---|
committer | Ayaan Masood <Ayaan.Masood@arm.com> | 2022-04-21 14:28:03 +0100 |
commit | a2ec5aa747633da72b6310ce7e5552c39f7f54bb (patch) | |
tree | 04886e488a7059064653c836c32f0e0ec69a4a74 /ethosu/vela/tensor.py | |
parent | f9267da3ad6251a7e04f501218380ac9a89953b7 (diff) | |
download | ethos-u-vela-a2ec5aa747633da72b6310ce7e5552c39f7f54bb.tar.gz |
MLBEDSW-5384 FC layers run on NPU if underlying shape is 2D
*Added generic function which checks if underlying shape of
FullyConnected operation is 2D and performs shape reduction
*Fully connected operation >2 dimensions now run on NPU if the above
case is satisfied
*constraint_fc_output_2d and rewrite_fully_connected_input refactored
*Added unit test to confirm this functionality
Signed-off-by: Ayaan Masood <Ayaan.Masood@arm.com>
Change-Id: I0e29c767e5b84841eb53bbc44464b36a454f7b38
Diffstat (limited to 'ethosu/vela/tensor.py')
-rw-r--r-- | ethosu/vela/tensor.py | 13 |
1 files changed, 13 insertions, 0 deletions
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py index 38b0e430..e9815845 100644 --- a/ethosu/vela/tensor.py +++ b/ethosu/vela/tensor.py @@ -823,6 +823,19 @@ class Tensor: else: return self.values.item(0) + def get_shape_as_2d(self, dimension_2_size: int) -> Optional[Shape4D]: + + elms = self.elements() + dimension_1_size = elms // dimension_2_size + # Checks if the reduction works and shape is not 1D + is_reducible = dimension_1_size * dimension_2_size == elms and not (len(self.shape) == 1) + + new_shape = None + if is_reducible: + new_shape = Shape4D([dimension_1_size, 1, 1, dimension_2_size]) + + return new_shape + def __lt__(self, other: "Tensor") -> bool: return self.equivalence_id < other.equivalence_id |