[Topi] Cortex-M DSP support

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,26 @@ 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 layout in ("NCW", "NCHW")
+            and "SMLAD" in isa
+            or not avg_pool
+            and "SSUB8" in isa
+            and "SEL" in isa
+            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"""
@@ -130,9 +150,9 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target):
         elif layout == "NHWC":
             if "SMLAD" in isa and kernel_layout == "HWOI":
                 strategy.add_implementation(
-                    wrap_compute_conv2d(topi.arm_cpu.conv2d_direct_simd),
-                    wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_direct_simd),
-                    name="conv2d_direct_simd.micro_dev",
+                    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",
                 )
             elif kernel_layout == "HWIO":
                 is_aarch64 = topi.arm_cpu.arm_utils.is_aarch64_arm()
@@ -415,3 +435,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 "SMLAD" in isa:
+        strategy.add_implementation(
+            wrap_compute_dense(topi.nn.dense),
+            wrap_topi_schedule(topi.arm_cpu.schedule_dense_direct_simd),
+            name="dense_direct_simd",
+        )
+    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 "SMLAD" in isa:
+            strategy.add_implementation(
+                wrap_compute_conv1d(topi.arm_cpu.conv1d_nwc_direct_simd),
+                wrap_topi_schedule(topi.arm_cpu.schedule_conv1d_nwc_direct_simd),
+                name="conv1d_direct_simd",
+            )
+        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/relay/qnn/op/legalizations.py

@@ -374,6 +374,8 @@ def _qnn_conv2d_legalize_arm_cpu(attrs, inputs, types):
         attrs["kernel_layout"],
         attrs["groups"],
     )
+
+    # Use int8 for Cortex-M7
     use_int8_on_arm = (not is_depthwise) and is_aarch64_arm() and attrs["data_layout"] == "NHWC"
     if use_int8_on_arm or is_fast_int8_on_arm():
         return helper_change_dtypes_to_be_same(attrs, inputs, types, relay.qnn.op.conv2d)

python/tvm/target/arm_isa.py

@@ -16,18 +16,24 @@
 # under the License.
 """Defines functions to analyze available opcodes in the ARM ISA."""
 
+import argparse
 
 ARM_ISA_MAP = {
-    "armv7e-m": ["SMLAD"],
+    "armv7e-m": ["SMLAD", "SSUB8", "SEL"],
+    "armv8-m": ["SMLAD", "SSUB8", "SEL"],
 }
 
 
 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]
+        parser = argparse.ArgumentParser()
+        parser.add_argument("-mcpu", type=str)
+        parser.add_argument("-march", type=str)
+        args, _ = parser.parse_known_args(str(target).split())
+        self._isa_map = ARM_ISA_MAP[args.march] if args.march in ARM_ISA_MAP else []
 
     def __contains__(self, instruction):
         return instruction in self._isa_map

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 *
@@ -27,3 +28,5 @@
 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,36 @@
+# 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 .cortex_m7.conv1d import direct_simd
+
+
+@autotvm.register_topi_compute("conv1d_nwc_direct_simd.arm_cpu")
+def conv1d_nwc_direct_simd(cfg, data, kernel, strides, padding, dilation, out_dtype):
+    """Compute conv1d with v7e-m DSP instructions."""
+    return direct_simd.conv1d_nwc_direct_simd_compute(
+        cfg, data, kernel, strides, padding, dilation, out_dtype
+    )
+
+
+@autotvm.register_topi_schedule("conv1d_nwc_direct_simd.arm_cpu")
+def schedule_conv1d_nwc_direct_simd(cfg, outs):
+    return direct_simd.conv1d_nwc_direct_simd_schedule(cfg, outs)

python/tvm/topi/arm_cpu/conv2d.py

@@ -505,15 +505,15 @@ def _callback(op):
     return s
 
 
-@autotvm.register_topi_compute("conv2d_direct_simd.arm_cpu")
-def conv2d_direct_simd(cfg, data, kernel, strides, padding, dilation, out_dtype):
-    """Compute conv2d with SIMD (v7e-m)."""
-    return direct_simd.conv2d_direct_simd_compute(
+@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 v7e-m DSP instructions."""
+    return direct_simd.conv2d_nhwc_direct_simd_compute(
         cfg, data, kernel, strides, padding, dilation, out_dtype
     )
 
 
-@autotvm.register_topi_schedule("conv2d_direct_simd.arm_cpu")
-def schedule_conv2d_direct_simd(cfg, outs):
-    """Create schedule for conv2d_direct_simd"""
-    return direct_simd.conv2d_direct_simd_nhwc_schedule(cfg, outs)
+@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)

python/tvm/topi/arm_cpu/cortex_m7/__init__.py

@@ -14,7 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Schedules specialized for cortex-m7."""
+"""Schedules specialized for v7e-m DSP instructions."""
 
 
 from . import conv2d
+from . import dense

python/tvm/topi/arm_cpu/cortex_m7/conv1d/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""Conv1d implementations for v7e-m DSP instructions."""
+
+from . import direct_simd