Anndrey24 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1608580761
##########
python/tvm/topi/arm_cpu/conv2d.py:
##########
@@ -655,3 +661,228 @@ def compute_conv2d_NHWC_hybrid_SVE(cfg, data, kernel,
strides, padding, dilation
def schedule_conv2d_NHWC_hybrid_SVE(cfg, outs):
"""Interface for hybrid schedule_conv2d_NHWC_hybrid_SVE"""
return schedule_conv2d_NHWC(cfg, outs, False)
+
+
[email protected]_topi_compute("conv2d_NHWC_hybrid_SME.arm_cpu")
+def compute_conv2d_NHWC_hybrid_SME(cfg, data, kernel, strides, padding,
dilation, out_dtype):
+ """Interface for hybrid compute_conv2d_NHWC_hybrid_SME"""
+ return compute_conv2d_NHWC(
+ cfg,
+ data,
+ kernel,
+ strides,
+ padding,
+ dilation,
+ out_dtype,
+ False,
+ True,
+ True,
+ )
+
+
+def schedule_conv2d_NHWC_hybrid_TIR(sch: tvm.tir.Schedule):
+ """
+ Perform TIR scheduling for conv2d NHWC.
+ """
+ # Get ordered buffer list
+ primfunc = sch.mod["main"]
+ buffer_names = primfunc.params
+ buffer_list = [primfunc.buffer_map[buf] for buf in buffer_names]
+ dtype = buffer_list[0].dtype
+
+ # Determine PrimFunc blocks
+ block_list = [
+ "data_pad",
+ "data_im2col",
+ "T_reshape",
+ "A_padded_K",
+ "A_padded_M",
+ "weight_flatten",
+ "C",
+ "conv2d_gemm_output",
+ ]
+ func_blocks = {}
+ for block in block_list:
+ func_blocks[block] = sch.get_block(block) if has_block(sch, block)
else None
+
+ gemm_block = func_blocks["C"]
+ b, m, n, k = sch.get_loops(gemm_block)
+
+ # Get tiling information
+ use_scalable_vectors =
sch.get(func_blocks["conv2d_gemm_output"]).annotations[
+ "use_scalable_vectors"
+ ]
+ use_sme = sch.get(func_blocks["conv2d_gemm_output"]).annotations["use_sme"]
+ M_padded = sch.get(m).extent
+ N_padded = sch.get(n).extent
+ K_padded = sch.get(k).extent
+ tile_M, tile_K = get_tiling_A(False, dtype, use_sme)
+ tile_N, _ = get_tiling_B_transformed(False, dtype, use_scalable_vectors,
use_sme)
+ tile_M = T.cast(tile_M, M_padded.dtype)
+ tile_N = T.cast(tile_N, N_padded.dtype)
+ tile_K = T.cast(tile_K, K_padded.dtype)
+
+ # GeMM
+ # Compute each tile_M x tile_N tile
+ # By summing up K outer products
+ if use_sme:
+ # pylint: disable=import-outside-toplevel
+ from tvm.topi.arm_cpu.pstate_attributes import SMEAttributes
+ from tvm.tir.tensor_intrin.arm_cpu import (
+ ARM_SME_2SVLx2SVL_TRANSPOSE_INTERLEAVE,
+ ARM_SME_2SVLx2SVL_GEMM_INTERLEAVED_MOPA,
+ ARM_SME_INIT,
+ get_sme_gemm_interleaved_mopa_2svlx2svl_intrin,
+ )
+
+ # Interleave the padded im2col matrix utilizing the matrix tile
+ interleave_t_A_block = sch.cache_read(gemm_block, 0, "global")
+ sch.transform_layout(interleave_t_A_block, ("write", 0), lambda b, m,
k: (b, k, m))
+ b, m, k = sch.get_loops(interleave_t_A_block)
+ mo, mi = sch.split(m, factors=(None, tile_M), disable_predication=True)
+ ko, ki = sch.split(k, factors=(None, tile_K), disable_predication=True)
+ sch.reorder(b, ko, mo, ki, mi)
Review Comment:
Yes, thanks for the catch!
--
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]
