From c464e06de99e2ec9de13357ebcc7bab849f12628 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <kparzysz@quicinc.com>
Date: Thu, 2 Sep 2021 14:06:42 -0500
Subject: [PATCH] [Hexagon] Add trivial conv2d operator to Hexagon relay
 strategy

This is just to enable relay codegen for conv2d on Hexagon, not for
performance.
---
 python/tvm/relay/op/strategy/__init__.py |  1 +
 python/tvm/relay/op/strategy/hexagon.py  | 44 ++++++++++++++++++++++++
 python/tvm/topi/__init__.py              |  1 +
 python/tvm/topi/hexagon/__init__.py      | 22 ++++++++++++
 python/tvm/topi/hexagon/conv2d.py        | 26 ++++++++++++++
 5 files changed, 94 insertions(+)
 create mode 100644 python/tvm/relay/op/strategy/hexagon.py
 create mode 100644 python/tvm/topi/hexagon/__init__.py
 create mode 100644 python/tvm/topi/hexagon/conv2d.py

diff --git a/python/tvm/relay/op/strategy/__init__.py b/python/tvm/relay/op/strategy/__init__.py
index 8d0543ba30af..cf915777ed0b 100644
--- a/python/tvm/relay/op/strategy/__init__.py
+++ b/python/tvm/relay/op/strategy/__init__.py
@@ -28,3 +28,4 @@
 from . import bifrost
 from . import rocm
 from . import intel_graphics
+from . import hexagon
diff --git a/python/tvm/relay/op/strategy/hexagon.py b/python/tvm/relay/op/strategy/hexagon.py
new file mode 100644
index 000000000000..cb1fec355917
--- /dev/null
+++ b/python/tvm/relay/op/strategy/hexagon.py
@@ -0,0 +1,44 @@
+# 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.
+
+"""Definition of Hexagon operator strategy."""
+
+# pylint: disable=invalid-name,unused-argument,wildcard-import,unused-wildcard-import
+
+from tvm import topi
+from .generic import *
+from .. import op as _op
+
+
+@conv2d_strategy.register("hexagon")
+def conv2d_strategy_hexagon(attrs, inputs, out_type, target):
+    """Conv2d strategy for Hexagon"""
+    strategy = _op.OpStrategy()
+    data_layout = attrs.data_layout
+    kernel_layout = attrs.kernel_layout
+
+    if data_layout == "NHWC" and kernel_layout == "HWIO":
+        strategy.add_implementation(
+            wrap_compute_conv2d(topi.nn.conv2d_nhwc),
+            wrap_topi_schedule(topi.hexagon.schedule_conv2d_nhwc),
+            name="conv2d.hexagon",
+        )
+        return strategy
+
+    raise RuntimeError(
+        "Unsupported layouts: data_layout:{}, kernel_layout:{}".format(data_layout, kernel_layout)
+    )
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index 9b843ae181fb..6b22cf13f5b9 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -61,6 +61,7 @@
 from . import sparse
 from . import hls
 from . import random
+from . import hexagon
 
 # error reporting
 from .utils import InvalidShapeError
diff --git a/python/tvm/topi/hexagon/__init__.py b/python/tvm/topi/hexagon/__init__.py
new file mode 100644
index 000000000000..3263819ccf3a
--- /dev/null
+++ b/python/tvm/topi/hexagon/__init__.py
@@ -0,0 +1,22 @@
+# 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.
+
+""" Schedules for Hexagon. """
+
+# pylint: disable=wildcard-import
+
+from .conv2d import *
diff --git a/python/tvm/topi/hexagon/conv2d.py b/python/tvm/topi/hexagon/conv2d.py
new file mode 100644
index 000000000000..8a484ae77e48
--- /dev/null
+++ b/python/tvm/topi/hexagon/conv2d.py
@@ -0,0 +1,26 @@
+# 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.
+
+""" Schedules for conv2d. """
+
+import tvm
+
+
+def schedule_conv2d_nhwc(outs):
+    """Schedule for Conv2d NHWC operator."""
+    s = tvm.te.create_schedule([x.op for x in outs])
+    return s