[Topi] Cortex-M DSP support (apache#9233)

Co-authored-by: Sergey Smirnov <Sergey.Smirnov@mir.dev>
Co-authored-by: Ekaterina Bern <Ekaterina.Bern@mir.dev>
Co-authored-by: Mikhail Trubnikov <Mikhail.Trubnikov@mir.dev>
Co-authored-by: German Tretiakov <german.tretiakov@mir.dev>
Co-authored-by: Ilya Gozman <Ilya.Gozman@mir.dev>
Co-authored-by: Alexey.Yazev <Alexey.Yazev@mir.dev>
Co-authored-by: Ilya Gozman <92577591+ilyag-grovety@users.noreply.github.com>

python/tvm/relay/op/strategy/arm_cpu.py

@@ -19,7 +19,7 @@
 import re
 import logging
 
-from tvm import topi
+from tvm import relay, topi
 from ....target import arm_isa
 from ....topi.generic import conv2d as conv2d_generic
 from .generic import *
@@ -49,6 +49,25 @@ def schedule_concatenate_arm_cpu(_, outs, target):
         return topi.arm_cpu.schedule_concatenate(outs)
 
 
+@schedule_pool.register(["arm_cpu"])
+def schedule_pool_arm_cpu(attrs, outs, target):
+    """schedule pooling ops arm cpu"""
+    layout = attrs.layout
+    isa = arm_isa.IsaAnalyzer(target)
+    avg_pool = isinstance(attrs, relay.op.op_attrs.AvgPool2DAttrs)
+    with target:
+        if (
+            avg_pool
+            and isa.has_dsp_support
+            and layout in ("NCW", "NCHW")
+            or not avg_pool
+            and isa.has_dsp_support
+            and layout in ("NWC", "NHWC")
+        ):
+            return topi.arm_cpu.schedule_pool(outs, layout)
+        return topi.generic.schedule_pool(outs, layout)
+
+
 @conv2d_strategy.register(["arm_cpu", "micro_dev"])
 def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target):
     """conv2d arm cpu strategy"""
@@ -128,11 +147,11 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target):
                 name="conv2d_hwcn.generic",
             )
         elif layout == "NHWC":
-            if "SMLAD" in isa and kernel_layout == "HWOI":
+            if isa.has_dsp_support and kernel_layout == "HWOI":
                 strategy.add_implementation(
-                    wrap_compute_conv2d(topi.arm_cpu.conv2d_nhwc_direct_simd),
-                    wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_nhwc_direct_simd),
-                    name="conv2d_nhwc_direct_simd.micro_dev",
+                    wrap_compute_conv2d(topi.arm_cpu.conv2d_nhwc_dsp),
+                    wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_nhwc_dsp),
+                    name="conv2d_nhwc_dsp.micro_dev",
                 )
             elif kernel_layout == "HWIO":
                 is_aarch64 = topi.arm_cpu.arm_utils.is_aarch64_arm()
@@ -415,3 +434,67 @@ def schedule_bitserial_dense_arm_cpu(attrs, inputs, out_type, target):
         name="bitserial_dense.arm_cpu",
     )
     return strategy
+
+
+@dense_strategy.register(["arm_cpu"])
+def schedule_dense_arm_cpu(attrs, inputs, out_type, target):
+    """dense arm cpu strategy"""
+    strategy = _op.OpStrategy()
+    isa = arm_isa.IsaAnalyzer(target)
+    if isa.has_dsp_support:
+        strategy.add_implementation(
+            wrap_compute_dense(topi.nn.dense),
+            wrap_topi_schedule(topi.arm_cpu.schedule_dense_dsp),
+            name="dense_dsp",
+        )
+    else:
+        strategy.add_implementation(
+            wrap_compute_dense(topi.nn.dense),
+            wrap_topi_schedule(topi.generic.schedule_dense),
+            name="dense.generic",
+        )
+    return strategy
+
+
+@conv1d_strategy.register("arm_cpu")
+def conv1d_strategy_arm_cpu(attrs, inputs, out_type, target):
+    """conv1d strategy"""
+    strategy = _op.OpStrategy()
+    layout = attrs.data_layout
+    kernel_layout = attrs.kernel_layout
+    dilation = get_const_tuple(attrs.dilation)
+    if dilation[0] < 1:
+        raise ValueError("dilation should be a positive value")
+
+    isa = arm_isa.IsaAnalyzer(target)
+
+    if kernel_layout == "WOI":
+        if layout == "NWC" and isa.has_dsp_support:
+            strategy.add_implementation(
+                wrap_compute_conv1d(topi.arm_cpu.conv1d_nwc_dsp),
+                wrap_topi_schedule(topi.arm_cpu.schedule_conv1d_nwc_dsp),
+                name="conv1d_dsp",
+            )
+        else:
+            raise RuntimeError(
+                "Unsupported kernel layout {} for conv1d {} for arm cpu.".format(
+                    kernel_layout, layout
+                )
+            )
+    elif layout == "NCW":
+        strategy.add_implementation(
+            wrap_compute_conv1d(topi.nn.conv1d_ncw),
+            wrap_topi_schedule(topi.generic.schedule_conv1d_ncw),
+            name="conv1d_ncw.generic",
+        )
+    elif layout == "NWC":
+        strategy.add_implementation(
+            wrap_compute_conv1d(topi.nn.conv1d_nwc),
+            wrap_topi_schedule(topi.generic.schedule_conv1d_nwc),
+            name="conv1d_nwc.generic",
+        )
+    else:
+        raise RuntimeError(
+            "Unsupported kernel layout {} for conv1d {} for arm cpu.".format(kernel_layout, layout)
+        )
+    return strategy

python/tvm/target/arm_isa.py

