[tvm] branch main updated: [ARM] Add dynamic matvec support (#13502)

[an...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

andrewzhaoluo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 95d2e9fa35 [ARM] Add dynamic matvec support (#13502)
95d2e9fa35 is described below

commit 95d2e9fa35524bbdafbe4ff758523eb571055d02
Author: AndrewZhaoLuo <an...@gmail.com>
AuthorDate: Mon Nov 28 19:36:42 2022 -0800

    [ARM] Add dynamic matvec support (#13502)
    
    * [ARM] Add dynamic matvec support
    
    * proper imports
    
    Co-authored-by: Tristan Konolige <tk...@octoml.ai>
---
 python/tvm/relay/op/strategy/arm_cpu.py | 18 +++++++++++++++++-
 python/tvm/relay/op/strategy/x86.py     |  9 ++++-----
 python/tvm/topi/utils.py                |  7 ++++++-
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py
index 5c25696a1e..261b979ded 100644
--- a/python/tvm/relay/op/strategy/arm_cpu.py
+++ b/python/tvm/relay/op/strategy/arm_cpu.py
@@ -21,7 +21,7 @@ import logging
 # pylint: disable=invalid-name,unused-argument,wildcard-import,unused-wildcard-import
 import re
 
-from tvm import relay, topi
+from tvm import relay, topi, tir
 
 from ....auto_scheduler import is_auto_scheduler_enabled
 from ....meta_schedule import is_meta_schedule_enabled
@@ -558,6 +558,22 @@ def schedule_dense_arm_cpu(attrs, inputs, out_type, target):
             name="dense_dsp.arm_cpu",
         )
     else:
+        # For dynamic matrix-vector multiply we use a hand written kernel.
+        if (
+            isinstance(inputs[0].shape[0], (int, tir.IntImm))
+            and inputs[0].shape[0] == 1
+            and (
+                topi.utils.is_dynamic_shape(inputs[0].shape)
+                or topi.utils.is_dynamic_shape(inputs[1].shape)
+            )
+        ):
+            strategy.add_implementation(
+                wrap_compute_dense(topi.x86.dense_dynamic),
+                wrap_topi_schedule(topi.x86.schedule_dense_dynamic),
+                name="dense_dynamic.x86",
+                plevel=20,
+            )
+            return strategy
         logger.warning("dense is not optimized for arm cpu.")
         strategy.add_implementation(
             wrap_compute_dense(
diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py
index 897f7c4e58..3e59209f58 100644
--- a/python/tvm/relay/op/strategy/x86.py
+++ b/python/tvm/relay/op/strategy/x86.py
@@ -507,10 +507,6 @@ def matmul_strategy_cpu(attrs, inputs, out_type, target):
     return strategy
 
 
-def is_dynamic_shape(shape):
-    return any([isinstance(x, (tir.Any, tir.SizeVar)) for x in shape])
-
-
 @dense_strategy.register("cpu")
 def dense_strategy_cpu(attrs, inputs, out_type, target):
     """dense x86 strategy"""
@@ -520,7 +516,10 @@ def dense_strategy_cpu(attrs, inputs, out_type, target):
     if (
         isinstance(inputs[0].shape[0], (int, tir.IntImm))
         and inputs[0].shape[0] == 1
-        and (is_dynamic_shape(inputs[0].shape) or is_dynamic_shape(inputs[1].shape))
+        and (
+            topi.utils.is_dynamic_shape(inputs[0].shape)
+            or topi.utils.is_dynamic_shape(inputs[1].shape)
+        )
     ):
         strategy.add_implementation(
             wrap_compute_dense(topi.x86.dense_dynamic),
diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py
index 8251dac413..7580eac021 100644
--- a/python/tvm/topi/utils.py
+++ b/python/tvm/topi/utils.py
@@ -23,7 +23,7 @@ from numbers import Integral
 import numpy as np
 import tvm
 from tvm import te
-from tvm.tir import bijective_layout, layout
+from tvm.tir import Any, SizeVar, bijective_layout, layout
 
 from . import cpp, tag
 
@@ -534,3 +534,8 @@ def is_target(names):
     names = [names] if isinstance(names, str) else names
     target = tvm.target.Target.current(allow_none=False)
     return any(name in target.keys for name in names)
+
+
+def is_dynamic_shape(shape):
+    """Checks if any part of a shape is dynamic"""
+    return any([isinstance(x, (Any, SizeVar)) for x in shape])