From 294e63124e91f5cee750e542a0bbd3e0f11e6c12 Mon Sep 17 00:00:00 2001
From: Tristan Konolige <tkonolige@octoml.ai>
Date: Wed, 23 Nov 2022 09:43:00 -0800
Subject: [PATCH] [ARM] Add dynamic matvec support

---
 python/tvm/relay/op/strategy/arm_cpu.py | 18 +++++++++++++++++-
 python/tvm/relay/op/strategy/x86.py     |  9 ++++-----
 python/tvm/topi/utils.py                |  5 +++++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py
index 5c25696a1ee1..261b979dedaf 100644
--- a/python/tvm/relay/op/strategy/arm_cpu.py
+++ b/python/tvm/relay/op/strategy/arm_cpu.py
@@ -21,7 +21,7 @@
 # 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 897f7c4e588f..3e59209f5822 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 8251dac4137b..704fc87fc27e 100644
--- a/python/tvm/topi/utils.py
+++ b/python/tvm/topi/utils.py
@@ -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, (tir.Any, tir.SizeVar)) for x in shape])