TOSA: Add support for Identity operation
Added support for Identity operation.
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: If00b30528932f7531807ce3914d6c1875ab72fa4
diff --git a/ethosu/vela/graph_optimiser_util.py b/ethosu/vela/graph_optimiser_util.py
index d2d3d83..73fbf6c 100644
--- a/ethosu/vela/graph_optimiser_util.py
+++ b/ethosu/vela/graph_optimiser_util.py
@@ -35,6 +35,7 @@
Op.QuantizedReshape,
Op.Squeeze,
Op.ExpandDims,
+ Op.Identity,
)
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index b426792..1e733d5 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -189,7 +189,7 @@
GreaterEqual = OperatorInfo()
HardSwish = OperatorInfo(indices=NNG_IFM_INDICES)
HashtableLookup = OperatorInfo()
- Identity = OperatorInfo()
+ Identity = OperatorInfo(indices=NNG_IFM_INDICES)
If = OperatorInfo()
L2Norm = OperatorInfo()
L2Pool2D = OperatorInfo()
diff --git a/ethosu/vela/tosa_graph_optimiser.py b/ethosu/vela/tosa_graph_optimiser.py
index d32955d..954ac68 100644
--- a/ethosu/vela/tosa_graph_optimiser.py
+++ b/ethosu/vela/tosa_graph_optimiser.py
@@ -182,7 +182,7 @@
def fix_sg_input_output_tosa(op, arch, nng):
- if not op.run_on_npu or op.type != Op.Reshape:
+ if not op.run_on_npu or op.type not in (Op.Reshape, Op.Identity):
return op
# For the Reshape operators we want to remove, tensors are removed.
@@ -306,8 +306,8 @@
assert op.ofm_shapes[0][axis_4D] == offset
-def remove_reshapes(op, arch):
- if op.run_on_npu and op.type == Op.Reshape:
+def remove_memory_ops(op, arch):
+ if op.run_on_npu and op.type in (Op.Reshape, Op.Identity):
bypass_memory_only_ops(op)
@@ -820,7 +820,7 @@
# Removal of reshapes
for sg in nng.subgraphs:
- rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [remove_reshapes])
+ rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [remove_memory_ops])
sg.refresh_after_modification()
# Decomposing of elementwise
diff --git a/ethosu/vela/tosa_mapping.py b/ethosu/vela/tosa_mapping.py
index f80a915..75f1c9c 100644
--- a/ethosu/vela/tosa_mapping.py
+++ b/ethosu/vela/tosa_mapping.py
@@ -223,7 +223,6 @@
TosaOp.SCATTER,
TosaOp.RESIZE,
TosaOp.CAST,
- TosaOp.IDENTITY,
TosaOp.CUSTOM,
TosaOp.COND_IF,
TosaOp.WHILE_LOOP,
@@ -316,7 +315,7 @@
# TODO TosaOp.CAST
TosaOp.RESCALE: (Op.Rescale, rescale_attrs, None, TOSA_IFM_INDICES),
TosaOp.CONST: (Op.Const, None, None, TOSA_NO_INDICES),
- # TODO TosaOp.IDENTITY
+ TosaOp.IDENTITY: (Op.Identity, None, None, TOSA_IFM_INDICES),
# TODO TosaOp.CUSTOM
# TODO TosaOp.COND_IF
# TODO TosaOp.WHILE_LOOP
diff --git a/ethosu/vela/tosa_supported_operators.py b/ethosu/vela/tosa_supported_operators.py
index 2692c05..5a85b0e 100644
--- a/ethosu/vela/tosa_supported_operators.py
+++ b/ethosu/vela/tosa_supported_operators.py
@@ -46,13 +46,21 @@
activation_ops = relu_ops | set((Op.Table,))
pad_ops = set((Op.Pad,))
- rank_unlimited_ops = set((Op.Concat, Op.Reshape))
+ rank_unlimited_ops = set((Op.Concat, Op.Reshape, Op.Identity))
rank6_limited_ops = elem_wise_ops
batch_enabled_ops = rank6_limited_ops | rank_unlimited_ops
large_tens_dims_enabled_ops = batch_enabled_ops | set((Op.SplitSliceRead,))
npu_post_ops = activation_ops
- supported_operators = mac_main_ops | type_conversion_ops | npu_post_ops | memory_only_ops | elem_wise_ops | pad_ops
+ supported_operators = (
+ mac_main_ops
+ | type_conversion_ops
+ | npu_post_ops
+ | memory_only_ops
+ | elem_wise_ops
+ | pad_ops
+ | set((Op.Identity,))
+ )
# Supported data types
# TODO will differ compared to TensorFlow Lite, currently set to the same