Anndrey24 commented on code in PR #17048:
URL: https://github.com/apache/tvm/pull/17048#discussion_r1627396909
##########
python/tvm/topi/arm_cpu/conv2d.py:
##########
@@ -680,6 +681,43 @@ def compute_conv2d_NHWC_hybrid_SME(cfg, data, kernel,
strides, padding, dilation
)
[email protected]_topi_compute("conv2d_NHWC_hybrid_SME_transposed_B.arm_cpu")
+def compute_conv2d_NHWC_SME_transposed_B(
+ cfg,
+ data,
+ kernel,
+ strides,
+ padding,
+ dilation,
+ out_dtype,
+ kernel_size,
+ output_channels,
+):
+ """Compute conv2d NHWC hybrid SME transposed B"""
+ N, K = get_const_tuple(kernel.shape)
+ tile_N, tile_K = get_tiling_B_transformed(False, data.dtype, True, True)
+ pad_N, pad_K = tvm.topi.arm_cpu.arm_utils.get_conv2d_weights_padding(N, K,
tile_N, tile_K)
+
+ kernel = tvm.topi.nn.pad(
+ kernel, pad_before=(0, 0), pad_after=(pad_N, pad_K),
name="weight_padding"
Review Comment:
Padding the B matrix should preferably happen before entering
`compute_conv2d_gemm_without_weight_transform()`, so I think the call to
`get_conv2d_weights_padding()` there is only to help re-compute `N_padded`.
Usually the weights are transformed in [this
function](https://github.com/apache/tvm/pull/17048/files#diff-04a1aab966320b6f63c390cc8b79f79b543a1a7a3f560ea3a97d18c2bb041a0aL618)
that gets called
[here](https://github.com/apache/tvm/pull/17048/files#diff-cfbbf96eaad867ef4c9e28936a21e8d7beae5bf917c12bfbdd48498006d1fcdbL541),
which is executed when the AlterOpLayout pass does not do any compile-time
transformation beforehand (e.g. SVE schedule) or when the pass just isn't run.
This schedule is slightly different from the rest of the conv2d `arm_cpu`
schedules as we have both compile-time transformations (transpose and reshape)
and a runtime transformation (padding to 2xSVL) for the kernel, but we could
also have to do everything at runtime if the AlterOpLayout pass isn't run,
which is why I preferred to separate the "partial" runtime transformation
(just pad) from the "full" runtime transformation (flatten, transpose and pad)
from `conv2d_gemm_weight_transform()`.
I could probably move the "partial" transformation path to
`conv2d_gemm_weight_transform()` too by adding an extra argument to the
function to indicate the calling context, but I feared it would make it less
intuitive / overcomplicated. What do you think?
--
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]
