[microNPU] Add support for scalar values (#9794)

* [microNPU] Add support for scalar values

PR #9515 enabled support for scalar constants, but didn't consider the
case of a scalar value where the underlying constant data does not have
a shape i.e. `constant.shape == []`. See the test case for a visual
differece when the scalar value is 1.

Change-Id: Id7a238cb5bf999dd5a8428c097202f9fb940a5f0

* Fix failing test by removing constant

Before this PR scalar constants were handled differently so this test
was able to pass. Now that scalar constants are handled in the same
manner as tensor constants, the test fails since unexpected tir is
produced in the compilation pipeline. Since the relay used in this test
case is not expected to be produced by higher levels of the compiler,
removing this constant for now.

Change-Id: I4ea5155778809041339e6faac05af3f72c3e3ea5

* clean up finding tensor from inputs

Change-Id: Ideccf84f8c9149148ff23e2406229cf637c982a3

Author: lhutton1

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

@@ -652,8 +652,8 @@ def callback(
             ifm2_zero_point=int(params.ifm2.q_params.zero_point),
             ofm_scale=float(params.ofm.q_params.scale_f32),
             ofm_zero_point=int(params.ofm.q_params.zero_point),
-            ifm_channels=params.ifm.shape[-1],
-            ifm2_channels=params.ifm2.shape[-1],
+            ifm_channels=params.ifm.shape[-1] if params.ifm.shape else 1,
+            ifm2_channels=params.ifm2.shape[-1] if params.ifm2.shape else 1,
             reversed_operands=params.reversed_operands,
             ofm_dtype=params.ofm.dtype,
             activation=activation,

python/tvm/relay/backend/contrib/ethosu/tir/compiler.py

@@ -123,12 +123,11 @@ def __init__(self):
 
     def visit_constant(self, const):
         if isinstance(const.checked_type, relay.ty.TensorType):
-            if const.checked_type.concrete_shape != ():
-                self.constants.append(const.data.asnumpy())
-                name = "p" + str(len(self.constants))
-                var = relay.var(type_annotation=const.checked_type, name_hint=name)
-                self.const_vars.append(var)
-                return var
+            self.constants.append(const.data.asnumpy())
+            name = "p" + str(len(self.constants))
+            var = relay.var(type_annotation=const.checked_type, name_hint=name)
+            self.const_vars.append(var)
+            return var
 
         return const
 

python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py

@@ -136,7 +136,10 @@ def _visit(tensor, reader, lut):
             if tensor not in planned:
                 planned.add(tensor)
                 if isinstance(tensor.op, tvm.te.PlaceholderOp) and tensor != lut:
-                    index = list(cached_func.inputs).index(tensor)
+                    # Find index of input using 'same_as' check to prevent equality
+                    # ambiguity when encountering a scalar.
+                    is_same = [var.same_as(tensor) for var in cached_func.inputs]
+                    index = is_same.index(True)
                     if index in const_dict:
                         sch.cache_read(tensor, "global", [reader])
 

tests/python/contrib/test_ethosu/test_codegen.py

@@ -629,12 +629,13 @@ def create_mod_from_relay():
 
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize("dtype", ["int8", "uint8"])
-def test_elementwise_add_from_constant_scalar(accel_type, dtype):
+@pytest.mark.parametrize("constant", [np.ones((1, 1, 1, 1)), np.array(1)])
+def test_elementwise_add_from_constant_scalar(accel_type, dtype, constant):
     ifm_shape = (1, 4, 4, 8)
 
     def create_relay_graph():
         inp = relay.var("input", shape=ifm_shape, dtype=dtype)
-        scalar = relay.const(np.ones((1, 1, 1, 1), dtype=dtype), dtype=dtype)
+        scalar = relay.const(constant, dtype=dtype)
         add = relay.qnn.op.add(
             inp,
             scalar,

tests/python/contrib/test_ethosu/test_compiler.py

@@ -34,8 +34,7 @@ def test_lower_to_tir():
         kernel_layout="HWIO",
         out_dtype="int32",
     )
-    multiply = relay.multiply(relay.const(-22, dtype="int32"), p2)
-    tile = relay.tile(multiply, reps=(1, 1, 1, 1001))
+    tile = relay.tile(p2, reps=(1, 1, 1, 1001))
     subtract = relay.subtract(conv, tile)
     func = subtract
     expr = relay.Function(relay.analysis.free_vars(func), func)