[Relay] Add support of conv2d with NHWC for Bifrost

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

@@ -65,6 +65,14 @@ def conv2d_strategy_bifrost(attrs, inputs, out_type, target):
                     wrap_topi_schedule(topi.bifrost.schedule_conv2d_nchw_spatial_pack),
                     name="conv2d_nchw_spatial_pack.bifrost",
                 )
+        elif layout == "NHWC":
+            assert kernel_layout == "HWIO"
+            # For now just reuse general Mali strategy.
+            strategy.add_implementation(
+                wrap_compute_conv2d(topi.mali.conv2d_nhwc_spatial_pack),
+                wrap_topi_schedule(topi.mali.schedule_conv2d_nhwc_spatial_pack),
+                name="conv2d_nhwc_spatial_pack.bifrost",
+            )
         else:
             raise RuntimeError("Unsupported conv2d layout {} for Mali(Bifrost)".format(layout))
     elif is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups):

tests/python/topi/python/test_topi_conv2d_nhwc.py

@@ -38,6 +38,10 @@
         topi.mali.conv2d_nhwc_spatial_pack,
         topi.mali.schedule_conv2d_nhwc_spatial_pack,
     ),
+    "bifrost": (
+        topi.mali.conv2d_nhwc_spatial_pack,
+        topi.mali.schedule_conv2d_nhwc_spatial_pack,
+    ),
     "hls": (topi.nn.conv2d_nhwc, topi.hls.schedule_conv2d_nhwc),
 }