Skip to content

Commit

Permalink
[Relay] Add support of conv2d with NHWC for Bifrost (#8430)
Browse files Browse the repository at this point in the history
Reuse generic Mali strategy for conv2d with NHWC in
Bifrost target.
  • Loading branch information
AnastasiaStulova authored Jul 10, 2021
1 parent 59b204d commit 972d7b5
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 0 deletions.
8 changes: 8 additions & 0 deletions python/tvm/relay/op/strategy/bifrost.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
4 changes: 4 additions & 0 deletions tests/python/topi/python/test_topi_conv2d_nhwc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}

Expand Down

0 comments on commit 972d7b5

Please sign in to comment.