TOSA: Add ifm ofm elem size into raw output
Change-Id: I645496536a6bddf2bd289a87be9d7cef11693954
Signed-off-by: Diqing Zhong <diqing.zhong@arm.com>
diff --git a/ethosu/vela/rawdata_writer.py b/ethosu/vela/rawdata_writer.py
index 76765e6..247cfe8 100644
--- a/ethosu/vela/rawdata_writer.py
+++ b/ethosu/vela/rawdata_writer.py
@@ -37,9 +37,11 @@
if custom_op:
ifm_shapes = []
+ ifm_elem_sizes = []
ifm_regions = []
ifm_offsets = []
ofm_shapes = []
+ ofm_elem_sizes = []
ofm_regions = []
ofm_offsets = []
cmd_stream_tensor, weight_tensor, scratch_tensor, scratch_fast_tensor = custom_op.inputs[:4]
@@ -50,10 +52,12 @@
ifm_shapes.append(ifm.shape)
ifm_regions.append(get_region(ifm.mem_type, arch))
ifm_offsets.append(ifm.address)
+ ifm_elem_sizes.append(ifm.element_size())
for ofm in custom_op.outputs:
ofm_shapes.append(ofm.shape)
ofm_regions.append(get_region(ofm.mem_type, arch))
ofm_offsets.append(ofm.address)
+ ofm_elem_sizes.append(ofm.element_size())
filename_sg = f"{filename}_sg{sg_idx}_vela.npz"
np.savez(
@@ -66,9 +70,11 @@
scratch_fast_shape=scratch_fast_tensor.shape,
scratch_fast_region=scratch_fast_region,
input_shape=ifm_shapes,
+ input_elem_size=ifm_elem_sizes,
input_region=ifm_regions,
input_offset=ifm_offsets,
output_shape=ofm_shapes,
+ output_elem_size=ofm_elem_sizes,
output_region=ofm_regions,
output_offset=ofm_offsets,
)
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index 8c5e277..ab63afd 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -434,7 +434,7 @@
def element_size(self) -> int:
if self.element_size_bytes == 0:
- return self.dtype.size_in_bits() / 8
+ return self.dtype.size_in_bits() // 8
return self.element_size_bytes
# Returns a copy, renamed to self.name + suffix