wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/python/s_tir/dlight/test_benchmark.py
+321
View File
@@ -0,0 +1,321 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=missing-docstring
# ruff: noqa: E501, F841
import tempfile
import pytest
pytest.importorskip("cloudpickle") # tvm.s_tir.dlight.benchmark imports cloudpickle
import tvm.testing
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.dlight.benchmark import (
benchmark,
benchmark_prim_func,
benchmark_relax_func,
extract_from_relax,
extract_func_info_from_prim_func,
extract_prim_func,
)
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
# The test function uses an undefined symbolic var in Relax.
# In principle, this should be attached to an argument.
# pylint: disable=no-self-argument,invalid-name,line-too-long,no-method-argument
# fmt: off
@I.ir_module(check_well_formed=False, s_tir=True)
class Module:
@T.prim_func(s_tir=True)
def full1(var_T_full: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
n = T.int64()
T_full = T.match_buffer(var_T_full, (T.int64(1), T.int64(32), T.int64(1), n), "float16")
# with T.sblock("root"):
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), T.int64(1), n):
with T.sblock("T_full"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
T.reads()
T.writes(T_full[v_ax0, v_ax1, v_ax2, v_ax3])
T_full[v_ax0, v_ax1, v_ax2, v_ax3] = T.float16(1.0)
@T.prim_func(s_tir=True)
def full2(var_T_full: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
n = T.int64()
T_full = T.match_buffer(var_T_full, (T.int64(1), T.int64(32), n, T.int64(128)), "float16")
# with T.sblock("root"):
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), n, T.int64(128)):
with T.sblock("T_full"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
T.reads()
T.writes(T_full[v_ax0, v_ax1, v_ax2, v_ax3])
T_full[v_ax0, v_ax1, v_ax2, v_ax3] = T.float16(1.0)
@T.prim_func(s_tir=True)
def matmul1(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16")):
T.func_attr({"op_pattern": 4, "tirx.noalias": True})
n = T.int64()
A = T.match_buffer(var_A, (T.int64(1), T.int64(32), T.int64(1), n), "float16")
B = T.match_buffer(var_B, (T.int64(1), T.int64(32), n, T.int64(128)), "float16")
# with T.sblock("root"):
for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(128), n):
with T.sblock("matmul"):
v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
T.reads(A[v_i0, v_i1, v_i2, v_k], B[v_i0, v_i1, v_k, v_i3])
T.writes(matmul[v_i0, v_i1, v_i2, v_i3])
with T.init():
matmul[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
matmul[v_i0, v_i1, v_i2, v_i3] = matmul[v_i0, v_i1, v_i2, v_i3] + A[v_i0, v_i1, v_i2, v_k] * B[v_i0, v_i1, v_k, v_i3]
@R.function
def test():
n = T.int64()
R.func_attr({"tir_var_upper_bound": {"n": 2048}})
cls = Module
with R.dataflow():
lv1 = R.call_tir(cls.full1,(), out_ty=R.Tensor((1, 32, 1, n), dtype="float16"))
lv1_1 = R.call_tir(cls.full1,(), out_ty=R.Tensor((1, 32, 1, n), dtype="float16"))
lv1_2 = R.call_tir(cls.full1,(), out_ty=R.Tensor((1, 32, 1, n), dtype="float16"))
lv2 = R.call_tir(cls.full2,(), out_ty=R.Tensor((1, 32, n, 128), dtype="float16"))
lv2_1 = R.call_tir(cls.full2,(), out_ty=R.Tensor((1, 32, n, 128), dtype="float16"))
lv3 = R.call_tir(cls.matmul1, (lv1, lv2), out_ty=R.Tensor((1, 32, 1, 128), dtype="float16"))
R.output(lv3)
return lv3
@T.prim_func(s_tir=True)
def cuda_workload(var_inp0: T.handle, inp1: T.Buffer((T.int64(4096), T.int64(4096)), "float32"), var_matmul: T.handle):
T.func_attr({"tirx.is_scheduled": True})
m = T.int64()
inp0 = T.match_buffer(var_inp0, (T.int64(1), m, T.int64(4096)))
matmul = T.match_buffer(var_matmul, (T.int64(1), m, T.int64(4096)))
# with T.sblock("root"):
matmul_reindex_pad_local = T.sblock_alloc_buffer((T.int64(1), (m + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="local")
inp0_reindex_pad_shared = T.sblock_alloc_buffer((T.int64(1), (m + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="shared")
inp1_reindex_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(4096), T.int64(4096)), scope="shared")
for ax0 in T.thread_binding(T.int64(1), thread="blockIdx.z"):
for ax1_0 in T.thread_binding((m + T.int64(31)) // T.int64(32), thread="blockIdx.x"):
for ax2_0 in T.thread_binding(T.int64(64), thread="blockIdx.y"):
for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.y"):
for ax1_1 in T.thread_binding(T.int64(1), thread="vthread.x"):
for ax2_2 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
for ax1_2 in T.thread_binding(T.int64(8), thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
for ax2_3_init, ax1_3_init in T.grid(T.int64(4), T.int64(4)):
with T.sblock("matmul_init"):
v0 = T.axis.spatial(T.int64(1), ax0)
v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax1_0 * T.int64(32) + ax1_1 * T.int64(32) + ax1_2 * T.int64(4) + ax1_3_init)
v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_init)
T.reads()
T.writes(matmul_reindex_pad_local[T.int64(0), v1, v2])
matmul_reindex_pad_local[T.int64(0), v1, v2] = T.float32(0)
for ax3_0 in range(T.int64(256)):
for ax0_ax1_ax2_fused_0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
for ax0_ax1_ax2_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
for ax0_ax1_ax2_fused_2 in range(T.int64(2)):
for ax0_ax1_ax2_fused_3 in T.vectorized(T.int64(2)):
with T.sblock("inp0_reindex_pad_shared"):
v0 = T.axis.spatial(T.int64(1), T.int64(0))
v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax1_0 * T.int64(32) + (ax0_ax1_ax2_fused_0 * T.int64(32) + ax0_ax1_ax2_fused_1 * T.int64(4) + ax0_ax1_ax2_fused_2 * T.int64(2) + ax0_ax1_ax2_fused_3) // T.int64(16))
v2 = T.axis.spatial(T.int64(4096), ax3_0 * T.int64(16) + (ax0_ax1_ax2_fused_0 * T.int64(32) + ax0_ax1_ax2_fused_1 * T.int64(4) + ax0_ax1_ax2_fused_2 * T.int64(2) + ax0_ax1_ax2_fused_3) % T.int64(16))
T.reads(inp0[v0, v1, v2])
T.writes(inp0_reindex_pad_shared[v0, v1, v2])
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
inp0_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < m, inp0[v0, v1, v2], T.float32(0))
for ax0_ax1_ax2_fused_0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
for ax0_ax1_ax2_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
for ax0_ax1_ax2_fused_2 in range(T.int64(4)):
for ax0_ax1_ax2_fused_3 in T.vectorized(T.int64(2)):
with T.sblock("inp1_reindex_shared"):
v0 = T.axis.spatial(T.int64(1), T.int64(0))
v1 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + (ax0_ax1_ax2_fused_0 * T.int64(64) + ax0_ax1_ax2_fused_1 * T.int64(8) + ax0_ax1_ax2_fused_2 * T.int64(2) + ax0_ax1_ax2_fused_3) // T.int64(16))
v2 = T.axis.spatial(T.int64(4096), ax3_0 * T.int64(16) + (ax0_ax1_ax2_fused_0 * T.int64(64) + ax0_ax1_ax2_fused_1 * T.int64(8) + ax0_ax1_ax2_fused_2 * T.int64(2) + ax0_ax1_ax2_fused_3) % T.int64(16))
T.reads(inp1[v2, v1])
T.writes(inp1_reindex_shared[v0, v1, v2])
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
inp1_reindex_shared[v0, v1, v2] = inp1[v2, v1]
for ax3_1, ax2_3, ax1_3 in T.grid(T.int64(16), T.int64(4), T.int64(4)):
with T.sblock("matmul_update"):
v0 = T.axis.spatial(T.int64(1), ax0)
v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax1_0 * T.int64(32) + ax1_1 * T.int64(32) + ax1_2 * T.int64(4) + ax1_3)
v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3)
v3 = T.axis.reduce(T.int64(4096), ax3_0 * T.int64(16) + ax3_1)
T.reads(matmul_reindex_pad_local[T.int64(0), v1, v2], inp0_reindex_pad_shared[T.int64(0), v1, v3], inp1_reindex_shared[T.int64(0), v2, v3])
T.writes(matmul_reindex_pad_local[T.int64(0), v1, v2])
matmul_reindex_pad_local[T.int64(0), v1, v2] = matmul_reindex_pad_local[T.int64(0), v1, v2] + inp0_reindex_pad_shared[T.int64(0), v1, v3] * inp1_reindex_shared[T.int64(0), v2, v3]
for ax0_1, ax1, ax2_0_1 in T.grid(T.int64(1), T.int64(4), T.int64(2)):
for ax2_1_1 in T.vectorized(T.int64(2)):
with T.sblock("matmul_reindex_pad_local"):
v0 = T.axis.spatial(T.int64(1), ax0_1)
v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax1_0 * T.int64(32) + ax1_2 * T.int64(4) + ax1)
v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_2 * T.int64(4) + ax2_0_1 * T.int64(2) + ax2_1_1)
T.reads(matmul_reindex_pad_local[v0, v1, v2])
T.writes(matmul[T.int64(0), v1, v2])
if v1 < m:
matmul[T.int64(0), v1, v2] = matmul_reindex_pad_local[v0, v1, v2]
# fmt: on
# pylint: enable=no-self-argument,invalid-name,line-too-long,no-method-argument
@pytest.mark.skip("requires CUDA")
def test_benchmark_prim_func_rpc():
with LocalRPC() as rpc:
rpc_config = ms.runner.RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
input_infos, _, _ = benchmark(
cuda_workload,
args=[
((1, "m", 4096), "float32"),
((4096, 4096), "float32"),
((1, "m", 4096), "float32"),
],
dym_var_sample={"m": 128},
target="nvidia/geforce-rtx-3070",
rpc_config=rpc_config,
)
assert input_infos == [
((1, 128, 4096), "float32"),
((4096, 4096), "float32"),
((1, 128, 4096), "float32"),
]
@pytest.mark.skip("requires CUDA")
def test_benchmark_prim_func_local():
input_infos, _, _ = benchmark(
cuda_workload,
args=[
((1, "m", 4096), "float32"),
((4096, 4096), "float32"),
((1, "m", 4096), "float32"),
],
dym_var_sample={"m": 128},
target="nvidia/geforce-rtx-3070",
)
assert input_infos == [
((1, 128, 4096), "float32"),
((4096, 4096), "float32"),
((1, 128, 4096), "float32"),
]
@pytest.mark.skip("requires CUDA")
def test_benchmark_prim_func_full_local():
with tvm.target.Target("nvidia/geforce-rtx-3070"):
benchmark_prim_func(
cuda_workload,
)
@pytest.mark.skip("requires CUDA")
def test_benchmark_prim_func_full_rpc():
with LocalRPC() as rpc:
rpc_config = ms.runner.RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
benchmark_prim_func(
cuda_workload,
target="nvidia/geforce-rtx-3070",
rpc_config=rpc_config,
evaluator_config=ms.runner.EvaluatorConfig(
number=10,
repeat=10,
min_repeat_ms=0,
enable_cpu_cache_flush=False,
),
)
def test_benchmark_relax_func():
with tvm.target.Target({"kind": "llvm", "num-cores": 4}):
benchmark_relax_func(Module, "test")
def test_extract_prim_func_full1():
print(
extract_prim_func(
model_name="TEST",
relax_func_name="test",
prim_func_name="full1",
func=Module["full1"], # type: ignore
func_args=[((1, 32, 1, "n"), "float16")],
dym_var_dict={"n": "int32"},
weight=2,
sample_number=10,
target={"kind": "llvm", "num-cores": 4},
)
)
def test_extract_prim_func_matmul1():
print(
extract_prim_func(
model_name="TEST",
relax_func_name="test",
prim_func_name="matmul1",
func=Module["matmul1"], # type: ignore
weight=2,
sample_number=10,
target={"kind": "llvm", "num-cores": 4},
)
)
def test_extract_from_relax():
with tvm.target.Target({"kind": "llvm", "num-cores": 4}):
with tempfile.TemporaryDirectory() as filepath:
extract_from_relax(
Module,
"TEST",
file_path=filepath,
)
def test_extract_func_info_from_prim_func():
assert (
str(extract_func_info_from_prim_func(cuda_workload))
== "([((1, m, 4096), 'float32'), ((4096, 4096), 'float32'), ((1, m, 4096), 'float32')], {'m': 'int64'})"
)
assert (
str(extract_func_info_from_prim_func(Module["full1"]))
== "([((1, 32, 1, n), 'float16')], {'n': 'int64'})"
)
assert (
str(extract_func_info_from_prim_func(Module["matmul1"]))
== "([((1, 32, 1, n), 'float16'), ((1, 32, n, 128), 'float16'), ((1, 32, 1, 128), 'float16')], {'n': 'int64'})"
)
assert (
str(extract_func_info_from_prim_func(Module["full2"]))
== "([((1, 32, n, 128), 'float16')], {'n': 'int64'})"
)
if __name__ == "__main__":
tvm.testing.main()