MLBEDSW-2637 Utilise new tensor and operator funcs
add_input_tensor, set_output_tensor, create_const_tensor and
create_reshape_tensor have recently been added.
This replaces all found existing instances with these new helper
functions
Signed-off-by: Michael McGeagh <michael.mcgeagh@arm.com>
Change-Id: If33be8dbf237b2087b562b03cdeb51da1f99a786
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index a9d5cce..582924c 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -28,6 +28,8 @@
from .operation import NpuBlockType
from .operation import Operation
from .softmax import SoftMax
+from .tensor import create_const_tensor
+from .tensor import create_reshape_tensor
from .tensor import QuantizationParameters
from .tensor import Tensor
@@ -84,7 +86,6 @@
tens.ops = []
new_op = Operation("SplitSliceRead", split_op.name)
new_op.inputs = [inp]
- new_op.outputs = [tens]
# For Split the offset cannot be extracted from the tensor so it has to
# be calculated from the index of the output tensor
@@ -102,7 +103,7 @@
new_op.attrs["split_start"] = offset_start
new_op.attrs["split_end"] = offset_end
new_op.run_on_npu = True
- tens.ops.append(new_op)
+ new_op.set_output_tensor(tens)
return tens
@@ -168,14 +169,12 @@
if len(op.inputs) < 4:
# Add bias/scale tensor filled with zeros
- scale_op = Operation("Const", op.name + "_bias")
scale_tens = Tensor([weight_sets], DataType.int32, op.name + "_bias_tens")
scale_tens.values = [0] * weight_sets
scale_tens.quant_values = [0] * weight_sets
- scale_tens.ops = [scale_op]
- scale_op.outputs = [scale_tens]
- scale_tens.consumer_list = [op]
- op.inputs.append(scale_tens)
+ scale_op = Operation("Const", op.name + "_bias")
+ scale_op.set_output_tensor(scale_tens)
+ op.add_input_tensor(scale_tens)
# Update strides
op.attrs.update({"stride_w": 1, "stride_h": 1, "strides": (1, 1, 1, 1)})
@@ -199,8 +198,7 @@
tens.quantization.zero_point = 0
tens.consumer_list = [op]
tens_op = op.inputs[1].ops[0]
- tens_op.outputs = [tens]
- tens.ops = [tens_op]
+ tens_op.set_output_tensor(tens)
# Set the add inputs
op.inputs[1] = op.inputs[0]
op.inputs[0] = tens
@@ -233,22 +231,7 @@
desired_shape = [batch_size, n_in_elems]
if inp.shape != desired_shape:
# mismatch, insert a reshape to fix this.
- reshape_name = op.name + "_reshape"
- new_shape_tens = Tensor([1], DataType.int32, reshape_name + "_shape")
- new_shape_tens.values = np.array(desired_shape)
- new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
- new_shape_tens.ops = [new_shape_tens_const]
- new_shape_tens_const.outputs = [new_shape_tens]
-
- reshape_op = Operation("Reshape", reshape_name)
- reshape_op.inputs = [inp, new_shape_tens]
- reshape_op.attrs["new_shape"] = desired_shape
- reshape_out = inp.clone("_reshaped")
- reshape_out.set_all_shapes(desired_shape)
- reshape_out.ops = [reshape_op]
- reshape_op.outputs = [reshape_out]
-
- op.inputs[0] = reshape_out
+ op.inputs[0] = create_reshape_tensor(inp, desired_shape)
return op
@@ -261,22 +244,16 @@
desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
# Construct 1 shape tensor to be used by all inserted reshape ops
- new_shape_name = op.name + "_reshape_shape"
- new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
- new_shape_tens.values = np.array(desired_shape)
- new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
- new_shape_tens.ops = [new_shape_tens_const]
- new_shape_tens_const.outputs = [new_shape_tens]
+ new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, desired_shape)
for idx, inp in enumerate(op.inputs):
- reshape_name = op.name + str(idx) + "_reshape"
- reshape_op = Operation("Reshape", reshape_name)
- reshape_op.inputs = [inp, new_shape_tens]
- reshape_op.attrs["new_shape"] = desired_shape
reshape_out = inp.clone("_reshaped")
reshape_out.set_all_shapes(desired_shape)
- reshape_out.ops = [reshape_op]
- reshape_op.outputs = [reshape_out]
+
+ reshape_op = Operation("Reshape", "{}{}_reshape".format(op.name, idx))
+ reshape_op.attrs["new_shape"] = desired_shape
+ reshape_op.inputs = [inp, new_shape_tens]
+ reshape_op.set_output_tensor(reshape_out)
op.inputs[idx] = reshape_out
@@ -335,22 +312,17 @@
reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
# Construct 1 shape tensor to be used by all inserted reshape ops
- new_shape_name = op.name + "_reshape_shape"
- new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
- new_shape_tens.values = np.array(tens.shape)
- new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
- new_shape_tens.ops = [new_shape_tens_const]
- new_shape_tens_const.outputs = [new_shape_tens]
+ new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
for idx, out_tens in enumerate(op.outputs):
- reshape_name = op.name + str(idx) + "_reshape"
- reshape_op = Operation("Reshape", reshape_name)
- reshape_op.outputs = [out_tens]
reshape_in = out_tens.clone("_reshaped")
reshape_in.set_all_shapes(reshape_input_shape)
reshape_in.ops = [op]
- out_tens.ops = [reshape_op]
+
+ reshape_op = Operation("Reshape", "{}{}_reshape".format(op.name, idx))
+ reshape_op.attrs["new_shape"] = reshape_input_shape
reshape_op.inputs = [reshape_in, new_shape_tens]
+ reshape_op.set_output_tensor(out_tens)
op.outputs[idx] = reshape_in
@@ -517,17 +489,12 @@
fc_ofm_tensor.set_all_shapes([1, fc_ofm_tensor.shape[-1]])
fc_ofm_tensor.ops = [op]
# Add a reshape after the new OFM to convert it back to the original 4D shape
- reshape_name = op.name + "_reshape_post"
- new_shape_tens = Tensor([1], DataType.int32, reshape_name + "_shape")
- new_shape_tens.values = np.array(orig_ofm_tensor.shape)
- new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
- new_shape_tens.ops = [new_shape_tens_const]
- new_shape_tens_const.outputs = [new_shape_tens]
+ reshape_name = op.name + "_reshape"
+ new_shape_tens = create_const_tensor(reshape_name + "_shape", [1], DataType.int32, orig_ofm_tensor.shape)
reshape_op = Operation("Reshape", reshape_name)
- reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
reshape_op.attrs["new_shape"] = orig_ofm_tensor.shape
- orig_ofm_tensor.ops = [reshape_op]
- reshape_op.outputs = [orig_ofm_tensor]
+ reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
+ reshape_op.set_output_tensor(orig_ofm_tensor)
# Replace this ops OFM to point to the 2D tensor
op.outputs[0] = fc_ofm_tensor
return op
@@ -542,8 +509,7 @@
act_op.inputs = [prep_op.inputs[0]]
act_op_out = act_op.inputs[0].clone("_acted")
act_op_out.quantization = op.outputs[0].quantization.clone()
- act_op_out.ops = [act_op]
- act_op.outputs = [act_op_out]
+ act_op.set_output_tensor(act_op_out)
prep_op.inputs[0] = act_op_out
prep_op.outputs[0].quantization = act_op_out.quantization.clone()