@@ -16,18 +16,24 @@
 # under the License.
 """Defines functions to analyze available opcodes in the ARM ISA."""
 
+import tvm.target
 
-ARM_ISA_MAP = {
-    "armv7e-m": ["SMLAD"],
-}
+
+ARM_MPROFILE_DSP_SUPPORT_LIST = [
+    "cortex-m7",
+    "cortex-m4",
+    "cortex-m33",
+    "cortex-m35p",
+    "cortex-m55",
+]
 
 
 class IsaAnalyzer(object):
+    """Checks ISA support for given target"""
+
     def __init__(self, target):
-        self.target = target
-        # TODO: actually parse -mcpu
-        arch = "armv7e-m"
-        self._isa_map = ARM_ISA_MAP[arch]
+        self.target = tvm.target.Target(target)
 
-    def __contains__(self, instruction):
-        return instruction in self._isa_map
+    @property
+    def has_dsp_support(self):
+        return self.target.mcpu is not None and self.target.mcpu in ARM_MPROFILE_DSP_SUPPORT_LIST

python/tvm/testing/plugin.py

@@ -49,6 +49,7 @@
     "llvm": "mark a test as requiring llvm",
     "ethosn": "mark a test as requiring ethosn",
     "hexagon": "mark a test as requiring hexagon",
+    "corstone300": "mark a test as requiring Corstone300 FVP",
 }
 
 

python/tvm/testing/utils.py

@@ -674,6 +674,18 @@ def requires_opencl(*args):
     return _compose(args, _requires_opencl)
 
 
+def requires_corstone300(*args):
+    """Mark a test as requiring the corstone300 FVP
+
+    Parameters
+    ----------
+    f : function
+        Function to mark
+    """
+    _requires_corstone300 = [pytest.mark.corstone300]
+    return _compose(args, _requires_corstone300)
+
+
 def requires_rocm(*args):
     """Mark a test as requiring the rocm runtime.
 

python/tvm/topi/arm_cpu/__init__.py

@@ -17,6 +17,7 @@
 # pylint: disable=wildcard-import
 """Schedule for ARM CPU"""
 
+from .conv1d import *
 from .conv2d import *
 from .depthwise_conv2d import *
 from .conv2d_transpose import *
@@ -25,5 +26,6 @@
 from .bitserial_conv2d import *
 from .bitserial_dense import *
 from .injective import *
-from . import cortex_m7
 from .group_conv2d import *
+from .pooling import *
+from .dense import *

python/tvm/topi/arm_cpu/conv1d.py

@@ -0,0 +1,37 @@
+# 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.
+# pylint: disable=invalid-name, unused-variable, no-else-return, unused-argument, import-outside-toplevel
+"""Conv1D schedule for ARM CPU"""
+from __future__ import absolute_import as _abs
+
+from tvm import autotvm
+
+from .mprofile.dsp.conv1d import (
+    conv1d_nwc_dsp_compute,
+    conv1d_nwc_dsp_schedule,
+)
+
+
+@autotvm.register_topi_compute("conv1d_nwc_dsp.arm_cpu")
+def conv1d_nwc_dsp(cfg, data, kernel, strides, padding, dilation, out_dtype):
+    """Compute conv1d with v7e-m DSP instructions."""
+    return conv1d_nwc_dsp_compute(cfg, data, kernel, strides, padding, dilation, out_dtype)
+
+
+@autotvm.register_topi_schedule("conv1d_nwc_dsp.arm_cpu")
+def schedule_conv1d_nwc_dsp(cfg, outs):
+    return conv1d_nwc_dsp_schedule(cfg, outs)

python/tvm/topi/arm_cpu/conv2d.py

@@ -33,7 +33,10 @@
     schedule_conv2d_spatial_pack_nchw,
     schedule_conv2d_spatial_pack_nhwc,
 )
-from .cortex_m7.conv2d import direct_simd
+from .mprofile.dsp.conv2d import (
+    conv2d_nhwc_dsp_compute,
+    conv2d_nhwc_dsp_schedule,
+)
 
 
 @autotvm.register_topi_compute("conv2d_nchw_spatial_pack.arm_cpu")
@@ -505,15 +508,13 @@ def _callback(op):
     return s
 
 
-@autotvm.register_topi_compute("conv2d_nhwc_direct_simd.arm_cpu")
-def conv2d_nhwc_direct_simd(cfg, data, kernel, strides, padding, dilation, out_dtype):
-    """Compute conv2d_nhwc with SIMD (v7e-m)."""
-    return direct_simd.conv2d_nhwc_direct_simd_compute(
-        cfg, data, kernel, strides, padding, dilation, out_dtype
-    )
+@autotvm.register_topi_compute("conv2d_nhwc_dsp.arm_cpu")
+def conv2d_nhwc_dsp(cfg, data, kernel, strides, padding, dilation, out_dtype):
+    """Compute conv2d_nhwc with v7e-m DSP instructions."""
+    return conv2d_nhwc_dsp_compute(cfg, data, kernel, strides, padding, dilation, out_dtype)
 
 
-@autotvm.register_topi_schedule("conv2d_nhwc_direct_simd.arm_cpu")
-def schedule_conv2d_nhwc_direct_simd(cfg, outs):
-    """Create schedule for conv2d_nhwc_direct_simd"""
-    return direct_simd.conv2d_nhwc_direct_simd_schedule(cfg, outs)
+@autotvm.register_topi_schedule("conv2d_nhwc_dsp.arm_cpu")
+def schedule_conv2d_nhwc_dsp(cfg, outs):
+    """Create schedule for conv2d_nhwc_dsp"""
+    return conv2d_nhwc_dsp_schedule(cfg, outs)