lhutton1 commented on code in PR #17048:
URL: https://github.com/apache/tvm/pull/17048#discussion_r1626149820
##########
python/tvm/relay/op/strategy/arm_cpu.py:
##########
@@ -583,11 +584,25 @@ def
conv2d_gemm_without_weight_transform_strategy_arm_cpu(attrs, inputs, out_typ
)
elif data.dtype in ["float32", "float16"]:
# Non-quantized cases
- strategy.add_implementation(
-
wrap_compute_conv2d_gemm(topi.arm_cpu.compute_conv2d_NHWC_hybrid_without_transform),
-
wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_NHWC_hybrid_without_transform),
- name="conv2d_NHWC_hybrid_without_transform.arm_cpu",
- )
+ if (
+ target.features.has_sme
+ and kernel.dtype == "float16"
+ and data.dtype == "float16"
+ and out_type.dtype == "float32"
+ ):
+ strategy.add_implementation(
+
wrap_compute_conv2d_gemm(topi.arm_cpu.compute_conv2d_NHWC_SME_transposed_B),
+ lambda: None,
+ name="conv2d_NHWC_hybrid_SME_transposed_B.arm_cpu",
Review Comment:
Nit: It might be worth leaving a note to explain why we expect B to be
transposed in this case
##########
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 seems to be typically applied
[here](https://github.com/apache/tvm/pull/17048/files#diff-42b1313a1be464c7f2c94f75d656be725f1ccb54b9391cd5b27c33009ac0e2d5R137),
curious why we need to add a separate stage of padding here?
##########
python/tvm/topi/arm_cpu/conv2d_alter_op.py:
##########
@@ -162,6 +162,30 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
inputs[0], new_kernel_expr, **new_attrs
)
+ if (
+ topi_tmpl == "conv2d_NHWC_hybrid_SME.arm_cpu"
+ and data_dtype == "float16"
+ and kernel_dtype == "float16"
+ and out_dtype == "float32"
+ ):
+ assert data_layout == "NHWC" and kernel_layout == "HWIO"
+ KH, KW, IC, OC = get_const_tuple(kernel.shape)
+ K = KH * KW * IC
+ N = OC
+ transposed_kernel_expr = relay.transpose(inputs[1], axes=[3, 0, 1, 2])
+ transposed_flattened_kernel_expr =
relay.reshape(transposed_kernel_expr, newshape=(N, K))
+ new_kernel_expr = transposed_flattened_kernel_expr
+ new_kernel = te.placeholder((N, K), kernel.dtype)
+ new_workload_name = "conv2d_NHWC_hybrid_SME_transposed_B.arm_cpu"
Review Comment:
Nit: also worth leaving a similar comment here
##########
python/tvm/topi/arm_cpu/conv2d.py:
##########
@@ -743,24 +787,40 @@ def schedule_conv2d_NHWC_hybrid_TIR(sch:
tvm.tir.Schedule):
ko, ki = sch.split(k, factors=(None, tile_K), disable_predication=True)
sch.parallel(b)
sch.reorder(b, ko, mo, ki, mi)
- sch.tensorize(ki, ARM_SME_2SVLx2SVL_FP32_TRANSPOSE_INTERLEAVE)
+ sch.tensorize(ki, transpose_interleave_intrin_name)
+
+ # Interleave the padded weights matrix utilizing the matrix tile
+ if in_dtype == "float16":
+ interleave_b_block = sch.cache_read(gemm_block, 1, "global")
+ sch.transform_layout(interleave_b_block, ("write", 0), lambda n,
k: (k, n))
+ n, k = sch.get_loops(interleave_b_block)
+ ko, ki = sch.split(k, factors=(None, tile_K),
disable_predication=True)
+ no, ni = sch.split(n, factors=(None, tile_N),
disable_predication=True)
+ sch.reorder(ko, no, ki, ni)
+ sch.tensorize(ki, transpose_interleave_intrin_name)
# Split and reorder the loops of the GeMM for tensorization
b, m, n, k = sch.get_loops(gemm_block)
+ tile_M, _ = get_tiling_A(False, out_dtype, True)
+ tile_N, _ = get_tiling_B_transformed(False, out_dtype, True, True)
+ tile_M = T.cast(tile_M, M_padded.dtype)
+ tile_N = T.cast(tile_N, N_padded.dtype)
mo, mi = sch.split(m, factors=(None, tile_M), disable_predication=True)
no, ni = sch.split(n, factors=(None, tile_N), disable_predication=True)
sch.parallel(b)
sch.reorder(b, mo, no, mi, ni, k)
- # Tensorize the GeMM output matrix initialization to zero
+ # Tensorize the GeMM initialization
init_block = sch.decompose_reduction(gemm_block, mi)
sch.tensorize(sch.get_loops(init_block)[-2], ARM_SME_INIT)
# Tensorize the GeMM update
- sme_gemm_interleaved_intrin_name =
ARM_SME_2SVLx2SVL_GEMM_INTERLEAVED_MOPA + f"_{K_padded}"
+ sme_gemm_interleaved_intrin_name = (
+ ARM_SME_2SVLx2SVL_GEMM_INTERLEAVED_MOPA + f"_{K_padded}_{in_dtype}"
Review Comment:
thanks for fixing this ;)
##########
python/tvm/topi/arm_cpu/conv2d.py:
##########
@@ -698,6 +737,8 @@ def schedule_conv2d_NHWC_hybrid_TIR(sch: tvm.tir.Schedule):
"A_padded_K",
"A_padded_M",
"weight_flatten",
+ "weight_padding",
+ "weight_transpose",
Review Comment:
Just to check, this "weight_transpose" block shouldn't exist if
"conv2d_alter_op" was performed correctly?
--
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]
