[microNPU][ETHOSU] Softmax int8 legalization support

python/tvm/autotvm/task/topi_integration.py

@@ -250,11 +250,15 @@ def wrapper(outs, *args, **kwargs):
 
 def get_workload(outs, task_name=None):
     """Retrieve the workload from outputs"""
+    visited = set()
 
     def traverse(tensors):
         """traverse all ops to find attached workload"""
         for t in tensors:
             op = t.op
+            if op in visited:
+                continue
+            visited.add(op)
             wkl = traverse(op.input_tensors)
             if wkl is not None:
                 return wkl

python/tvm/relay/backend/contrib/ethosu/__init__.py

@@ -21,3 +21,4 @@
 from . import codegen
 from . import vela_api
 from . import tir_to_cs_translator
+from . import softmax_rewriter

python/tvm/relay/backend/contrib/ethosu/legalize.py

@@ -32,6 +32,7 @@
 from tvm.relay.backend.contrib.ethosu import op as ethosu_ops  # type: ignore
 from tvm.relay.backend.contrib.ethosu import vela_api
 from tvm.relay.backend.contrib.ethosu import util
+from tvm.relay.backend.contrib.ethosu.softmax_rewriter import SoftmaxRewriter
 from tvm.relay.op.contrib import ethosu as ethosu_patterns  # type: ignore
 
 
@@ -1479,6 +1480,7 @@ def transform_npu_function(self, _, func: relay.Function) -> relay.Function:
             LeakyReLURewriter(),
             MeanRewriter(),
             SumRewriter(),
+            SoftmaxRewriter(),
             ConcatRewriter(),
             SigmoidRewriter(),
             RequantizeRewriter(),