ekalda commented on code in PR #14981:
URL: https://github.com/apache/tvm/pull/14981#discussion_r1211974551
##########
python/tvm/relay/op/strategy/arm_cpu.py:
##########
@@ -132,13 +132,16 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type,
target):
plevel=15,
)
else:
+ # TODO(@FranklandJack)
+ # Investigate why this producing output tensor of
+ # incorrect dimensions in
+ # test_runtime_module_based_interface.py
Review Comment:
Are there any other tests that exercise the same schedule that pass? I'm
quite worried about commenting out an Arm specific schedule, essentially
defaulting to use x86 schedule for this data and kernel layout.
##########
python/tvm/relay/op/strategy/arm_cpu.py:
##########
@@ -284,8 +288,21 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type,
target):
name="depthwise_conv2d_nchw.x86",
)
elif layout == "NHWC":
- assert kernel_layout == "HWOI"
- if target.features.has_asimd:
+ # TODO(@FranklandJack)
+ # Handle HWOI in arm_cpu scheudles/compute definition.
+ if kernel_layout == "HWOI":
+ logger.warning(
+ """depthwise_conv2d with layout NHWC and HWOI
+ kernel layout is not optimized for arm_cpu
target.
+ """
+ )
+ strategy.add_implementation(
+ wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc,
need_kernel_layout=True),
+
wrap_topi_schedule(conv2d_generic.schedule_depthwise_conv2d_nhwc),
+ name="depthwise_conv2d_nhwc.generic",
+ )
+
+ elif target.features.has_asimd:
Review Comment:
It looks like the logic is changing quite a bit there... So if we have NHWC
with HWOI, we add first the generic schedule, then if there is vector
extension, we also add the arm specific schedule? Why is that? Also the DSP
schedule below requires HWOI, but doesn't check for it explicitly since it is
relying on the assert that is being removed.
##########
python/tvm/topi/arm_cpu/conv2d.py:
##########
@@ -509,3 +510,24 @@ def conv2d_nhwc_dsp(cfg, data, kernel, strides, padding,
dilation, out_dtype):
def schedule_conv2d_nhwc_dsp(cfg, outs):
"""Create schedule for conv2d_nhwc_dsp"""
return conv2d_nhwc_dsp_schedule(cfg, outs)
+
+
+@conv2d_infer_layout.register("arm_cpu")
+def _conv2d_infer_layout(workload, cfg):
+ _, data, kernel, strides, padding, dilation, layout, _, dtype = workload
+ batch_size, in_channel, in_height, in_width = data[1]
+ out_channel, _, k_height, k_width = kernel[1]
+ idxdiv = tvm.tir.indexdiv
+
+ pt, pl, pb, pr = get_pad_tuple(padding, (k_height, k_width))
+ hdilation, wdilation = dilation if isinstance(dilation, (tuple, list))
else (dilation, dilation)
+ dilated_kernel_h = (k_height - 1) * hdilation + 1
+ dilated_kernel_w = (k_width - 1) * wdilation + 1
+ out_height = idxdiv(in_height + pt + pb - dilated_kernel_h, strides[0]) + 1
+ out_width = idxdiv(in_width + pl + pr - dilated_kernel_w, strides[1]) + 1
+ tile_ic, tile_oc = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
+ in_shape = (batch_size, idxdiv(in_channel, tile_ic), in_height, in_width,
tile_ic)
+ in_layout = f"NCHW{tile_ic}c"
+ out_shape = (batch_size, idxdiv(out_channel, tile_oc), out_height,
out_width, tile_oc)
+ out_layout = f"NCHW{tile_oc}c"
+ return ((in_shape, in_layout),), ((out_shape, out_layout),)
Review Comment:
Is this function identical to the ones in `x86` and `intel_graphics`? If
yes, then while you are at it, maybe move the body of that function to some
common location (e.g. `python/tvm/topi/utils.py`) to reduce the amount of
duplicated code.
##########
python/tvm/relay/op/strategy/arm_cpu.py:
##########
@@ -284,8 +288,21 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type,
target):
name="depthwise_conv2d_nchw.x86",
)
elif layout == "NHWC":
- assert kernel_layout == "HWOI"
- if target.features.has_asimd:
+ # TODO(@FranklandJack)
+ # Handle HWOI in arm_cpu scheudles/compute definition.
Review Comment:
Nit:
```suggestion
# Handle HWOI in arm_cpu schedules/compute definition.
```
##########
tests/python/unittest/test_meta_schedule_relay_integration.py:
##########
@@ -420,6 +420,7 @@ def test_extract_task_arm_conv2d_nchwc():
def test_meta_schedule_te2primfunc_argument_order_and_lowering():
+ target = Target("llvm --num-cores=16")
Review Comment:
Just curious, what does that change do?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
