commit a0c3624899edc601525a589643c802469003f89d
author Fredrik Svedberg <fredrik.svedberg@arm.com> Wed Jun 03 15:43:31 2020 +0200
committer Louis Verhaard <louis.verhaard@arm.com> Wed Aug 05 16:26:04 2020 +0200
tree 4fc52db04cd29901b3e5d4a7425a7a641e9647fb
parent 9a03fdff316662be69a1adc4e391e43bc6519b08

[MLBEDSW-2335] SoftMax int16

Added graph rewrite of Softmax for int16.

Change-Id: Id7885af6056a23e8b8362fb61ae94283251eb398
Signed-off-by: Fredrik Svedberg <fredrik.svedberg@arm.com>

ethosu/vela/graph_optimiser.py

diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 355b16f..9c6e1f5 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -27,6 +27,7 @@
 from .numeric_util import full_shape
 from .operation import NpuBlockType
 from .operation import Operation
+from .softmax import SoftMax
 from .tensor import QuantizationParameters
 from .tensor import Tensor
 
@@ -357,7 +358,7 @@
         if "Conv" in op.type:
             kernel_size = op.inputs[1].shape[:2]
             input_shape = op.inputs[0].shape
-        elif "Pool" in op.type or "ResizeBilinear" == op.type:
+        elif "Pool" in op.type or op.type in ("ResizeBilinear", "ReduceSum"):
             kernel_size = op.attrs["ksize"][1:3]
             input_shape = op.inputs[0].shape
         elif op.type == "ExtractImagePatches":
@@ -401,9 +402,10 @@
 )
 depthwise_op = set(("DepthwiseConv2dNative", "DepthwiseConv2dBiasAct",))
 pool_op = set(
-    ("AvgPool", "MaxPool", "QuantizedAvgPool", "QuantizedMaxPool", "AvgPoolAct", "MaxPoolAct", "ResizeBilinear",)
+    ("AvgPool", "MaxPool", "QuantizedAvgPool", "QuantizedMaxPool", "AvgPoolAct", "MaxPoolAct", "ResizeBilinear")
 )
-elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum", "LeakyRelu", "Abs"))
+reduce_sum_ops = set(("ReduceSum",))
+elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum", "LeakyRelu", "Abs", "CLZ", "SHL", "SHR"))
 binary_elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum"))
 activation_ops = set(("Relu", "Relu6", "ReluN1To1", "Sigmoid", "Tanh"))
 memory_only_ops = set(("Reshape",))
@@ -437,6 +439,8 @@
         npu_block_type = NpuBlockType.Pooling
     elif op.type in elementwise_op:
         npu_block_type = NpuBlockType.ElementWise
+    elif op.type in reduce_sum_ops:
+        npu_block_type = NpuBlockType.ReduceSum
 
     op.attrs["npu_block_type"] = npu_block_type
     return op
@@ -573,6 +577,13 @@
     return op
 
 
+def convert_softmax(op, arch):
+    if op.type == "Softmax" and op.run_on_npu:
+        softmax = SoftMax(op)
+        op = softmax.get_graph()
+    return op
+
+
 def convert_mul_max_to_abs_or_lrelu(op, arch):
     r"""Whenever there is a subgraph with this topology:
 
@@ -671,6 +682,7 @@
         # then do any rewrites of supported operators
         convert_depthwise_to_conv,
         convert_conv_to_fc,
+        convert_softmax,
         fixup_fully_connected_input,
         fixup_pack_input,
         fixup_conv2d_backprop,