TOSA: Added RESHAPE, SLICE and CONCAT
Added support for Data layout ops
RESHAPE, SLICE and CONCAT.
-No support for bool_t
-Support limited to Rank <= 4 and N = 1
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: I487ac494b6506a2a6ba947ee758aa193194dd796
diff --git a/ethosu/vela/graph_optimiser_util.py b/ethosu/vela/graph_optimiser_util.py
index d01d4a1..8095f08 100644
--- a/ethosu/vela/graph_optimiser_util.py
+++ b/ethosu/vela/graph_optimiser_util.py
@@ -212,6 +212,29 @@
cons.set_input_tensor(ifm, ifm_idx)
+def move_splitsliceread_to_consumer(op, cons_op):
+ assert op.type == Op.SplitSliceRead
+
+ if cons_op.ifm == op.ofm:
+ cons_op.read_offsets[0] = op.read_offsets[0]
+ cons_op.read_shapes[0] = op.read_shapes[0]
+ cons_op.set_input_tensor(op.ifm, cons_op.type.info.indices.ifms[0])
+ cons_op.ifm_shapes[0] = op.ifm_shapes[0]
+ elif cons_op.type.is_binary_elementwise_op() and cons_op.ifm2 == op.ofm:
+ cons_op.read_offsets[1] = op.read_offsets[0]
+ cons_op.read_shapes[1] = op.read_shapes[0]
+ cons_op.set_input_tensor(op.ifm, cons_op.type.info.indices.ifms[1])
+ cons_op.ifm_shapes[1] = op.ifm_shapes[0]
+
+ if "skirt" in cons_op.attrs:
+ assert cons_op.attrs["explicit_padding"] == cons_op.attrs["skirt"]
+ cons_op.attrs["skirt"] = None
+ cons_op.attrs["force_padding"] = True
+ op.ofm.consumer_list.remove(cons_op)
+ op.ofm.ops = []
+ op.ifm.consumer_list.remove(op)
+
+
def check_reshapes(op, arch):
if op.run_on_npu and op.type == Op.Reshape:
ofm = op.ofm
diff --git a/ethosu/vela/operation_util.py b/ethosu/vela/operation_util.py
index 4a4fd33..0fbed46 100644
--- a/ethosu/vela/operation_util.py
+++ b/ethosu/vela/operation_util.py
@@ -44,6 +44,12 @@
return op
+def create_add_nop(name: str) -> Operation:
+ op = Operation(Op.Add, name)
+ op.run_on_npu = True
+ return op
+
+
def create_depthwise_maxpool(
name: str,
ifm: Tensor,
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index ef39aea..7526f46 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -34,6 +34,7 @@
from .graph_optimiser_util import calc_explicit_padding
from .graph_optimiser_util import convert_depthwise_to_conv
from .graph_optimiser_util import fix_sg_input_output
+from .graph_optimiser_util import move_splitsliceread_to_consumer
from .graph_optimiser_util import needed_total_padding
from .graph_optimiser_util import set_ifm_ofm_op_shapes
from .graph_optimiser_util import set_tensor_equivalence
@@ -193,24 +194,7 @@
):
# SplitSliceRead can be performed by tensor consumer
cons_op = op.ofm.consumer_list[0]
- if cons_op.ifm == op.ofm:
- cons_op.read_offsets[0] = op.read_offsets[0]
- cons_op.read_shapes[0] = op.read_shapes[0]
- cons_op.set_input_tensor(op.ifm, cons_op.type.info.indices.ifms[0])
- cons_op.ifm_shapes[0] = op.ifm_shapes[0]
- elif cons_op.type.is_binary_elementwise_op() and cons_op.ifm2 == op.ofm:
- cons_op.read_offsets[1] = op.read_offsets[0]
- cons_op.read_shapes[1] = op.read_shapes[0]
- cons_op.set_input_tensor(op.ifm, cons_op.type.info.indices.ifms[1])
- cons_op.ifm_shapes[1] = op.ifm_shapes[0]
-
- if "skirt" in cons_op.attrs:
- assert cons_op.attrs["explicit_padding"] == cons_op.attrs["skirt"]
- cons_op.attrs["skirt"] = None
- cons_op.attrs["force_padding"] = True
- op.ofm.consumer_list.remove(cons_op)
- op.ofm.ops = []
- op.ifm.consumer_list.remove(op)
+ move_splitsliceread_to_consumer(op, cons_op)
else:
avgpool_op = create_avgpool_nop(op.name + "_avgpool")
avgpool_op.add_input_tensor(op.ifm)
diff --git a/ethosu/vela/tosa_graph_optimiser.py b/ethosu/vela/tosa_graph_optimiser.py
index 169da40..f3cddad 100644
--- a/ethosu/vela/tosa_graph_optimiser.py
+++ b/ethosu/vela/tosa_graph_optimiser.py
@@ -22,14 +22,17 @@
from .graph_optimiser_util import bypass_reshape_and_squeeze_ops
from .graph_optimiser_util import calc_explicit_padding
from .graph_optimiser_util import convert_depthwise_to_conv
-from .graph_optimiser_util import fix_sg_input_output
+from .graph_optimiser_util import move_splitsliceread_to_consumer
from .graph_optimiser_util import needed_total_padding
from .graph_optimiser_util import set_ifm_ofm_op_shapes
from .graph_optimiser_util import set_tensor_equivalence
from .operation import ExplicitScaling
from .operation import NpuBlockType
from .operation import Op
+from .operation_util import create_add_nop
from .operation_util import create_avgpool_nop
+from .shape4d import Shape4D
+from .tensor import create_const_tensor
def replace_rescale_with_avg_pool(rescale_op):
@@ -103,12 +106,157 @@
removed = True
if not removed:
- print("Cannot remove Transpose, and handling of Transpose is not supported")
+ print("Warning: Cannot remove Transpose, and handling of Transpose is not supported")
assert False
return op
+# TODO can we change to add for both TFLite and TOSA?
+def insert_add_copy_op_after_tens(tens):
+ tens_cons_list_copy = tens.consumer_list.copy()
+ copy_tens = tens.clone()
+
+ name = tens.name + "_add"
+ ifm2 = create_const_tensor(
+ name + "_zero_scalar",
+ [1],
+ copy_tens.dtype,
+ [0],
+ copy_tens.dtype.as_numpy_type(),
+ quantization=copy_tens.quantization,
+ )
+ copy_op = create_add_nop(name)
+ copy_op.add_input_tensor(tens)
+ copy_op.add_input_tensor(ifm2)
+ copy_op.set_output_tensor(copy_tens)
+ copy_op.set_ifm_ofm_shapes()
+ copy_op.run_on_npu = True
+
+ # Set copy_ifm consumers
+ for tens_cons in tens_cons_list_copy:
+ if tens_cons is not None:
+ for ifm_idx, cons_inp in enumerate(tens_cons.inputs):
+ if cons_inp == tens:
+ tens_cons.set_input_tensor(copy_tens, ifm_idx)
+
+ DebugDatabase.add_optimised(tens.ops[0], copy_op)
+
+
+def fix_sg_input_output_tosa(op, arch, nng):
+ if not op.run_on_npu or op.type != Op.Reshape:
+ return op
+
+ # For the Reshape operators we want to remove, tensors are removed.
+ # But in order to to do this, they cannot be outputs of the sg,
+ # this need to be fixed prior to the removal.
+ # Solution is to add a copy op, to maintain the original tensor.
+ # This is also valid when reshape ifm/ofm is produced respectively
+ # consumed by CPU
+
+ # Check if operator ifm/ofm are sg ifm/ofm
+ ifm_is_sg_ifm = op.ifm.ops[0].type in (Op.Placeholder, Op.SubgraphInput, Op.Const)
+ ifm_is_sg_ofm = any(ifm_cons is None for ifm_cons in op.ifm.consumer_list)
+ ofm_is_sg_ofm = any(ofm_cons is None for ofm_cons in op.ofm.consumer_list)
+ # Check if ifm/ofm is produced repectivly consumed by CPU
+ ifm_is_cpu_produced = any(ifm_prod is not None and not ifm_prod.run_on_npu for ifm_prod in op.ifm.ops)
+ ofm_is_cpu_consumed = any(ofm_cons is not None and not ofm_cons.run_on_npu for ofm_cons in op.ofm.consumer_list)
+
+ if (ifm_is_sg_ofm or ifm_is_sg_ifm or ifm_is_cpu_produced) and (ofm_is_sg_ofm or ofm_is_cpu_consumed):
+ # Both ifm and ofm need to persist, but only ifm need a copy, in order to remove the Reshape
+ insert_add_copy_op_after_tens(op.ifm)
+
+ return op
+
+
+def create_add_for_concat(concat_op, name, ifm, ifm_shape: Shape4D, write_offset: Shape4D):
+ """Creates an add op for the given concat op/input feature map"""
+ ofm = concat_op.ofm
+ ifm2 = create_const_tensor(
+ name + "_zero_scalar", [1], ofm.dtype, [0], ofm.dtype.as_numpy_type(), quantization=ofm.quantization
+ )
+ add_op = create_add_nop(name)
+
+ add_op.inputs = [ifm, ifm2]
+ add_op.outputs = [ofm]
+ add_op.write_offset = write_offset
+ add_op.write_shape = ifm_shape
+ ofm.ops.append(add_op)
+ DebugDatabase.add_optimised(concat_op, add_op)
+ add_op.ifm_shapes.append(ifm_shape)
+ add_op.ifm_shapes.append(Shape4D(ifm2.shape))
+ add_op.ofm_shapes.append(concat_op.ofm_shapes[0])
+ add_op.memory_function = Op.ConcatSliceWrite
+ return add_op
+
+
+# TODO Could be further optimized checking the type of the consumer,
+# rather than just mimic the TFLite behaviour depending on type.
+# TOSA bool_t not considered yet
+def remove_splitsliceread(op, arch):
+
+ if op.type == Op.SplitSliceRead:
+ # Check if it is possible to put the SplitSliceRead on the tensor consumer, or if an avgpool need to be inserted
+ if (
+ len(op.ofm.consumer_list) == 1
+ and op.ofm.consumer_list[0] is not None
+ and op.ofm.consumer_list[0].run_on_npu
+ and op.ofm.consumer_list[0].type != Op.Reshape
+ and op.ofm_shapes[0] == Shape4D.from_list(op.ofm.shape)
+ and op.ofm.dtype in (DataType.uint8, DataType.int8, DataType.int16)
+ ):
+ # SplitSliceRead can be performed by tensor consumer
+ cons_op = op.ofm.consumer_list[0]
+ move_splitsliceread_to_consumer(op, cons_op)
+ else:
+ name = op.name + "_add"
+ ofm = op.ofm
+ ifm2 = create_const_tensor(
+ name + "_zero_scalar", [1], ofm.dtype, [0], ofm.dtype.as_numpy_type(), quantization=ofm.quantization
+ )
+ add_op = create_add_nop(name)
+ add_op.inputs = [op.ifm, ifm2]
+ add_op.outputs = [ofm]
+ op.ofm.ops.remove(op)
+ op.ofm.ops.append(add_op)
+ add_op.ifm_shapes.append(op.ifm_shapes[0])
+ add_op.ifm_shapes.append(Shape4D(ifm2.shape))
+ add_op.ofm_shapes.append(op.ofm_shapes[0])
+ add_op.read_offsets[0] = op.read_offsets[0]
+ add_op.read_shapes[0] = op.read_shapes[0]
+
+ op.ifm.consumer_list.remove(op)
+ DebugDatabase.add_optimised(op, add_op)
+
+
+def rewrite_concat_ops(op, arch):
+ if not op.run_on_npu or not op.type == Op.Concat:
+ return
+
+ axis_4D = 0
+ ofm = op.ofm
+ ofm.ops = []
+ offset = 0
+
+ inputs = op.inputs
+ axis = op.attrs["axis"]
+
+ for idx, inp in enumerate(inputs):
+ op.ifm_shapes[idx] = Shape4D(inp.shape)
+ if axis >= 0:
+ axis_4D = axis + (4 - len(inp.shape))
+ else:
+ axis_4D = axis
+ write_offset = [0, 0, 0, 0]
+ write_offset[axis_4D] = offset
+ concat_end = offset + op.ifm_shapes[idx][axis_4D]
+ create_add_for_concat(op, op.name + str(idx) + "_add", inp, op.ifm_shapes[idx], Shape4D.from_list(write_offset))
+ offset = concat_end
+ assert ofm.shape[axis] == offset
+
+ return op
+
+
def remove_reshapes(op, arch):
if op.run_on_npu and op.type == Op.Reshape:
bypass_reshape_and_squeeze_ops(op)
@@ -271,9 +419,14 @@
# Handle sg input output
for idx, sg in enumerate(nng.subgraphs):
nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
- nng, sg, arch, [], [fix_sg_input_output], rewrite_unsupported=False,
+ nng, sg, arch, [], [fix_sg_input_output_tosa], rewrite_unsupported=False,
)
+ # Rewrite concat ops
+ for idx, sg in enumerate(nng.subgraphs):
+ rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [rewrite_concat_ops])
+ sg.refresh_after_modification()
+
# Removal of reshapes
for sg in nng.subgraphs:
rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [remove_reshapes])
@@ -293,6 +446,12 @@
nng, sg, arch, [], [rewrite_activation, add_padding_fields],
)
+ # Removal of Slice, need to be done after optimisation has been performed,
+ # since ifm/ofm_shapes are of importance to this function
+ for sg in nng.subgraphs:
+ rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [remove_splitsliceread])
+ sg.refresh_after_modification()
+
# Post-processing step 2
for idx, sg in enumerate(nng.subgraphs):
nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(nng, sg, arch, [], [fixup_quantization],)
diff --git a/ethosu/vela/tosa_mapping.py b/ethosu/vela/tosa_mapping.py
index 377f455..6efc479 100644
--- a/ethosu/vela/tosa_mapping.py
+++ b/ethosu/vela/tosa_mapping.py
@@ -154,7 +154,7 @@
"TransposeConv2dAttribute", (("outpad", is_vec), ("stride", is_vec), ("dilation", is_vec), ("out_shape", is_vec))
)
relun_attrs = AttrSerializer("ReluNAttribute", ("max_int"))
-axis_attrs = AttrSerializer("AxisAttribute", ("axis"))
+axis_attrs = AttrSerializer("AxisAttribute", ("axis",))
reshape_attrs = AttrSerializer("ReshapeAttribute", (("shape", is_vec),))
slice_attrs = AttrSerializer("SliceAttribute", (("begin", is_vec), ("size", is_vec)))
tile_attrs = AttrSerializer("TileAttribute", (("multiplies", is_vec),))
@@ -218,10 +218,8 @@
TosaOp.REDUCE_MIN,
TosaOp.REDUCE_PRODUCT,
TosaOp.REDUCE_SUM,
- TosaOp.CONCAT,
TosaOp.PAD,
TosaOp.REVERSE,
- TosaOp.SLICE,
TosaOp.TILE,
TosaOp.GATHER,
TosaOp.SCATTER,
@@ -241,7 +239,7 @@
TOSA_IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
# TOSA_CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
# TOSA_TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
-# TOSA_CONCAT_INDICES = TensorIndices([1, 2], [], [])
+TOSA_CONCAT_INDICES = TensorIndices([1, 2], [], [])
# TOSA_SPLIT_IFM_INDICES = TensorIndices([1], [], [])
# TOSA_BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])
@@ -299,11 +297,11 @@
# TODO TosaOp.REDUCE_MIN
# TODO TosaOp.REDUCE_PRODUCT
# TODO TosaOp.REDUCE_SUM
- # TODO TosaOp.CONCAT
+ TosaOp.CONCAT: (Op.Concat, axis_attrs, None, TOSA_CONCAT_INDICES),
# TODO TosaOp.PAD
TosaOp.RESHAPE: (Op.Reshape, reshape_attrs, None, TOSA_IFM_INDICES),
# TODO TosaOp.REVERSE
- # TODO TosaOp.SLICE
+ TosaOp.SLICE: (Op.SplitSliceRead, slice_attrs, None, TOSA_IFM_INDICES),
# TODO TosaOp.TILE
TosaOp.TRANSPOSE: (Op.Transpose, None, None, TOSA_IFM_IFM2_INDICES),
# TODO TosaOp.GATHER
diff --git a/ethosu/vela/tosa_reader.py b/ethosu/vela/tosa_reader.py
index 2925ab4..94ba350 100644
--- a/ethosu/vela/tosa_reader.py
+++ b/ethosu/vela/tosa_reader.py
@@ -29,6 +29,7 @@
from .reader_util import clone_and_reshape_tensor
from .reader_util import decode_str
from .reader_util import fixup_tensors
+from .shape4d import Shape4D
from .tensor import QuantizationParameters
from .tensor import shape_num_elements
from .tensor import Tensor
@@ -186,6 +187,9 @@
op.rescale = [1, shift]
if op.type.is_depthwise_conv2d_op():
op.attrs["depth_multiplier"] = op.weights.shape[3]
+ if op.type == Op.SplitSliceRead:
+ op.read_offsets[0] = Shape4D.from_list(list(op.attrs["begin"]), 0)
+ op.read_shapes[0] = op.attrs["size"]
elif op.type == Op.Transpose:
op.attrs["perms"] = perms.values
diff --git a/ethosu/vela/tosa_supported_operators.py b/ethosu/vela/tosa_supported_operators.py
index d7a1ebc..c619f2f 100644
--- a/ethosu/vela/tosa_supported_operators.py
+++ b/ethosu/vela/tosa_supported_operators.py
@@ -38,7 +38,7 @@
fc_vector_products = set((Op.FullyConnected,))
mac_main_ops = convolution_like_ops | pooling_ops | fc_vector_products
- memory_only_ops = set((Op.Reshape, Op.Transpose,))
+ memory_only_ops = set((Op.Reshape, Op.Transpose, Op.Concat, Op.SplitSliceRead,))
binary_elem_wise_add_mul_sub = set((Op.Add, Op.Mul, Op.RescaleMul, Op.Sub,))
@@ -53,7 +53,7 @@
# Supported data types
# TODO will differ compared to TensorFlow Lite, currently set to the same
- supported_op_dtypes = set((DataType.uint8, DataType.int8, DataType.int16, DataType.int32))
+ supported_op_dtypes = set((DataType.uint8, DataType.int8, DataType.int16, DataType.int32)) # TODO add bool
def __init__(self):
# Setup the generic constraints. Note: the order matters
@@ -109,7 +109,7 @@
valid = op.ifm.ops and op.ifm.ops[0].type == Op.Const
return valid, "Op has ifm with non-constant data"
- # TODO duplicates tflite_supported operators, but support for depth multiplier should be added at a later stage
+ # TODO duplicates TFLite_supported operators, but support for depth multiplier should be added at a later stage
@staticmethod
def constraint_depth_multiplier(op):
"For depth multipliers > 1, IFM channels must be 1 and OFM channels must be equal to the depth multiplier"