Added few tests for scalar_to_tensor_constant pass

Change-Id: Ib916e4b29be124ccdf4c88936f10499b6642f141

Author: Ashutosh Parkhi

python/tvm/relay/op/contrib/cmsisnn.py

@@ -56,10 +56,8 @@ def partition_for_cmsisnn(mod, params=None, **opts):
             transform.MergeComposite(pattern_table()),
             transform.AnnotateTarget("cmsis-nn"),
             transform.PartitionGraph(),
-            transform.InferType(),
             GenerateCMSISNNConstants(),
             ScalarToTensorConstants(),
-            transform.InferType(),
             ExtractConstantsFromPartitionedFunction(),
             transform.InferType(),
         ]
@@ -228,8 +226,6 @@ def binary_op_pattern(op):
 
     def check_qnn_binary_op(pattern):
         """Check if multiply is supported by CMSIS-NN."""
-        import numpy as np
-
         arg0 = pattern.args[0]
         arg1 = pattern.args[1]
         both_args_scalar = False

src/relay/backend/contrib/cmsisnn/extract_constants.cc

@@ -229,7 +229,7 @@ transform::Pass ExtractConstantsFromPartitionedFunction() {
   runtime::TypedPackedFunc<IRModule(IRModule, transform::PassContext)> pass_func =
       [=](IRModule m, transform::PassContext pc) { return ExtractConstants(m); };
   return tvm::transform::CreateModulePass(pass_func, 0, "ExtractConstantsFromPartitionedFunction",
-                                          {});
+                                          {"InferType"});
 }
 
 TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.ExtractConstantsFromPartitionedFunction")

src/relay/backend/contrib/cmsisnn/scalar_to_tensor_constant.cc

