From 438e5638e5a99ee75e0434818f48493dc1f0a373 Mon Sep 17 00:00:00 2001 From: Patrik Gustavsson Date: Tue, 1 Sep 2020 12:23:25 +0200 Subject: MLBEDSW-2903 Split mapping to tensor Split mapping to tensor Signed-off-by: Patrik Gustavsson Change-Id: Ic143f3b4d37f6904edd8f119eff1d108f70b5026 --- ethosu/vela/high_level_command_stream_generator.py | 29 +++++++++++++++------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/ethosu/vela/high_level_command_stream_generator.py b/ethosu/vela/high_level_command_stream_generator.py index d5a6341b..50b913d8 100644 --- a/ethosu/vela/high_level_command_stream_generator.py +++ b/ethosu/vela/high_level_command_stream_generator.py @@ -50,18 +50,29 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id npu_block_type = ps.npu_block_type split_offsets = [None, None] # offset for [ifm, ifm2] - ifm_idx = 0 - for op in ps.ops: - if op.type == "SplitSliceRead": - split_offsets[ifm_idx] = op.attrs["split_start"] - ps.primary_op.attrs["fused_memory_function"] = op.type - ifm_idx += 1 - if len(ps.inputs) == 2 and npu_block_type == NpuBlockType.ElementWise: - # Ensure correct imf and ifm2 order + # Ensure correct ifm and ifm2 order if match_tensor(ps.inputs[0], ps.primary_op.inputs[1]) and match_tensor(ps.inputs[1], ps.primary_op.inputs[0]): ps.ifm_tensor, ps.ifm2_tensor = ps.ifm2_tensor, ps.ifm_tensor - split_offsets[0], split_offsets[1] = split_offsets[1], split_offsets[0] + + for op in ps.ops: + if op.type == "SplitSliceRead": + ps.primary_op.attrs["fused_memory_function"] = op.type + assert len(op.inputs) == 1 + if match_tensor(ps.ifm_tensor, op.inputs[0]): + split_offsets[0] = op.attrs["split_start"] + elif match_tensor(ps.ifm2_tensor, op.inputs[0]): + split_offsets[1] = op.attrs["split_start"] + else: + assert False + else: + ifm_idx = 0 + for op in ps.ops: + if op.type == "SplitSliceRead": + assert ifm_idx < 2 + split_offsets[ifm_idx] = op.attrs["split_start"] + ps.primary_op.attrs["fused_memory_function"] = op.type + ifm_idx += 1 ifm_tensor = ps.ifm_tensor ifm2_tensor = ps.ifm2_tensor -- cgit v1.2.1