MLBEDSW-8491: Add support for Mirror pad
Change-Id: I3c13118e14195a5fb8e522a38b205b75fb07b74b
Signed-off-by: Rickard Bolin <rickard.bolin@arm.com>
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index cc947bc..ad979bd 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -1881,6 +1881,130 @@
return op
+def convert_mirror_pad(op: Operation, arch, nng):
+ if op.type != Op.MirrorPad or not op.run_on_npu:
+ return op
+
+ _, (top, bot), (left, right), _ = op.ifm2.values
+ mode = op.attrs["mode"] # 0 = reflect, 1 = symmetric
+
+ ifm = op.ifm
+ ofm = op.ofm
+ ofm.ops = []
+ elem_size = 2 if ofm.dtype == DataType.int16 else 1
+ n, h, w, c = ifm.shape
+ _, oh, ow, _ = ofm.shape
+ # Force linear format on OFM to allow negative stride multipliers
+ ofm.force_linear_format = True
+
+ # Intermediate ofm needed to store ifm padded with top and bot values as input to the left and right padding
+ intermediate_ofm_tens = Tensor([n, h + top + bot, w, c], ofm.dtype, "intermediate_ofm_tens")
+ intermediate_ofm_tens.quantization = op.outputs[0].quantization.clone()
+ intermediate_ofm_tens.force_linear_format = True
+
+ # If there is no left or right padding, we can write directly to the ofm without an intermediate tensor
+ if not (left or right):
+ intermediate_ofm_tens = ofm
+
+ # Initial op to copy the ifm into the middle of the intermediate ofm
+ avg_pool_init = create_avgpool_nop("init_pool")
+ avg_pool_init.write_shape = Shape4D(n, h, w, c)
+ avg_pool_init.write_offset = Shape4D(0, top, 0, 0)
+ avg_pool_init.read_shapes[0] = Shape4D(n, h, w, c)
+ avg_pool_init.read_offsets[0] = Shape4D(0, 0, 0, 0)
+ avg_pool_init.add_input_tensor(ifm)
+ avg_pool_init.set_output_tensor(intermediate_ofm_tens)
+ avg_pool_init.set_ifm_ofm_shapes()
+ DebugDatabase.add_optimised(op, avg_pool_init)
+
+ # Create pools with negative stride to mirror edges and offset to write at padding positions
+ avg_pool_pad = create_avgpool_nop("pad_pool")
+ for i, pad_amount in enumerate([top, bot, left, right]):
+ # Clear input from previous cloned op
+ avg_pool_pad.inputs = []
+ if not pad_amount:
+ continue
+
+ if i == 0: # top
+ # Set read and write shape width to full ifm width and height to "top" pad size
+ avg_pool_pad.write_shape = Shape4D(n, top, w, c)
+ avg_pool_pad.read_shapes[0] = Shape4D(n, top, w, c)
+ # Leave read offset as default to read the top chunk of the ifm
+ # For reflect mode, shift height offset down one step to "skip" the edge
+ avg_pool_pad.read_offsets[0] = Shape4D(0, 0, 0, 0) if mode == 1 else Shape4D(0, 1, 0, 0)
+ # Offset the base address of tile 0 to start writing just above the ifm that was copied into the middle of
+ # the ofm and use negative height striding to mirror the above ifm chunk
+ avg_pool_pad.tile_base_offsets_ofm[0] = ((top - 1) * w) * c * elem_size
+ if i == 1: # bot
+ # Set read and write shape width to full ifm width and height to "bot" pad size
+ avg_pool_pad.write_shape = Shape4D(n, bot, w, c)
+ avg_pool_pad.read_shapes[0] = Shape4D(n, bot, w, c)
+ # Set read offset to read the bottom chunk of the ifm
+ # For reflect mode, shift height offset up one step to "skip" the edge
+ avg_pool_pad.read_offsets[0] = Shape4D(0, h - bot, 0, 0) if mode == 1 else Shape4D(0, h - bot - 1, 0, 0)
+ # Offset the base address of tile 0 to start writing at the very bottom of the ofm and use negative height
+ # striding to mirror the above ifm chunk
+ avg_pool_pad.tile_base_offsets_ofm[0] = (oh - 1) * w * c * elem_size
+ if i == 2: # left
+ # Set read and write shape height to full intermediate ofm height and width to "left" pad size
+ avg_pool_pad.write_shape = Shape4D(n, h + top + bot, left, c)
+ avg_pool_pad.read_shapes[0] = Shape4D(n, h + top + bot, left, c)
+ # Leave read offset as default to read the leftmost chunk of the intermediate ofm
+ # For reflect mode, shift width offset one step to the right to "skip" the edge
+ avg_pool_pad.read_offsets[0] = Shape4D(0, 0, 0, 0) if mode == 1 else Shape4D(0, 0, 1, 0)
+ # Offset the base address of tile 0 to start writing just left of the intermediate ofm and use negative
+ # width striding to mirror the above ifm chunk
+ avg_pool_pad.tile_base_offsets_ofm[0] = (left - 1) * c * elem_size
+ if i == 3: # right
+ # Set read and write shape height to full intermediate ofm height and width to "right" pad size
+ avg_pool_pad.write_shape = Shape4D(n, h + top + bot, right, c)
+ avg_pool_pad.read_shapes[0] = Shape4D(n, h + top + bot, right, c)
+ # Set read offset to read the rightmost chunk of the intermediate ofm
+ # For reflect mode, shift width offset one step to the left to "skip" the edge
+ avg_pool_pad.read_offsets[0] = Shape4D(0, 0, w - right, 0) if mode == 1 else Shape4D(0, 0, w - right - 1, 0)
+ # Offset the base address of tile 0 to start writing at the rightmost part of the ofm and use negative
+ # width striding to mirror the above ifm chunk
+ avg_pool_pad.tile_base_offsets_ofm[0] = (ow - 1) * c * elem_size
+
+ # Write offset (0,0,0,0) for all convs
+ avg_pool_pad.write_offset = Shape4D(0, 0, 0, 0)
+
+ if i in [0, 1]: # negative height stride for top and bot, negative width stride for left and right
+ avg_pool_pad.ofm_stride_multiplier = [1, -1, 1] # C/H/W
+ # top and bot reads from ifm and writes to intermediate ofm
+ avg_pool_pad.add_input_tensor(ifm)
+ intermediate_ofm_tens.ops.append(avg_pool_pad)
+ avg_pool_pad.outputs = [intermediate_ofm_tens]
+ else:
+ avg_pool_pad.ofm_stride_multiplier = [1, 1, -1] # C/H/W
+ # left and right reads from intermediate ofm and writes to ofm
+ avg_pool_pad.add_input_tensor(intermediate_ofm_tens)
+ ofm.ops.append(avg_pool_pad)
+ avg_pool_pad.outputs = [ofm]
+
+ avg_pool_pad.set_ifm_ofm_shapes()
+ DebugDatabase.add_optimised(op, avg_pool_pad)
+
+ # Clone operation for next padding direction
+ avg_pool_pad = avg_pool_pad.clone(f"_{i}")
+
+ if left or right:
+ # Copy intermediate ofm into final ofm
+ avg_pool_final_copy = create_avgpool_nop("avg_pool_final_copy")
+ avg_pool_final_copy.write_shape = Shape4D(n, h + top + bot, w, c)
+ avg_pool_final_copy.write_offset = Shape4D(0, 0, left, 0)
+ avg_pool_final_copy.read_shapes[0] = Shape4D(n, h + top + bot, w, c)
+ avg_pool_final_copy.read_offsets[0] = Shape4D(0, 0, 0, 0)
+
+ avg_pool_final_copy.add_input_tensor(intermediate_ofm_tens)
+ ofm.ops.append(avg_pool_final_copy)
+ avg_pool_final_copy.outputs = [ofm]
+ avg_pool_final_copy.set_ifm_ofm_shapes()
+ DebugDatabase.add_optimised(op, avg_pool_final_copy)
+
+ return op
+
+
def convert_pad(op: Operation, arch, nng):
"""
Rewrites PAD operator to an average pool that copies the IFM to the OFM
@@ -2899,6 +3023,7 @@
convert_mul_max_to_abs_or_lrelu,
convert_lrelu,
convert_avg_pool_to_conv2d,
+ convert_mirror_pad,
fixup_strided_conv,
convert_hardswish_to_lut,
rewrite_fully_connected_input,