You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Anndrey24 (via GitHub)" <gi...@apache.org> on 2024/05/16 09:23:24 UTC
[PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Anndrey24 opened a new pull request, #17003:
URL: https://github.com/apache/tvm/pull/17003
This commit adds a scalable `arm_cpu` conv2d NHWC schedule for fp32 which generates SME instructions by using the tensor intrinsics introduced in #16921.
Alongside the SME schedule, the logic of the TE schedule `schedule_conv2d_gemm_native()` for both non-scalable and scalable vector implementations has also been translated into the new TIR schedule. This means that the TE compute definition `compute_conv2d_NHWC_hybrid()` is now compatible with both the original TE schedules (e.g. `schedule_conv2d_NHWC_hybrid()`) and the newly introduced TIR schedule `schedule_conv2d_NHWC_hybrid_TIR()`. The corresponding TOPI test has been extended to reflect that.
cc @ekalda @lhutton1
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "lhutton1 (via GitHub)" <gi...@apache.org>.
lhutton1 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1607969433
##########
tests/python/codegen/test_target_codegen_aarch64.py:
##########
@@ -756,6 +772,43 @@ def check_correct_assembly(dtype):
assert len(compute_ops) > 0
assert len(stores) > 0
+ with tvm.target.Target(target):
+ check_correct_assembly(dtype, *conv2d_impl)
+
+
+@pytest.mark.skipif(
+ llvm_version_major() < 16, reason="Test requires an LLVM version of at least 16 to target SVE"
Review Comment:
*SME
##########
tests/python/topi/test_topi_conv2d_nhwc.py:
##########
@@ -117,21 +140,34 @@ def test_conv2d_nhwc_gemm_fp32(device, ref_data, dtype, stride, padding, dilatio
A = te.placeholder(a_np.shape, name="A", dtype=dtype)
W = te.placeholder(w_np.shape, name="W", dtype=dtype)
- target, compute, schedule = device
- dev = tvm.device(target, 0)
+ target_string, compute, schedule, use_tir_schedule = device
+ dev = tvm.device(target_string, 0)
+ target = tvm.target.Target(target_string)
- with tvm.target.Target(target) as target:
- B = compute(A, W, stride, padding, dilation, dtype)
- s = schedule([B])
+ if (target.features.has_sve and llvm_version_major() < 15) or (
+ target.features.has_sme and llvm_version_major() < 16
+ ):
+ return
+
+ with target:
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
+ B = compute(A, W, stride, padding, dilation, dtype)
b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
- func = tvm.build(s, [A, W, B], target)
+ if use_tir_schedule:
+ primfunc = te.create_prim_func([A, W, B], index_dtype_override="int64")
Review Comment:
curious about the need for `index_dtype_override`? Is this because topi outputs int64 types by default?
##########
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)
+
+
+@autotvm.register_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:
is it possible to parallelize over batches here as well?
##########
tests/python/relay/strategy/arm_cpu/test_conv2d.py:
##########
@@ -16,8 +16,22 @@
# under the License.
"""Tests for arm_cpu schedules for regular conv2d."""
+from tests.python.relay.strategy.arm_cpu.scalable_utils import (
+ calculate_extra_workspace_size_from_scalable_extents,
+)
+import tvm
+import pytest
Review Comment:
nit: best practise seems to be to place 3rd party packages before tvm packages
##########
tests/python/relay/strategy/arm_cpu/test_conv2d.py:
##########
@@ -107,5 +121,127 @@ class TestConv2d_NCHW_Spatial_Pack(Conv2dTests):
schedule_name = parameter("conv2d_nchw_spatial_pack.arm_cpu")
+dtype = tvm.testing.parameter("float32")
+
+batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation = tvm.testing.parameters(
+ # Pad M, N, K
+ (1, 1, 1, 1, 1, 1, "SAME", 1),
+ (1, 1, 3, 15, 1, 1, "SAME", 1),
+ # Pad M, K
+ (1, 3, 9, 16, 3, 1, "SAME", 1),
+ # Pad M, N
+ (1, 2, 9, 15, 4, 1, "SAME", 1),
+ # Pad K, N
+ (1, 7, 4, 15, 3, 1, "SAME", 1),
+ # Pad M
+ (1, 2, 9, 16, 4, 1, "SAME", 1),
+ # Pad K
+ (1, 7, 4, 16, 3, 1, "SAME", 1),
+ # Pad N
+ (1, 2, 4, 15, 4, 1, "SAME", 1),
+ (1, 2, 4, 20, 1, 1, "SAME", 1),
+ # Large workloads
+ (1, 128, 32, 128, 3, 1, "SAME", 1),
+ (4, 64, 16, 64, 5, 2, "SAME", 1),
+ (1, 128, 32, 128, 3, 1, "VALID", 1),
+ (4, 64, 16, 64, 5, 2, "VALID", 1),
+ (1, 64, 16, 64, 3, 2, (0, 0, 1, 1), 1),
+ (1, 64, 16, 64, 3, 2, (1, 1, 2, 2), 1),
+ (1, 64, 16, 64, 5, 2, (3, 3, 2, 2), 1),
+ (1, 64, 16, 64, 3, 2, (0, 1, 2, 3), 1),
+ (1, 64, 32, 64, 3, 1, "SAME", 2),
+ (1, 64, 32, 64, 3, 1, (1, 1, 2, 2), 2),
+)
+
+
+@tvm.testing.fixture(cache_return_value=True)
Review Comment:
curious if we need to cache the return value for each set of parameters? they should always be different right?
##########
tests/python/topi/test_topi_conv2d_nhwc.py:
##########
@@ -117,21 +140,34 @@ def test_conv2d_nhwc_gemm_fp32(device, ref_data, dtype, stride, padding, dilatio
A = te.placeholder(a_np.shape, name="A", dtype=dtype)
W = te.placeholder(w_np.shape, name="W", dtype=dtype)
- target, compute, schedule = device
- dev = tvm.device(target, 0)
+ target_string, compute, schedule, use_tir_schedule = device
+ dev = tvm.device(target_string, 0)
+ target = tvm.target.Target(target_string)
- with tvm.target.Target(target) as target:
- B = compute(A, W, stride, padding, dilation, dtype)
- s = schedule([B])
+ if (target.features.has_sve and llvm_version_major() < 15) or (
+ target.features.has_sme and llvm_version_major() < 16
+ ):
+ return
Review Comment:
we should use pytest.skip() so that the tests appear as skipped in the CI logs as opposed to silently not running (assuming the tvm test framework doesn't already do this in some way)
##########
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)
+
+
+@autotvm.register_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)
+ sch.tensorize(ki, ARM_SME_2SVLx2SVL_TRANSPOSE_INTERLEAVE)
+
+ # Split and reorder the loops of the GeMM for tensorization
+ b, m, n, k = sch.get_loops(gemm_block)
+ 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
+ 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}"
+ tvm.tir.TensorIntrin.register(
+ sme_gemm_interleaved_intrin_name,
+ *get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(K_padded),
+ override=True,
+ )
+ sch.tensorize(mi, sme_gemm_interleaved_intrin_name)
+
+ # Add pstate annotations
+ root_block = sch.get_block("root")
+ sch.annotate(
+ root_block, SMEAttributes.STREAMING_MODE, SMEAttributes.StreamingModeValues.ENABLED
+ )
+ sch.annotate(root_block, SMEAttributes.ZA_STORAGE, SMEAttributes.ZAStorageValues.NEW)
+ elif use_scalable_vectors:
+ mo, mi = sch.split(m, [None, tile_M])
Review Comment:
are these schedules expected to perform similarly/differently to the te variants?
##########
tests/python/topi/test_topi_conv2d_nhwc.py:
##########
@@ -117,21 +140,34 @@ def test_conv2d_nhwc_gemm_fp32(device, ref_data, dtype, stride, padding, dilatio
A = te.placeholder(a_np.shape, name="A", dtype=dtype)
W = te.placeholder(w_np.shape, name="W", dtype=dtype)
- target, compute, schedule = device
- dev = tvm.device(target, 0)
+ target_string, compute, schedule, use_tir_schedule = device
+ dev = tvm.device(target_string, 0)
+ target = tvm.target.Target(target_string)
- with tvm.target.Target(target) as target:
- B = compute(A, W, stride, padding, dilation, dtype)
- s = schedule([B])
+ if (target.features.has_sve and llvm_version_major() < 15) or (
+ target.features.has_sme and llvm_version_major() < 16
+ ):
+ return
+
+ with target:
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
+ B = compute(A, W, stride, padding, dilation, dtype)
b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
- func = tvm.build(s, [A, W, B], target)
+ if use_tir_schedule:
+ primfunc = te.create_prim_func([A, W, B], index_dtype_override="int64")
+ sch = schedule(tvm.tir.Schedule(primfunc))
+ func = tvm.build(sch.mod["main"], target)
+ else:
+ s = schedule([B])
+ func = tvm.build(s, [A, W, B], target)
# Run only on AArch64 devices
# Do not run SVE schedules on non-SVE devices
- build_only = platform.machine() != "aarch64" or (
- target.features.has_sve and not tvm.testing.requires_aarch64_sve.run_time_check()
+ build_only = (
+ platform.machine() != "aarch64"
+ or (target.features.has_sve and not tvm.testing.requires_aarch64_sve.run_time_check())
+ or target.features.has_sme
Review Comment:
it might be nice to add a similar run_time_check for sme, if/when we're lucky enough to have it available in CI
##########
tests/python/relay/strategy/arm_cpu/test_conv2d.py:
##########
@@ -107,5 +121,127 @@ class TestConv2d_NCHW_Spatial_Pack(Conv2dTests):
schedule_name = parameter("conv2d_nchw_spatial_pack.arm_cpu")
+dtype = tvm.testing.parameter("float32")
+
+batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation = tvm.testing.parameters(
+ # Pad M, N, K
+ (1, 1, 1, 1, 1, 1, "SAME", 1),
+ (1, 1, 3, 15, 1, 1, "SAME", 1),
+ # Pad M, K
+ (1, 3, 9, 16, 3, 1, "SAME", 1),
+ # Pad M, N
+ (1, 2, 9, 15, 4, 1, "SAME", 1),
+ # Pad K, N
+ (1, 7, 4, 15, 3, 1, "SAME", 1),
+ # Pad M
+ (1, 2, 9, 16, 4, 1, "SAME", 1),
+ # Pad K
+ (1, 7, 4, 16, 3, 1, "SAME", 1),
+ # Pad N
+ (1, 2, 4, 15, 4, 1, "SAME", 1),
+ (1, 2, 4, 20, 1, 1, "SAME", 1),
+ # Large workloads
+ (1, 128, 32, 128, 3, 1, "SAME", 1),
+ (4, 64, 16, 64, 5, 2, "SAME", 1),
+ (1, 128, 32, 128, 3, 1, "VALID", 1),
+ (4, 64, 16, 64, 5, 2, "VALID", 1),
+ (1, 64, 16, 64, 3, 2, (0, 0, 1, 1), 1),
+ (1, 64, 16, 64, 3, 2, (1, 1, 2, 2), 1),
+ (1, 64, 16, 64, 5, 2, (3, 3, 2, 2), 1),
+ (1, 64, 16, 64, 3, 2, (0, 1, 2, 3), 1),
+ (1, 64, 32, 64, 3, 1, "SAME", 2),
+ (1, 64, 32, 64, 3, 1, (1, 1, 2, 2), 2),
+)
+
+
+@tvm.testing.fixture(cache_return_value=True)
+def ref_data(dtype, batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation):
+ in_height = in_width = in_size
+ a_shape = (batch, in_height, in_width, in_channel)
+ w_shape = (kernel, kernel, in_channel, num_filter)
+
+ a_np = np.random.uniform(size=a_shape).astype(dtype)
+ w_np = np.random.uniform(size=w_shape).astype(dtype)
+ return a_np, w_np
+
+
+@pytest.mark.skipif(
+ llvm_version_major() < 16, reason="SME is not supported in earlier versions of LLVM"
+)
+@tvm.testing.requires_aprofile_aem_fvp
+def test_conv2d_fp32(target, ref_data, dtype, stride, padding, dilation):
+ a_np, w_np = ref_data
+ dw_np = tvm.topi.testing.dilate_python(w_np, (dilation, dilation, 1, 1))
+
+ kernel_size = get_const_tuple(w_np.shape[:2])
+ out_channels = w_np.shape[3]
+
+ x = relay.var("data", shape=a_np.shape, dtype=dtype)
+ weight = relay.const(w_np, dtype=dtype)
+ conv2d = relay.nn.conv2d(
+ x,
+ weight,
+ channels=out_channels,
+ kernel_size=kernel_size,
+ strides=stride,
+ dilation=dilation,
+ padding=get_pad_tuple(padding, dw_np.shape[:2]),
+ data_layout="NHWC",
+ kernel_layout="HWIO",
+ out_dtype=dtype,
+ )
+
+ func = relay.Function(relay.analysis.free_vars(conv2d), conv2d)
+
+ ir_mod = tvm.IRModule.from_expr(func)
+ ir_mod = tvm.relay.transform.InferType()(ir_mod)
+
+ inputs = {"data": a_np}
+ params = {}
+ ref_outputs = generate_ref_data(ir_mod, inputs, params)
+
+ target = tvm.target.Target("llvm -mtriple=aarch64-none-elf -mattr=+v9.2a,+sme")
+ runtime = tvm.relay.backend.Runtime("crt", {"system-lib": True})
+ executor = tvm.relay.backend.Executor(
+ "aot",
+ {
+ "interface-api": "packed",
+ "unpacked-api": False,
+ },
+ )
+
+ with tvm.transform.PassContext(
+ opt_level=3, config=AOT_APROFILE_AEM_RUNNER.pass_config
+ ), tvm.meta_schedule.database.ScheduleFnDatabase(arm_cpu_tir_strategy):
Review Comment:
nit: I've opted to add `target` to the context here otherwise some SVE analyzer warnings appear. For some reason relay.build doesn't seem to add the target that's passed to the context, we should probably investigate further at some point
##########
tests/python/relay/strategy/arm_cpu/test_conv2d.py:
##########
@@ -107,5 +121,127 @@ class TestConv2d_NCHW_Spatial_Pack(Conv2dTests):
schedule_name = parameter("conv2d_nchw_spatial_pack.arm_cpu")
+dtype = tvm.testing.parameter("float32")
+
+batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation = tvm.testing.parameters(
+ # Pad M, N, K
+ (1, 1, 1, 1, 1, 1, "SAME", 1),
+ (1, 1, 3, 15, 1, 1, "SAME", 1),
+ # Pad M, K
+ (1, 3, 9, 16, 3, 1, "SAME", 1),
+ # Pad M, N
+ (1, 2, 9, 15, 4, 1, "SAME", 1),
+ # Pad K, N
+ (1, 7, 4, 15, 3, 1, "SAME", 1),
+ # Pad M
+ (1, 2, 9, 16, 4, 1, "SAME", 1),
+ # Pad K
+ (1, 7, 4, 16, 3, 1, "SAME", 1),
+ # Pad N
+ (1, 2, 4, 15, 4, 1, "SAME", 1),
+ (1, 2, 4, 20, 1, 1, "SAME", 1),
+ # Large workloads
+ (1, 128, 32, 128, 3, 1, "SAME", 1),
+ (4, 64, 16, 64, 5, 2, "SAME", 1),
+ (1, 128, 32, 128, 3, 1, "VALID", 1),
+ (4, 64, 16, 64, 5, 2, "VALID", 1),
+ (1, 64, 16, 64, 3, 2, (0, 0, 1, 1), 1),
+ (1, 64, 16, 64, 3, 2, (1, 1, 2, 2), 1),
+ (1, 64, 16, 64, 5, 2, (3, 3, 2, 2), 1),
+ (1, 64, 16, 64, 3, 2, (0, 1, 2, 3), 1),
+ (1, 64, 32, 64, 3, 1, "SAME", 2),
+ (1, 64, 32, 64, 3, 1, (1, 1, 2, 2), 2),
+)
+
+
+@tvm.testing.fixture(cache_return_value=True)
+def ref_data(dtype, batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation):
+ in_height = in_width = in_size
+ a_shape = (batch, in_height, in_width, in_channel)
+ w_shape = (kernel, kernel, in_channel, num_filter)
+
+ a_np = np.random.uniform(size=a_shape).astype(dtype)
+ w_np = np.random.uniform(size=w_shape).astype(dtype)
+ return a_np, w_np
+
+
+@pytest.mark.skipif(
+ llvm_version_major() < 16, reason="SME is not supported in earlier versions of LLVM"
+)
+@tvm.testing.requires_aprofile_aem_fvp
+def test_conv2d_fp32(target, ref_data, dtype, stride, padding, dilation):
+ a_np, w_np = ref_data
Review Comment:
nit: we should set a random seed for numpy to avoid flakiness in CI
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2136283203
Seems we have a breakage/flaky likely due to this PR https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-arm/detail/main/1980/pipeline
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "ekalda (via GitHub)" <gi...@apache.org>.
ekalda merged PR #17003:
URL: https://github.com/apache/tvm/pull/17003
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
Anndrey24 commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2135473130
Resolved the conflict!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
Anndrey24 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1608579653
##########
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)
+
+
+@autotvm.register_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)
+ sch.tensorize(ki, ARM_SME_2SVLx2SVL_TRANSPOSE_INTERLEAVE)
+
+ # Split and reorder the loops of the GeMM for tensorization
+ b, m, n, k = sch.get_loops(gemm_block)
+ 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
+ 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}"
+ tvm.tir.TensorIntrin.register(
+ sme_gemm_interleaved_intrin_name,
+ *get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(K_padded),
+ override=True,
+ )
+ sch.tensorize(mi, sme_gemm_interleaved_intrin_name)
+
+ # Add pstate annotations
+ root_block = sch.get_block("root")
+ sch.annotate(
+ root_block, SMEAttributes.STREAMING_MODE, SMEAttributes.StreamingModeValues.ENABLED
+ )
+ sch.annotate(root_block, SMEAttributes.ZA_STORAGE, SMEAttributes.ZAStorageValues.NEW)
+ elif use_scalable_vectors:
+ mo, mi = sch.split(m, [None, tile_M])
Review Comment:
They should theoretically perform the exact same as the TE schedules, just using `tir.Schedule()` instead.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
Anndrey24 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1608559925
##########
tests/python/relay/strategy/arm_cpu/test_conv2d.py:
##########
@@ -107,5 +121,127 @@ class TestConv2d_NCHW_Spatial_Pack(Conv2dTests):
schedule_name = parameter("conv2d_nchw_spatial_pack.arm_cpu")
+dtype = tvm.testing.parameter("float32")
+
+batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation = tvm.testing.parameters(
+ # Pad M, N, K
+ (1, 1, 1, 1, 1, 1, "SAME", 1),
+ (1, 1, 3, 15, 1, 1, "SAME", 1),
+ # Pad M, K
+ (1, 3, 9, 16, 3, 1, "SAME", 1),
+ # Pad M, N
+ (1, 2, 9, 15, 4, 1, "SAME", 1),
+ # Pad K, N
+ (1, 7, 4, 15, 3, 1, "SAME", 1),
+ # Pad M
+ (1, 2, 9, 16, 4, 1, "SAME", 1),
+ # Pad K
+ (1, 7, 4, 16, 3, 1, "SAME", 1),
+ # Pad N
+ (1, 2, 4, 15, 4, 1, "SAME", 1),
+ (1, 2, 4, 20, 1, 1, "SAME", 1),
+ # Large workloads
+ (1, 128, 32, 128, 3, 1, "SAME", 1),
+ (4, 64, 16, 64, 5, 2, "SAME", 1),
+ (1, 128, 32, 128, 3, 1, "VALID", 1),
+ (4, 64, 16, 64, 5, 2, "VALID", 1),
+ (1, 64, 16, 64, 3, 2, (0, 0, 1, 1), 1),
+ (1, 64, 16, 64, 3, 2, (1, 1, 2, 2), 1),
+ (1, 64, 16, 64, 5, 2, (3, 3, 2, 2), 1),
+ (1, 64, 16, 64, 3, 2, (0, 1, 2, 3), 1),
+ (1, 64, 32, 64, 3, 1, "SAME", 2),
+ (1, 64, 32, 64, 3, 1, (1, 1, 2, 2), 2),
+)
+
+
+@tvm.testing.fixture(cache_return_value=True)
Review Comment:
You're right, since we're not reusing the values for multiple devices, there is no need to cache them here. I'll remove the argument!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
Anndrey24 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1608503953
##########
tests/python/topi/test_topi_conv2d_nhwc.py:
##########
@@ -117,21 +140,34 @@ def test_conv2d_nhwc_gemm_fp32(device, ref_data, dtype, stride, padding, dilatio
A = te.placeholder(a_np.shape, name="A", dtype=dtype)
W = te.placeholder(w_np.shape, name="W", dtype=dtype)
- target, compute, schedule = device
- dev = tvm.device(target, 0)
+ target_string, compute, schedule, use_tir_schedule = device
+ dev = tvm.device(target_string, 0)
+ target = tvm.target.Target(target_string)
- with tvm.target.Target(target) as target:
- B = compute(A, W, stride, padding, dilation, dtype)
- s = schedule([B])
+ if (target.features.has_sve and llvm_version_major() < 15) or (
+ target.features.has_sme and llvm_version_major() < 16
+ ):
+ return
+
+ with target:
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
+ B = compute(A, W, stride, padding, dilation, dtype)
b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
- func = tvm.build(s, [A, W, B], target)
+ if use_tir_schedule:
+ primfunc = te.create_prim_func([A, W, B], index_dtype_override="int64")
Review Comment:
I think it had to do with me initially using `tile_N = T.int64(tile_N)` inside the schedule since that seems to be the index dtype when lowering from Relay (e.g. in `test_conv2d.py`), so I also made the TOPI tests use int64. However, after changing the line to `tile_N = T.cast(tile_N, N_padded.dtype)` this is not an issue anymore, so I will remove it!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2119233292
Failed to re-run CI in https://github.com/apache/tvm/actions/runs/9147913719
<details>
```
Traceback (most recent call last):
File "ci/scripts/github/github_tvmbot.py", line 595, in comment_failure
raise item
File "ci/scripts/github/github_tvmbot.py", line 701, in run
pr.rerun_jenkins_ci()
File "ci/scripts/github/github_tvmbot.py", line 554, in rerun_jenkins_ci
post(url, auth=("tvm-bot", TVM_BOT_JENKINS_TOKEN))
File "/home/runner/work/tvm/tvm/ci/scripts/jenkins/git_utils.py", line 53, in post
with request.urlopen(req, data) as response:
File "/usr/lib/python3.8/urllib/request.py", line 222, in urlopen
return opener.open(url, data, timeout)
File "/usr/lib/python3.8/urllib/request.py", line 531, in open
response = meth(req, response)
File "/usr/lib/python3.8/urllib/request.py", line 640, in http_response
response = self.parent.error(
File "/usr/lib/python3.8/urllib/request.py", line 569, in error
return self._call_chain(*args)
File "/usr/lib/python3.8/urllib/request.py", line 502, in _call_chain
result = func(*args)
File "/usr/lib/python3.8/urllib/request.py", line 649, in http_error_default
raise HTTPError(req.full_url, code, msg, hdrs, fp)
urllib.error.HTTPError: HTTP Error 502: Bad Gateway
```
with response
```
<html>
<head><title>502 Bad Gateway</title></head>
<body>
<center><h1>502 Bad Gateway</h1></center>
</body>
</html>
```
</details>
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
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)
+
+
+@autotvm.register_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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "Anndrey24 (via GitHub)" <gi...@apache.org>.
Anndrey24 commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2119232995
@tvm-bot rerun
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "ekalda (via GitHub)" <gi...@apache.org>.
ekalda commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2135675436
Thanks @Anndrey24 and @lhutton1, this is now merged!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [SME][TOPI] Add conv2d NHWC SME fp32 schedule [tvm]
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #17003:
URL: https://github.com/apache/tvm/pull/17003#issuecomment-2136289066
Thanks @Anndrey24 @lhutton1 @ekalda .
I created a temp revert, https://github.com/apache/tvm/pull/17038 to unblock the ci, if there is an alternative fix that would also be good, eitherway we followup with a redo quickly.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org