@@ -71,9 +71,6 @@ class ScalarToTensorConstantMutator : public MixedModeMutator {
     // in the binary op. This needs to be done only when one of the arguments is a scalar.
     if (auto* opnode = call->op.as<OpNode>()) {
       String op_name = opnode->name;
-      if (op_name != "qnn.mul" && op_name != "qnn.add") {
-        return final_call;
-      }
       Array<Expr> new_args;
       for (uint32_t i = 0; i < call->args.size(); ++i) {
         Expr arg = call->args[i];
@@ -179,7 +176,7 @@ IRModule ScalarToTensorConstant(const IRModule& mod) {
 transform::Pass ScalarToTensorConstantPass() {
   runtime::TypedPackedFunc<IRModule(IRModule, transform::PassContext)> pass_func =
       [=](IRModule m, transform::PassContext pc) { return ScalarToTensorConstant(m); };
-  return tvm::transform::CreateModulePass(pass_func, 0, "ScalarToTensorConstant", {});
+  return tvm::transform::CreateModulePass(pass_func, 0, "ScalarToTensorConstant", {"InferType"});
 }
 
 TVM_REGISTER_GLOBAL("relay.ext.cmsisnn.transform.ScalarToTensorConstants")

tests/python/contrib/test_cmsisnn/test_extract_constants.py

@@ -23,15 +23,6 @@
 import tvm
 from tvm import relay
 
-from utils import (
-    make_module,
-    count_num_calls,
-    get_range_for_dtype_str,
-    get_same_padding,
-    get_conv2d_qnn_params,
-    make_qnn_relu,
-)
-
 tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
 
 

tests/python/contrib/test_cmsisnn/test_scalar_to_tensor_constant.py

@@ -0,0 +1,187 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""CMSIS-NN integration tests: scalar_to_tensor_constant pass"""
+import itertools
+import math
+import numpy as np
+import pytest
+import tvm
+from tvm import relay
+
+tvm._ffi._init_api("relay.ext.cmsisnn.transform", __name__)
+
+
+class CheckFunctionsForConstants(tvm.relay.ExprVisitor):
+    def __init__(self):
+        super().__init__()
+        self.num_constants_ = 0
+
+    def visit_call(self, call):
+        super().visit_call(call)
+        for arg in call.args:
+            if isinstance(arg, relay.Constant) and arg.data.numpy().ndim > 0:
+                self.num_constants_ += 1
+
+    def check_num_constants(self, func):
+        assert self.num_constants_ == 0, "Functions should not have constant arguments in Calls"
+
+
+def set_external_func_attr(func, compiler, ext_symbol):
+    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+    func = func.with_attr("Compiler", compiler)
+    func = func.with_attr("global_symbol", ext_symbol)
+    return func
+
+
+def set_composite_func_attr(func, name):
+    func = func.with_attr("Composite", name)
+    return func
+
+
+@tvm.testing.requires_cmsisnn
+def test_single_scalar_position_0():
+    x0 = relay.var("x0", shape=None)
+    x1 = relay.var("x1", shape=(8, 8))
+    z1 = x0 + x1
+    lf = relay.Function([x0, x1], z1, relay.TensorType((8, 8), "float32"))
+    lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
+
+    y0 = relay.expr.const(3, "float32")
+    y1 = relay.var("y1", shape=(8, 8))
+    c0 = relay.Call(lf, [y0, y1])
+    ef = relay.Function([y1], c0, relay.TensorType((8, 8), "float32"))
+
+    x = relay.var("x", shape=(8, 8))
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint)
+    c = relay.Call(ev, [x])
+    mf = relay.Function([x], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = relay.transform.InferType()(mod)
+    mod = ScalarToTensorConstants()(mod)
+    check_for_constants = CheckFunctionsForConstants()
+    check_for_constants.visit_call(mod[ev].body)
+    assert (
+        check_for_constants.num_constants_ == 1
+    ), "Scalar constant wasn't converted into tensor constant"
+
+
+@tvm.testing.requires_cmsisnn
+def test_single_scalar_position_1():
+    x0 = relay.var("x0", shape=(8, 8))
+    x1 = relay.var("x1", shape=None)
+    z1 = x0 + x1
+    lf = relay.Function([x0, x1], z1, relay.TensorType((8, 8), "float32"))
+    lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
+
+    y0 = relay.var("y0", shape=(8, 8))
+    y1 = relay.expr.const(3, "float32")
+    c0 = relay.Call(lf, [y0, y1])
+    ef = relay.Function([y0], c0, relay.TensorType((8, 8), "float32"))
+
+    x = relay.var("x", shape=(8, 8))
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint)
+    c = relay.Call(ev, [x])
+    mf = relay.Function([x], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = relay.transform.InferType()(mod)
+    mod = ScalarToTensorConstants()(mod)
+    check_for_constants = CheckFunctionsForConstants()
+    check_for_constants.visit_call(mod[ev].body)
+    assert (
+        check_for_constants.num_constants_ == 1
+    ), "Scalar constant wasn't converted into tensor constant"
+
+
+@tvm.testing.requires_cmsisnn
+def test_two_scalars():
+    x1 = relay.var("x1", shape=None)
+    x2 = relay.var("x2", shape=None)
+    z1 = x1 + x2
+    lf = relay.Function([x1, x2], z1, relay.TensorType((), "float32"))
+    lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
+
+    y0 = relay.expr.const(5, "float32")
+    y1 = relay.expr.const(3, "float32")
+    c0 = relay.Call(lf, [y0, y1])
+    ef = relay.Function([], c0, relay.TensorType((), "float32"))
+
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint)
+    c = relay.Call(ev, [])
+    mf = relay.Function([], c, relay.TensorType((), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = relay.transform.InferType()(mod)
+    mod = ScalarToTensorConstants()(mod)
+    check_for_constants = CheckFunctionsForConstants()
+    check_for_constants.visit_call(mod[ev].body)
+    assert (
+        check_for_constants.num_constants_ == 0
+    ), "Scalar constant wasn't converted into tensor constant"
+
+
+@tvm.testing.requires_cmsisnn
+def test_two_tensor_constants():
+    x0 = relay.var("x0", shape=(8, 8))
+    x1 = relay.var("x1", shape=(8, 8))
+    z1 = x0 + x1
+    lf = relay.Function([x0, x1], z1, relay.TensorType((8, 8), "float32"))
+    lf = set_composite_func_attr(lf, "cmsis-nn.qnn_add")
+
+    y0 = relay.const(np.random.uniform(0, 1, (8, 8)).astype("float32"), "float32")
+    y1 = relay.const(np.random.uniform(0, 1, (8, 8)).astype("float32"), "float32")
+    c0 = relay.Call(lf, [y0, y1])
+    ef = relay.Function([], c0, relay.TensorType((8, 8), "float32"))
+
+    ev = relay.GlobalVar("external_function")
+    ef = set_external_func_attr(ef, "cmsis-nn", ev.name_hint)
+    c = relay.Call(ev, [])
+    mf = relay.Function([], c, relay.TensorType((8, 8), "float32"))
+    mv = relay.GlobalVar("main")
+
+    mod = tvm.IRModule()
+    mod[ev] = ef
+    mod[mv] = mf
+
+    mod = relay.transform.InferType()(mod)
+    mod = ScalarToTensorConstants()(mod)
+    check_for_constants = CheckFunctionsForConstants()
+    check_for_constants.visit_call(mod[ev].body)
+    assert (
+        check_for_constants.num_constants_ == 2
+    ), "Scalar constant wasn't converted into tensor constant"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))