chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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()
|
||||
@@ -0,0 +1,590 @@
|
||||
# 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, F401, F841
|
||||
import pytest
|
||||
|
||||
import tvm.testing
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_gemv_basic():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p_lv1614: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int32()
|
||||
lv1638 = T.match_buffer(p_lv1638, (1, 32, n, 128), "float16")
|
||||
lv1614 = T.match_buffer(p_lv1614, (1, 1, 1, n), "float16")
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (1, 32, 1, n))
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
|
||||
var_T_divide_intermediate = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
|
||||
var_T_maximum_intermediate = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
|
||||
var_T_minimum_intermediate = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
|
||||
for i0, i1, i2, i3, k in T.grid(1, 32, 1, n, 128):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
|
||||
T.reads(lv1637[v_i0, v_i1, v_i2, v_k], lv1638[v_i0, v_i1, v_i3, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + lv1637[v_i0, v_i1, v_i2, v_k] * lv1638[v_i0, v_i1, v_i3, v_k]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 32, 1, n):
|
||||
with T.sblock("T_divide"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
T.writes(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
var_T_divide_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] * T.float16(0.088397790055248615)
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 32, 1, n):
|
||||
with T.sblock("T_maximum"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
T.writes(var_T_maximum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
var_T_maximum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = T.max(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2, v_ax3], T.float16(-65504))
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 32, 1, n):
|
||||
with T.sblock("T_minimum"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(var_T_maximum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3], lv1614[v_ax0, 0, v_ax2, v_ax3])
|
||||
T.writes(var_T_minimum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
var_T_minimum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = T.min(var_T_maximum_intermediate[v_ax0, v_ax1, v_ax2, v_ax3], lv1614[v_ax0, 0, v_ax2, v_ax3])
|
||||
for i0, i1, i2, i3 in T.grid(1, 32, 1, n):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(var_T_minimum_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
T.writes(var_compute_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
var_compute_intermediate[v_i0, v_i1, v_i2, v_i3] = T.Cast("float32", var_T_minimum_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv1637: T.Buffer((1, 32, 1, 128), "float16"), p_lv1638: T.handle, p_lv1614: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int32()
|
||||
lv1638 = T.match_buffer(p_lv1638, (1, 32, n, 128), "float16")
|
||||
lv1614 = T.match_buffer(p_lv1614, (1, 1, 1, n), "float16")
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (1, 32, 1, n))
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((1, 32, 1, n), "float16")
|
||||
for ax0_fused in range(32):
|
||||
for ax1_fused_0 in T.parallel((n + 63) // 64):
|
||||
for ax1_fused_1 in T.vectorized(64):
|
||||
for ax2_fused_0 in T.serial(2, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax2_fused_1, u_0, u_1 in T.grid(64, 1, 1):
|
||||
with T.sblock("NT_matmul"):
|
||||
v0 = T.axis.spatial(32, ax0_fused)
|
||||
v1 = T.axis.spatial(n, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
v2 = T.axis.reduce(128, ax2_fused_0 * 64 + ax2_fused_1)
|
||||
T.where(ax1_fused_0 * 64 + ax1_fused_1 < n)
|
||||
T.reads(lv1637[0, v0, 0, v2], lv1638[0, v0, v1, v2])
|
||||
T.writes(var_NT_matmul_intermediate[0, v0, 0, v1])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[0, v0, 0, v1] = T.float16(0.0)
|
||||
var_NT_matmul_intermediate[0, v0, 0, v1] = var_NT_matmul_intermediate[0, v0, 0, v1] + lv1637[0, v0, 0, v2] * lv1638[0, v0, v1, v2]
|
||||
for ax0, ax1 in T.grid(32, n):
|
||||
with T.sblock("compute"):
|
||||
v0, v1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(var_NT_matmul_intermediate[0, v0, 0, v1], lv1614[0, 0, 0, v1])
|
||||
T.writes(var_compute_intermediate[0, v0, 0, v1])
|
||||
var_compute_intermediate[0, v0, 0, v1] = T.Cast("float32", T.min(T.max(var_NT_matmul_intermediate[0, v0, 0, v1] * T.float16(0.088397790055248615), T.float16(-65504.0)), lv1614[0, 0, 0, v1]))
|
||||
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_decode_gemv_256_threads():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv571: T.Buffer((22016, 512), "uint32"), lv572: T.Buffer((22016, 128), "float16"), lv1654: T.Buffer((1, 1, 4096), "float16"), var_NT_matmul_intermediate: T.Buffer((1, 1, 22016), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate = T.sblock_alloc_buffer((22016, 4096), "float16")
|
||||
for i, j in T.grid(22016, 4096):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv571[v_i, v_j // 8], lv572[v_i, v_j // 32])
|
||||
T.writes(p_output0_intermediate[v_i, v_j])
|
||||
p_output0_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv571[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv572[v_i, v_j // 32]
|
||||
for i0, i1, i2, k in T.grid(1, 1, 22016, 4096):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv1654[v_i0, v_i1, v_k], p_output0_intermediate[v_i2, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv1654[v_i0, v_i1, v_k] * p_output0_intermediate[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv571: T.Buffer((22016, 512), "uint32"), lv572: T.Buffer((22016, 128), "float16"), lv1654: T.Buffer((1, 1, 4096), "float16"), var_NT_matmul_intermediate: T.Buffer((1, 1, 22016), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
for u_fused in range(1):
|
||||
for ax0_fused_0 in T.parallel(172):
|
||||
for ax0_fused_1 in T.vectorized(128):
|
||||
for ax1_0_fused_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1, 8):
|
||||
with T.sblock("NT_matmul"):
|
||||
v0 = T.axis.spatial(22016, ax0_fused_0 * 128 + ax0_fused_1)
|
||||
v1 = T.axis.reduce(4096, ax1_0_fused_0 * 512 + ax1_0_fused_1 * 8 + ax1_1_0 * 8 + ax1_1_1)
|
||||
T.reads(lv1654[0, 0, v1], lv571[v0, v1 // 8], lv572[v0, v1 // 32])
|
||||
T.writes(var_NT_matmul_intermediate[0, 0, v0])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[0, 0, v0] = T.float16(0.0)
|
||||
var_NT_matmul_intermediate[0, 0, v0] = var_NT_matmul_intermediate[0, 0, v0] + lv1654[0, 0, v1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv571[v0, v1 // 8], T.Cast("uint32", v1 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv572[v0, v1 // 32])
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_decode_gemv1():
|
||||
# fmt: off
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv571: T.Buffer((22016, 512), "uint32"), lv572: T.Buffer((22016, 128), "float16"), lv1654: T.Buffer((1, 1, 4096), "float16"), var_NT_matmul_intermediate: T.Buffer((1, 1, 22016), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate = T.sblock_alloc_buffer((22016, 4096), "float16")
|
||||
for i, j in T.grid(22016, 4096):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv571[v_i, v_j // 8], lv572[v_i, v_j // 32])
|
||||
T.writes(p_output0_intermediate[v_i, v_j])
|
||||
p_output0_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv571[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv572[v_i, v_j // 32]
|
||||
for i0, i1, i2, k in T.grid(1, 1, 22016, 4096):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv1654[v_i0, v_i1, v_k], p_output0_intermediate[v_i2, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv1654[v_i0, v_i1, v_k] * p_output0_intermediate[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv571: T.Buffer((22016, 512), "uint32"), lv572: T.Buffer((22016, 128), "float16"), lv1654: T.Buffer((1, 1, 4096), "float16"), var_NT_matmul_intermediate: T.Buffer((1, 1, 22016), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
for u_fused in range(1):
|
||||
for ax0_fused_0 in T.parallel(172):
|
||||
for ax0_fused_1 in T.vectorized(128):
|
||||
for ax1_0_fused_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1, 8):
|
||||
with T.sblock("NT_matmul"):
|
||||
v0 = T.axis.spatial(22016, ax0_fused_0 * 128 + ax0_fused_1)
|
||||
v1 = T.axis.reduce(4096, ax1_0_fused_0 * 512 + ax1_0_fused_1 * 8 + ax1_1_0 * 8 + ax1_1_1)
|
||||
T.reads(lv1654[0, 0, v1], lv571[v0, v1 // 8], lv572[v0, v1 // 32])
|
||||
T.writes(var_NT_matmul_intermediate[0, 0, v0])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[0, 0, v0] = T.float16(0.0)
|
||||
var_NT_matmul_intermediate[0, 0, v0] = var_NT_matmul_intermediate[0, 0, v0] + lv1654[0, 0, v1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv571[v0, v1 // 8], T.Cast("uint32", v1 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv572[v0, v1 // 32])
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_decode_gemv2():
|
||||
# fmt: off
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv771: T.Buffer((32000, 512), "uint32"), lv772: T.Buffer((32000, 128), "float16"), lv3216: T.Buffer((1, 1, 4096), "float16"), p_output0_intermediate: T.Buffer((1, 1, 32000), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((32000, 4096), "float16")
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((1, 1, 32000), "float16")
|
||||
for i, j in T.grid(32000, 4096):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv771[v_i, v_j // 8], lv772[v_i, v_j // 32])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv771[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv772[v_i, v_j // 32]
|
||||
for i0, i1, i2, k in T.grid(1, 1, 32000, 4096):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv3216[v_i0, v_i1, v_k], p_output0_intermediate_1[v_i2, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv3216[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_i2, v_k]
|
||||
for i0, i1, i2 in T.grid(1, 1, 32000):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
T.writes(p_output0_intermediate[v_i0, v_i1, v_i2])
|
||||
p_output0_intermediate[v_i0, v_i1, v_i2] = T.Cast("float32", var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv771: T.Buffer((32000, 512), "uint32"), lv772: T.Buffer((32000, 128), "float16"), lv3216: T.Buffer((1, 1, 4096), "float16"), p_output0_intermediate: T.Buffer((1, 1, 32000), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((1, 1, 32000), "float16")
|
||||
for u_fused in range(1):
|
||||
for ax0_fused_0 in T.parallel(250):
|
||||
for ax0_fused_1 in T.vectorized(128):
|
||||
for ax1_0_fused_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(64, 1, 8):
|
||||
with T.sblock("NT_matmul"):
|
||||
v0 = T.axis.spatial(32000, ax0_fused_0 * 128 + ax0_fused_1)
|
||||
v1 = T.axis.reduce(4096, ax1_0_fused_0 * 512 + ax1_0_fused_1 * 8 + ax1_1_0 * 8 + ax1_1_1)
|
||||
T.reads(lv3216[0, 0, v1], lv771[v0, v1 // 8], lv772[v0, v1 // 32])
|
||||
T.writes(var_NT_matmul_intermediate[0, 0, v0])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[0, 0, v0] = T.float16(0.0)
|
||||
var_NT_matmul_intermediate[0, 0, v0] = var_NT_matmul_intermediate[0, 0, v0] + lv3216[0, 0, v1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv771[v0, v1 // 8], T.Cast("uint32", v1 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv772[v0, v1 // 32])
|
||||
for ax0 in range(32000):
|
||||
with T.sblock("compute"):
|
||||
v0 = T.axis.spatial(32000, ax0)
|
||||
T.reads(var_NT_matmul_intermediate[0, 0, v0])
|
||||
T.writes(p_output0_intermediate[0, 0, v0])
|
||||
p_output0_intermediate[0, 0, v0] = T.Cast("float32", var_NT_matmul_intermediate[0, 0, v0])
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_decode_gemv3():
|
||||
# fmt: off
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv575: T.Buffer((T.int64(4096), T.int64(1376)), "uint32"), lv576: T.Buffer((T.int64(4096), T.int64(344)), "float16"), lv574: T.Buffer((T.int64(1), T.int64(1), T.int64(11008)), "float16"), lv570: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((T.int64(4096), T.int64(11008)), "float16")
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")
|
||||
for i, j in T.grid(T.int64(4096), T.int64(11008)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv575[v_i, v_j // T.int64(8)], lv576[v_i, v_j // T.int64(32)])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv575[v_i, v_j // T.int64(8)], T.Cast("uint32", v_j % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv576[v_i, v_j // T.int64(32)]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(4096), T.int64(11008)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv574[v_i0, v_i1, v_k], p_output0_intermediate_1[v_i2, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv574[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_i2, v_k]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(lv570[v_ax0, v_ax1, v_ax2], var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = lv570[v_ax0, v_ax1, v_ax2] + var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv575: T.Buffer((T.int64(4096), T.int64(1376)), "uint32"), lv576: T.Buffer((T.int64(4096), T.int64(344)), "float16"), lv574: T.Buffer((T.int64(1), T.int64(1), T.int64(11008)), "float16"), lv570: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")
|
||||
for u_fused in range(1):
|
||||
for ax0_fused_0 in T.parallel(T.int64(64)):
|
||||
for ax0_fused_1 in T.vectorized(T.int64(64)):
|
||||
for ax1_0_fused_0 in T.serial(T.int64(11), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0_fused_1, ax1_1_0, ax1_1_1 in T.grid(T.int64(128), T.int64(1), T.int64(8)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax0_fused_0 * T.int64(64) + ax0_fused_1)
|
||||
v1 = T.axis.reduce(T.int64(11008), (ax1_0_fused_0 * T.int64(128) + ax1_0_fused_1) * T.int64(8) + ax1_1_0 * T.int64(8) + ax1_1_1)
|
||||
T.where(ax1_0_fused_0 * T.int64(128) + ax1_0_fused_1 < T.int64(1376))
|
||||
T.reads(lv574[T.int64(0), T.int64(0), v1], lv575[v0, v1 // T.int64(8)], lv576[v0, v1 // T.int64(32)])
|
||||
T.writes(var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0] = T.float16(0.0)
|
||||
var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0] = var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0] + lv574[T.int64(0), T.int64(0), v1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv575[v0, v1 // T.int64(8)], T.Cast("uint32", v1 % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv576[v0, v1 // T.int64(32)])
|
||||
for ax0 in range(T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax0)
|
||||
T.reads(lv570[T.int64(0), T.int64(0), v0], var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0])
|
||||
T.writes(p_output0_intermediate[T.int64(0), T.int64(0), v0])
|
||||
p_output0_intermediate[T.int64(0), T.int64(0), v0] = lv570[T.int64(0), T.int64(0), v0] + var_NT_matmul_intermediate[T.int64(0), T.int64(0), v0]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_autogptq_decode_gemv():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(lv9: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), lv10: T.Buffer((T.int64(32), T.int64(512)), "uint32"), lv11: T.Buffer((T.int64(32), T.int64(4096)), "float16"), lv12: T.Buffer((T.int64(4096),), "uint32"), lv8: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), lv1613: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
decode_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16")
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv9[v_i // T.int64(8), v_j], lv10[lv12[v_i], v_j // T.int64(8)], lv12[v_i], lv11[lv12[v_i], v_j])
|
||||
T.writes(decode_intermediate[v_i, v_j])
|
||||
decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv9[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8) * T.int64(4))), T.uint32(15))) - (T.Cast("float16", T.bitwise_and(T.shift_right(lv10[lv12[v_i], v_j // T.int64(8)], T.Cast("uint32", v_j % T.int64(8) * T.int64(4))), T.uint32(15))) + T.float16(1))) * lv11[lv12[v_i], v_j]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv8[v_i0, v_i1, v_k], decode_intermediate[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv8[v_i0, v_i1, v_k] * decode_intermediate[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(lv1613[v_ax0, v_ax1, v_ax2], var_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = lv1613[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
# fmt: on
|
||||
|
||||
# The GeMV rule does not yet support the inner dim being grouped.
|
||||
# So the rule is expected to skip transforming this function.
|
||||
mod = tvm.IRModule({"main": func})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], func)
|
||||
|
||||
|
||||
def test_outer_reduction_adreno():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
lv575: T.Buffer((1376, 4096), "uint32"),
|
||||
lv576: T.Buffer((344, 4096), "float16"),
|
||||
lv574: T.Buffer((1, 1, 11008), "float16"),
|
||||
lv570: T.Buffer((1, 1, 4096), "float16"),
|
||||
p_output0_intermediate: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((11008, 4096), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((1, 1, 4096), "float16")
|
||||
for i, j in T.grid(11008, 4096):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv575[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15)))- T.float16(7)) * lv576[v_i // 32, v_j]
|
||||
for i0, i1, i2, k in T.grid(1, 1, 4096, 11008):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv574[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = lv570[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv575: T.Buffer((1376, 4096), "uint32"), lv576: T.Buffer((344, 4096), "float16"), lv574: T.Buffer((1, 1, 11008), "float16"), lv570: T.Buffer((1, 1, 4096), "float16"), p_output0_intermediate: T.Buffer((1, 1, 4096), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((11008, 4096), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((1, 1, 4096), "float16")
|
||||
for i, j in T.grid(11008, 4096):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv575[v_i // 8, v_j], lv576[v_i // 32, v_j])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv575[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv576[v_i // 32, v_j]
|
||||
for i0, i1, i2, k in T.grid(1, 1, 4096, 11008):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv574[v_i0, v_i1, v_k], p_output0_intermediate_1[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0.0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv574[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(lv570[v_ax0, v_ax1, v_ax2], var_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = lv570[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_outer_reduction_adreno_dynamic():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(p_lv612: T.handle, p_lv613: T.handle, lv1607: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
v = T.int64()
|
||||
lv612 = T.match_buffer(p_lv612, (T.int64(512), v), "uint32")
|
||||
lv613 = T.match_buffer(p_lv613, (T.int64(128), v), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (T.int64(1), T.int64(1), v))
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((T.int64(4096), v), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), v), "float16")
|
||||
for i, j in T.grid(T.int64(4096), v):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv612[v_i // T.int64(8), v_j], lv613[v_i // T.int64(32), v_j])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv612[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv613[v_i // T.int64(32), v_j]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), v, T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv1607[v_i0, v_i1, v_k], p_output0_intermediate_1[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv1607[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_k, v_i2]
|
||||
for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), v):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
T.writes(p_output0_intermediate[v_i0, v_i1, v_i2])
|
||||
p_output0_intermediate[v_i0, v_i1, v_i2] = T.Cast("float32", var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(p_lv612: T.handle, p_lv613: T.handle, lv1607: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
v = T.int64()
|
||||
lv612 = T.match_buffer(p_lv612, (T.int64(512), v), "uint32")
|
||||
lv613 = T.match_buffer(p_lv613, (T.int64(128), v), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (T.int64(1), T.int64(1), v))
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((T.int64(4096), v), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), v), "float16")
|
||||
for i, j in T.grid(T.int64(4096), v):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv612[v_i // T.int64(8), v_j], lv613[v_i // T.int64(32), v_j])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv612[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7.0)) * lv613[v_i // T.int64(32), v_j]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), v, T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv1607[v_i0, v_i1, v_k], p_output0_intermediate_1[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0.0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv1607[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_k, v_i2]
|
||||
for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), v):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
T.writes(p_output0_intermediate[v_i0, v_i1, v_i2])
|
||||
p_output0_intermediate[v_i0, v_i1, v_i2] = T.Cast("float32", var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_blockized_gemv():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(x: T.Buffer((1, 4096), "float16"), w: T.Buffer((8, 16384, 4096), "float16"), indptr: T.Buffer((2,), "int32"), o: T.Buffer((2, 16384), "float16")):
|
||||
# with T.sblock("root"):
|
||||
for expert_id in T.thread_binding(2, thread="blockIdx.y"):
|
||||
with T.sblock("gemv_o"):
|
||||
v_expert_id_o = T.axis.spatial(2, expert_id)
|
||||
vi_o = T.axis.spatial(1, 0)
|
||||
vj_o = T.axis.reduce(1, 0)
|
||||
T.reads(x[0, 0:4096], w[indptr[v_expert_id_o], 0:16384, 0:4096], indptr[v_expert_id_o])
|
||||
T.writes(o[v_expert_id_o, 0:16384])
|
||||
for i, j in T.grid(16384, 4096):
|
||||
with T.sblock("gemv"):
|
||||
vi_i, vj_i = T.axis.remap("SR", [i, j])
|
||||
T.reads(x[0, vj_i], w[indptr[v_expert_id_o], vi_i, vj_i], indptr[v_expert_id_o])
|
||||
T.writes(o[v_expert_id_o, vi_i])
|
||||
with T.init():
|
||||
o[v_expert_id_o, vi_i] = T.float16(0)
|
||||
o[v_expert_id_o, vi_i] = o[v_expert_id_o, vi_i] + x[0, vj_i] * w[indptr[v_expert_id_o], vi_i, vj_i]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(x: T.Buffer((1, 4096), "float16"), w: T.Buffer((8, 16384, 4096), "float16"), indptr: T.Buffer((2,), "int32"), o: T.Buffer((2, 16384), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
for expert_id in T.thread_binding(2, thread="blockIdx.y"):
|
||||
with T.sblock("gemv_o"):
|
||||
v_expert_id_o = T.axis.spatial(2, expert_id)
|
||||
vi_o = T.axis.spatial(1, 0)
|
||||
vj_o = T.axis.reduce(1, 0)
|
||||
T.reads(x[0, 0:4096], w[indptr[v_expert_id_o], 0:16384, 0:4096], indptr[v_expert_id_o])
|
||||
T.writes(o[v_expert_id_o, 0:16384])
|
||||
for u_fused in range(1):
|
||||
for ax0_fused_0 in T.parallel(128):
|
||||
for ax0_fused_1 in T.vectorized(128):
|
||||
for ax1_fused_0 in T.serial(64, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_fused_1, u_0, u_1 in T.grid(64, 1, 1):
|
||||
with T.sblock("gemv"):
|
||||
v0 = T.axis.spatial(16384, ax0_fused_0 * 128 + ax0_fused_1)
|
||||
v1 = T.axis.reduce(4096, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
T.reads(x[0, v1], w[indptr[v_expert_id_o], v0, v1], indptr[v_expert_id_o])
|
||||
T.writes(o[v_expert_id_o, v0])
|
||||
with T.init():
|
||||
o[v_expert_id_o, v0] = T.float16(0.0)
|
||||
o[v_expert_id_o, v0] = o[v_expert_id_o, v0] + x[0, v1] * w[indptr[v_expert_id_o], v0, v1]
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_func_to_skip():
|
||||
@T.prim_func(s_tir=True)
|
||||
def before(var_A: T.handle, var_exclusive_scan_thrust: T.handle, seq_len: T.int64):
|
||||
data_buf = T.match_buffer(var_A, (seq_len * T.int64(8),), "int32", align=8)
|
||||
output_buf = T.match_buffer(
|
||||
var_exclusive_scan_thrust, (seq_len * T.int64(8),), "int32", align=8
|
||||
)
|
||||
with T.sblock("exclusive_scan_thrust"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.call_packed(
|
||||
"tvm.contrib.thrust.sum_scan",
|
||||
T.tvm_stack_make_array(
|
||||
data_buf.data, T.tvm_stack_make_shape(seq_len * T.int64(8)), 0, 1, 0, T.int64(0)
|
||||
),
|
||||
T.tvm_stack_make_array(
|
||||
output_buf.data,
|
||||
T.tvm_stack_make_shape(seq_len * T.int64(8)),
|
||||
0,
|
||||
1,
|
||||
0,
|
||||
T.int64(0),
|
||||
),
|
||||
T.bool(False),
|
||||
)
|
||||
|
||||
# This function should be skipped.
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("llvm"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.cpu.GEMV())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], before)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,275 @@
|
||||
# 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
|
||||
"""Tests for CPU DLight Reduction schedule rule."""
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import te, tirx, topi
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.s_tir.dlight.cpu import Reduction
|
||||
from tvm.target import Target
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _llvm_target():
|
||||
return Target({"kind": "llvm"})
|
||||
|
||||
|
||||
def _rvv_target():
|
||||
return Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"mtriple": "riscv64-linux-gnu",
|
||||
"mcpu": "generic-rv64",
|
||||
"mabi": "lp64d",
|
||||
"mattr": ["+64bit", "+m", "+a", "+f", "+d", "+c", "+v"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _build_softmax(batch, features, fast=False):
|
||||
A = te.placeholder((batch, features), dtype="float32", name="A")
|
||||
B = topi.nn.fast_softmax(A, axis=1) if fast else topi.nn.softmax(A, axis=1)
|
||||
func = te.create_prim_func([A, B])
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def _apply_and_check(mod, target):
|
||||
"""Apply Reduction rule and verify it was applied."""
|
||||
rule = Reduction()
|
||||
result = rule.apply(mod["main"], target, False)
|
||||
assert result is not None, "Reduction rule should apply"
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test: schedule applicability
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fast", [False, True], ids=["softmax", "fast_softmax"])
|
||||
@pytest.mark.parametrize(
|
||||
"batch,features",
|
||||
[
|
||||
(1, 10),
|
||||
(1, 128),
|
||||
(14, 185),
|
||||
(32, 256),
|
||||
(64, 512),
|
||||
(128, 1024),
|
||||
(1, 30522),
|
||||
],
|
||||
)
|
||||
def test_reduction_applies(batch, features, fast):
|
||||
"""Reduction rule should apply to softmax/fast_softmax of various shapes."""
|
||||
mod = _build_softmax(batch, features, fast=fast)
|
||||
target = _llvm_target()
|
||||
_apply_and_check(mod, target)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test: scheduled TIR structure
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_softmax_schedule_structure():
|
||||
"""Verify the scheduled TIR has expected structure:
|
||||
- parallel on batch axis
|
||||
- vectorized innermost loops for injective blocks
|
||||
- split+unroll for reduction blocks
|
||||
"""
|
||||
mod = _build_softmax(14, 185, fast=False)
|
||||
target = _llvm_target()
|
||||
sch = _apply_and_check(mod, target)
|
||||
scheduled_mod = sch.mod
|
||||
|
||||
# Check that tirx.is_scheduled is NOT set (only set by ApplyDefaultSchedule)
|
||||
# but the schedule should be valid
|
||||
assert scheduled_mod is not None
|
||||
|
||||
# Verify via ApplyDefaultSchedule path
|
||||
with target:
|
||||
scheduled = dl.ApplyDefaultSchedule(Reduction())(mod)
|
||||
func = scheduled["main"]
|
||||
|
||||
# Check tirx.is_scheduled is set
|
||||
assert func.attrs and func.attrs.get("tirx.is_scheduled", False)
|
||||
|
||||
|
||||
def test_fast_softmax_schedule_structure():
|
||||
"""fast_softmax should keep T_fast_exp as a separate vectorizable block."""
|
||||
mod = _build_softmax(14, 185, fast=True)
|
||||
target = _llvm_target()
|
||||
sch = _apply_and_check(mod, target)
|
||||
script = str(sch.mod)
|
||||
|
||||
# fast_exp block should exist (not inlined)
|
||||
assert "T_fast_exp" in script or "T_softmax_delta" in script
|
||||
# Should have T.parallel
|
||||
assert "T.parallel" in script
|
||||
# Should have T.vectorized
|
||||
assert "T.vectorized" in script
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test: LLVM IR quality (cross-compile to RISC-V RVV)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _codegen_llvm_ir(mod, target):
|
||||
"""Lower and codegen to LLVM IR (no linking)."""
|
||||
bound = tirx.transform.BindTarget(target.with_host(target))(mod)
|
||||
pipeline, finalize_host, _ = tirx.get_tir_pipeline("default")
|
||||
lowered = pipeline(bound)
|
||||
from tvm.tirx.build import split_host_device_mods
|
||||
|
||||
host_mod, _ = split_host_device_mods(lowered)
|
||||
host_mod = finalize_host()(host_mod)
|
||||
built = tvm.target.codegen.build_module(host_mod, target)
|
||||
return built.inspect_source("ll")
|
||||
|
||||
|
||||
def _codegen_asm(mod, target):
|
||||
"""Lower and codegen to assembly (no linking)."""
|
||||
bound = tirx.transform.BindTarget(target.with_host(target))(mod)
|
||||
pipeline, finalize_host, _ = tirx.get_tir_pipeline("default")
|
||||
lowered = pipeline(bound)
|
||||
from tvm.tirx.build import split_host_device_mods
|
||||
|
||||
host_mod, _ = split_host_device_mods(lowered)
|
||||
host_mod = finalize_host()(host_mod)
|
||||
built = tvm.target.codegen.build_module(host_mod, target)
|
||||
return built.inspect_source("s")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fast", [False, True], ids=["softmax", "fast_softmax"])
|
||||
def test_rvv_code_size_reduction(fast):
|
||||
"""Scheduled RVV code should be smaller than unscheduled.
|
||||
|
||||
The original issue (apache/tvm#18569) shows RVV softmax is 1.34x slower
|
||||
than scalar, partly due to LLVM generating bloated code with excessive
|
||||
unrolling. The schedule should reduce code size significantly.
|
||||
"""
|
||||
target = _rvv_target()
|
||||
mod = _build_softmax(14, 185, fast=fast)
|
||||
|
||||
# Unscheduled
|
||||
ir_unsched = _codegen_llvm_ir(mod, target)
|
||||
n_unsched = len(ir_unsched.splitlines())
|
||||
|
||||
# Scheduled
|
||||
with target:
|
||||
mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
|
||||
ir_sched = _codegen_llvm_ir(mod_sched, target)
|
||||
n_sched = len(ir_sched.splitlines())
|
||||
|
||||
# Scheduled should be meaningfully smaller (at least 30% reduction)
|
||||
ratio = n_sched / n_unsched
|
||||
assert ratio < 0.75, (
|
||||
f"Expected >=25% code reduction, got {(1 - ratio) * 100:.1f}% "
|
||||
f"({n_unsched} -> {n_sched} lines)"
|
||||
)
|
||||
|
||||
|
||||
# The arith analyzer no longer proves vscale-bearing inequalities via
|
||||
# substitution (CanProveVscaleExpressionFromKnownValues was deleted). This
|
||||
# weakens simplification of scalable-vector index expressions, which can
|
||||
# prevent the RVV vectorization schedule from producing scalable vector ops.
|
||||
@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution")
|
||||
def test_rvv_fast_softmax_vectorizes_exp():
|
||||
"""fast_softmax + schedule should produce RVV vector instructions
|
||||
for the polynomial exp approximation (no scalar exp calls)."""
|
||||
target = _rvv_target()
|
||||
mod = _build_softmax(14, 185, fast=True)
|
||||
with target:
|
||||
mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
|
||||
ir = _codegen_llvm_ir(mod_sched, target)
|
||||
|
||||
# Should have zero scalar exp calls (fast_exp uses polynomial)
|
||||
scalar_exp = sum(1 for line in ir.splitlines() if "llvm.exp.f32" in line)
|
||||
assert scalar_exp == 0, f"Expected 0 scalar exp calls, got {scalar_exp}"
|
||||
|
||||
# Should have scalable vector operations
|
||||
n_svec = ir.count("<vscale x")
|
||||
assert n_svec > 0, "Expected scalable vector operations in LLVM IR"
|
||||
|
||||
|
||||
def test_rvv_asm_instruction_reduction():
|
||||
"""Scheduled RVV assembly should have fewer total instructions
|
||||
than both unscheduled RVV and scalar RV."""
|
||||
rvv = _rvv_target()
|
||||
rv = Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"mtriple": "riscv64-linux-gnu",
|
||||
"mcpu": "generic-rv64",
|
||||
"mabi": "lp64d",
|
||||
"mattr": ["+64bit", "+m", "+a", "+f", "+d", "+c"],
|
||||
}
|
||||
)
|
||||
|
||||
mod = _build_softmax(14, 185, fast=True)
|
||||
|
||||
# Scalar baseline
|
||||
asm_rv = _codegen_asm(mod, rv)
|
||||
n_rv = len(
|
||||
[
|
||||
line
|
||||
for line in asm_rv.splitlines()
|
||||
if line.strip() and not line.strip().startswith((".", "#", "/"))
|
||||
]
|
||||
)
|
||||
|
||||
# RVV unscheduled
|
||||
asm_rvv = _codegen_asm(mod, rvv)
|
||||
n_rvv = len(
|
||||
[
|
||||
line
|
||||
for line in asm_rvv.splitlines()
|
||||
if line.strip() and not line.strip().startswith((".", "#", "/"))
|
||||
]
|
||||
)
|
||||
|
||||
# RVV scheduled
|
||||
with rvv:
|
||||
mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
|
||||
asm_sched = _codegen_asm(mod_sched, rvv)
|
||||
n_sched = len(
|
||||
[
|
||||
line
|
||||
for line in asm_sched.splitlines()
|
||||
if line.strip() and not line.strip().startswith((".", "#", "/"))
|
||||
]
|
||||
)
|
||||
|
||||
# Scheduled should be smaller than both unscheduled RVV and scalar
|
||||
assert n_sched < n_rvv, (
|
||||
f"Scheduled ({n_sched}) should have fewer instructions than unscheduled RVV ({n_rvv})"
|
||||
)
|
||||
assert n_sched <= n_rv * 1.1, (
|
||||
f"Scheduled ({n_sched}) should not be much larger than scalar RV ({n_rv})"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,112 @@
|
||||
# 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
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_conv3d():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
A: T.Buffer((14308, 3, 2, 14, 14), "float16"),
|
||||
W: T.Buffer((1280, 3, 2, 14, 14), "float16"),
|
||||
C: T.Buffer((14308, 1280, 1, 1, 1), "float16"),
|
||||
):
|
||||
pad_A = T.sblock_alloc_buffer((14308, 3, 2, 14, 14), "float16")
|
||||
for i0, i1, i2, i3, i4 in T.grid(14308, 3, 2, 14, 14):
|
||||
with T.sblock("pad_A"):
|
||||
v_i0, v_i1, v_i2, v_i3, v_i4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
pad_A[v_i0, v_i1, v_i2, v_i3, v_i4] = A[v_i0, v_i1, v_i2, v_i3, v_i4]
|
||||
for nn, ff, yy, xx, zz, rc, ry, rx, rz in T.grid(14308, 1280, 1, 1, 1, 3, 2, 14, 14):
|
||||
with T.sblock("C"):
|
||||
v_nn, v_ff, v_yy, v_xx, v_zz, v_rc, v_ry, v_rx, v_rz = T.axis.remap("SSSSSRRRR", [nn, ff, yy, xx, zz, rc, ry, rx, rz])
|
||||
with T.init():
|
||||
C[v_nn, v_ff, v_yy, v_xx, v_zz] = T.float16(0.0)
|
||||
C[v_nn, v_ff, v_yy, v_xx, v_zz] += pad_A[v_nn, v_rc, v_yy * 2 + v_ry, v_xx * 14 + v_rx, v_zz * 14 + v_rz]* W[v_ff, v_rc, v_ry, v_rx, v_rz]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.Buffer((14308, 3, 2, 14, 14), "float16"), W: T.Buffer((1280, 3, 2, 14, 14), "float16"), C: T.Buffer((14308, 1280, 1, 1, 1), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
C_reindex_pad_local = T.sblock_alloc_buffer((1, 14336, 1280), "float16", scope="local")
|
||||
pad_A_reindex_pad_shared = T.sblock_alloc_buffer((1, 14336, 1184), "float16", scope="shared")
|
||||
W_reindex_pad_shared = T.sblock_alloc_buffer((1, 1280, 1184), "float16", scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(20, thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding(448, thread="blockIdx.x"):
|
||||
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
|
||||
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
|
||||
for ax2_2 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_2 in T.thread_binding(8, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
|
||||
for ax2_3_1_init in T.vectorized(2):
|
||||
with T.sblock("C_init"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(14336, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3_init)
|
||||
v2 = T.axis.spatial(1280, ax0_ax2_0_fused * 64 + ax2_1 * 64 + ax2_2 * 4 + ax2_3_0_init * 2 + ax2_3_1_init)
|
||||
C_reindex_pad_local[0, v1, v2] = T.float16(0.0)
|
||||
for ax3_0 in range(74):
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(2):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("pad_A_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(14336, ax1_0 * 32 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(1184, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 16)
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
pad_A_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < 14308 and v2 < 1176, A[v1, v2 // 392, v2 // 196 % 2, v2 // 14 % 14, v2 % 14], T.float16(0.0))
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(4):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("W_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(1280, ax0_ax2_0_fused * 64 + (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(1184, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 16)
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
W_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v2 < 1176, W[v1, v2 // 392, v2 // 196 % 2, v2 // 14 % 14, v2 % 14], T.float16(0.0))
|
||||
for ax3_1, ax1_3, ax2_3_0 in T.grid(16, 4, 2):
|
||||
for ax2_3_1 in T.vectorized(2):
|
||||
with T.sblock("C_update"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(14336, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3)
|
||||
v2 = T.axis.spatial(1280, ax0_ax2_0_fused * 64 + ax2_1 * 64 + ax2_2 * 4 + ax2_3_0 * 2 + ax2_3_1)
|
||||
v3 = T.axis.reduce(1184, ax3_0 * 16 + ax3_1)
|
||||
C_reindex_pad_local[0, v1, v2] = C_reindex_pad_local[0, v1, v2] + pad_A_reindex_pad_shared[0, v1, v3] * W_reindex_pad_shared[0, v2, v3]
|
||||
for ax0, ax1, ax2_0 in T.grid(1, 4, 2):
|
||||
for ax2_1_1 in T.vectorized(2):
|
||||
with T.sblock("C_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(14336, ax1_0 * 32 + ax1_2 * 4 + ax1)
|
||||
v2 = T.axis.spatial(1280, ax0_ax2_0_fused * 64 + ax2_2 * 4 + ax2_0 * 2 + ax2_1_1)
|
||||
T.where(ax1_0 * 32 + ax1_2 * 4 + ax1 < 14308)
|
||||
C[v1, v2, 0, 0, 0] = C_reindex_pad_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,262 @@
|
||||
# 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, E741, F841
|
||||
import tvm.testing
|
||||
from tvm.ir import assert_structural_equal
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_fallback():
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
B = T.sblock_alloc_buffer((1, 1, 32, 128), "float16")
|
||||
for i, j, k, l in T.grid(1, 1, 32, 128):
|
||||
with T.sblock("T_transpose"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
B[vi, vj, vk, vl] = A[vi, vk, vj, vl]
|
||||
for i, j, k in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_reshape"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
C[vi, vj, vk] = B[0, 0, vk % 4096 // 128, vk % 128]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
for ax0_fused_0 in T.thread_binding(4, thread="blockIdx.x"):
|
||||
for ax0_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("T_reshape"):
|
||||
v0 = T.axis.spatial(4096, ax0_fused_0 * 1024 + ax0_fused_1)
|
||||
T.reads(A[0, v0 // 128, 0, v0 % 128])
|
||||
T.writes(C[0, 0, v0])
|
||||
C[0, 0, v0] = A[0, v0 // 128, 0, v0 % 128]
|
||||
|
||||
target = Target("nvidia/geforce-rtx-3090-ti")
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Fallback(),
|
||||
)(Before)
|
||||
assert_structural_equal(mod, After)
|
||||
|
||||
|
||||
def test_fallback_reduction():
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 6144), "float32"), B: T.Buffer((1,), "float32")):
|
||||
for ax0, ax1 in T.grid(1, 6144):
|
||||
with T.sblock("block"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.reduce(6144, ax1)
|
||||
T.reads(A[v0, v1])
|
||||
T.writes(B[v0])
|
||||
with T.init():
|
||||
B[v0] = T.float32(0)
|
||||
B[v0] = B[v0] + T.Cast("float32", A[v0, v1])
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Expected:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 6144), "float32"), B: T.Buffer((1,), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
for ax0_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
||||
for ax0_fused_1 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
|
||||
with T.sblock("block_init"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
T.where(ax0_fused_0 * T.int64(1024) + ax0_fused_1 < T.int64(1))
|
||||
T.reads()
|
||||
T.writes(B[0])
|
||||
B[0] = T.float32(0)
|
||||
for ax1 in range(6144):
|
||||
with T.sblock("block_update"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.reduce(6144, ax1)
|
||||
T.where(ax0_fused_0 * T.int64(1024) + ax0_fused_1 < T.int64(1))
|
||||
T.reads(B[0], A[0, v1])
|
||||
T.writes(B[0])
|
||||
B[0] = B[0] + T.Cast("float32", A[0, v1])
|
||||
|
||||
with Target("apple/m1-gpu"):
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Fallback(),
|
||||
)(Module)
|
||||
assert_structural_equal(mod, Expected)
|
||||
|
||||
|
||||
def test_fallback_irregular_spatial():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(
|
||||
var_pages: T.handle,
|
||||
var_page_table_indptr: T.handle,
|
||||
var_page_table_values: T.handle,
|
||||
var_values: T.handle,
|
||||
seq_id: T.int32,
|
||||
):
|
||||
nhead = T.int32()
|
||||
nlayer = T.int32()
|
||||
seqlen = T.int32()
|
||||
npage = T.int32()
|
||||
page_size = T.int32()
|
||||
num_total_pages = T.int32()
|
||||
num_total_seqs_plus_1 = T.int32()
|
||||
|
||||
pages = T.match_buffer(var_pages, (num_total_pages, nlayer, nhead, page_size), "float16")
|
||||
page_table_indptr = T.match_buffer(var_page_table_indptr, (num_total_seqs_plus_1,), "int32")
|
||||
page_table_values = T.match_buffer(var_page_table_values, (npage,), "int32")
|
||||
values = T.match_buffer(var_values, (nlayer, nhead, seqlen), "float16")
|
||||
|
||||
for l, h, pos in T.grid(nlayer, nhead, seqlen):
|
||||
with T.sblock("block"):
|
||||
vl, vh, vp = T.axis.remap("SSS", [l, h, pos])
|
||||
values[vl, vh, vp] = pages[
|
||||
page_table_values[page_table_indptr[seq_id] + T.floordiv(vp, page_size)],
|
||||
vl,
|
||||
vh,
|
||||
T.floormod(vp, page_size),
|
||||
]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_pages: T.handle, var_page_table_indptr: T.handle, var_page_table_values: T.handle, var_values: T.handle, seq_id: T.int32):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
nhead = T.int32()
|
||||
nlayer = T.int32()
|
||||
seqlen = T.int32()
|
||||
npage = T.int32()
|
||||
page_size = T.int32()
|
||||
num_total_pages = T.int32()
|
||||
num_total_seqs_plus_1 = T.int32()
|
||||
|
||||
pages = T.match_buffer(var_pages, (num_total_pages, nlayer, nhead, page_size), "float16")
|
||||
page_table_indptr = T.match_buffer(var_page_table_indptr, (num_total_seqs_plus_1,), "int32")
|
||||
page_table_values = T.match_buffer(var_page_table_values, (npage,), "int32")
|
||||
values = T.match_buffer(var_values, (nlayer, nhead, seqlen), "float16")
|
||||
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding((nlayer * nhead * seqlen + 1023) // 1024, thread="blockIdx.x"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("block"):
|
||||
v0 = T.axis.spatial(nlayer, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) // (nhead * seqlen))
|
||||
v1 = T.axis.spatial(nhead, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % (nhead * seqlen) // seqlen)
|
||||
v2 = T.axis.spatial(seqlen, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % seqlen)
|
||||
T.where(ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1 < nlayer * nhead * seqlen)
|
||||
T.reads(pages[page_table_values[page_table_indptr[seq_id] + v2 // page_size], v0, v1, v2 % page_size], page_table_values[page_table_indptr[seq_id] + v2 // page_size], page_table_indptr[seq_id])
|
||||
T.writes(values[v0, v1, v2])
|
||||
values[v0, v1, v2] = pages[page_table_values[page_table_indptr[seq_id] + v2 // page_size], v0, v1, v2 % page_size]
|
||||
# fmt: on
|
||||
|
||||
target = Target("nvidia/geforce-rtx-3090-ti")
|
||||
with target:
|
||||
mod = tvm.IRModule({"main": func})
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Fallback(),
|
||||
)(mod)
|
||||
assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_gpu_fallback_ignores_non_gpu_functions():
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
# This function has no "target" attribute, and is scheduled
|
||||
# using the `Target.current`.
|
||||
@T.prim_func(s_tir=True)
|
||||
def gpu_func(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
B = T.sblock_alloc_buffer((1, 1, 32, 128), "float16")
|
||||
for i, j, k, l in T.grid(1, 1, 32, 128):
|
||||
with T.sblock("T_transpose"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
B[vi, vj, vk, vl] = A[vi, vk, vj, vl]
|
||||
for i, j, k in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_reshape"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
C[vi, vj, vk] = B[0, 0, vk % 4096 // 128, vk % 128]
|
||||
|
||||
# This function is identical, except that it is explicitly
|
||||
# annotated with the "target" attribute, and is scheduled
|
||||
# based on the annotation's target.
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_func(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
T.func_attr({"target": T.target("llvm")})
|
||||
B = T.sblock_alloc_buffer((1, 1, 32, 128), "float16")
|
||||
for i, j, k, l in T.grid(1, 1, 32, 128):
|
||||
with T.sblock("T_transpose"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
B[vi, vj, vk, vl] = A[vi, vk, vj, vl]
|
||||
for i, j, k in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_reshape"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
C[vi, vj, vk] = B[0, 0, vk % 4096 // 128, vk % 128]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def gpu_func(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
for ax0_fused_0 in T.thread_binding(4, thread="blockIdx.x"):
|
||||
for ax0_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("T_reshape"):
|
||||
v0 = T.axis.spatial(4096, ax0_fused_0 * 1024 + ax0_fused_1)
|
||||
T.reads(A[0, v0 // 128, 0, v0 % 128])
|
||||
T.writes(C[0, 0, v0])
|
||||
C[0, 0, v0] = A[0, v0 // 128, 0, v0 % 128]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_func(
|
||||
A: T.Buffer((1, 32, 1, 128), "float16"),
|
||||
C: T.Buffer((1, 1, 4096), "float16"),
|
||||
):
|
||||
T.func_attr({"target": T.target("llvm")})
|
||||
B = T.sblock_alloc_buffer((1, 1, 32, 128), "float16")
|
||||
for i, j, k, l in T.grid(1, 1, 32, 128):
|
||||
with T.sblock("T_transpose"):
|
||||
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
||||
B[vi, vj, vk, vl] = A[vi, vk, vj, vl]
|
||||
for i, j, k in T.grid(1, 1, 4096):
|
||||
with T.sblock("T_reshape"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
C[vi, vj, vk] = B[0, 0, vk % 4096 // 128, vk % 128]
|
||||
|
||||
with Target("cuda"):
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Fallback(),
|
||||
)(Before)
|
||||
assert_structural_equal(mod, After)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,695 @@
|
||||
# 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 tvm
|
||||
import tvm.testing
|
||||
from tvm.ir import IRModule, assert_structural_equal
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _check(mod_before: IRModule, mod_after: IRModule):
|
||||
target = Target("nvidia/geforce-rtx-3090-ti")
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.GeneralReduction(),
|
||||
)(mod_before)
|
||||
assert_structural_equal(mod, mod_after)
|
||||
|
||||
|
||||
def test_softmax_1():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p_lv44: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n, m = T.int64(), T.int64()
|
||||
lv44 = T.match_buffer(p_lv44, (T.int64(1), T.int64(32), n, m))
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (T.int64(1), T.int64(32), n, m), "float16")
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n))
|
||||
T_softmax_exp = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n, m))
|
||||
T_softmax_expsum = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n))
|
||||
var_T_softmax_norm_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n, m))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(32), n, m):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv44[v_i0, v_i1, v_i2, v_k])
|
||||
T.writes(T_softmax_maxelem[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
T_softmax_maxelem[v_i0, v_i1, v_i2] = T.float32(-3.4028234663852886e+38)
|
||||
T_softmax_maxelem[v_i0, v_i1, v_i2] = T.max(T_softmax_maxelem[v_i0, v_i1, v_i2], lv44[v_i0, v_i1, v_i2, v_k])
|
||||
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(32), n, m):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(lv44[v_i0, v_i1, v_i2, v_i3], T_softmax_maxelem[v_i0, v_i1, v_i2])
|
||||
T.writes(T_softmax_exp[v_i0, v_i1, v_i2, v_i3])
|
||||
T_softmax_exp[v_i0, v_i1, v_i2, v_i3] = T.exp(lv44[v_i0, v_i1, v_i2, v_i3] - T_softmax_maxelem[v_i0, v_i1, v_i2])
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(32), n, m):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_i2, v_k])
|
||||
T.writes(T_softmax_expsum[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
T_softmax_expsum[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
T_softmax_expsum[v_i0, v_i1, v_i2] = T_softmax_expsum[v_i0, v_i1, v_i2] + T_softmax_exp[v_i0, v_i1, v_i2, v_k]
|
||||
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(32), n, m):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_i2, v_i3], T_softmax_expsum[v_i0, v_i1, v_i2])
|
||||
T.writes(var_T_softmax_norm_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
T.sblock_attr({"axis": 3})
|
||||
var_T_softmax_norm_intermediate[v_i0, v_i1, v_i2, v_i3] = T_softmax_exp[v_i0, v_i1, v_i2, v_i3] / T_softmax_expsum[v_i0, v_i1, v_i2]
|
||||
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(32), n, m):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(var_T_softmax_norm_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
T.writes(var_compute_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
var_compute_intermediate[v_i0, v_i1, v_i2, v_i3] = T.Cast("float16", var_T_softmax_norm_intermediate[v_i0, v_i1, v_i2, v_i3])
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p_lv44: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n, m = T.int64(), T.int64()
|
||||
lv44 = T.match_buffer(p_lv44, (T.int64(1), T.int64(32), n, m))
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (T.int64(1), T.int64(32), n, m), "float16")
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n), scope="shared")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(32), n), scope="shared")
|
||||
for ax0_ax1_fused in T.thread_binding(n * T.int64(32), thread="blockIdx.x"):
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial((m + T.int64(255)) // T.int64(256), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v0 = T.axis.spatial(T.int64(32), ax0_ax1_fused // n + ax0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused % n + ax1)
|
||||
v2 = T.axis.reduce(m, ax2_fused_0 * T.int64(256) + ax2_fused_1)
|
||||
T.where(ax2_fused_0 * T.int64(256) + ax2_fused_1 < m)
|
||||
T.reads(lv44[T.int64(0), v0, v1, v2])
|
||||
T.writes(T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[T.int64(0), v0, v1] = T.float32(-3.4028234663852886e+38)
|
||||
T_softmax_maxelem_shared[T.int64(0), v0, v1] = T.max(T_softmax_maxelem_shared[T.int64(0), v0, v1], lv44[T.int64(0), v0, v1, v2])
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial((m + T.int64(255)) // T.int64(256), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v0 = T.axis.spatial(T.int64(32), ax0_ax1_fused // n + ax0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused % n + ax1)
|
||||
v2 = T.axis.reduce(m, ax2_fused_0 * T.int64(256) + ax2_fused_1)
|
||||
T.where(ax2_fused_0 * T.int64(256) + ax2_fused_1 < m)
|
||||
T.reads(lv44[T.int64(0), v0, v1, v2], T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
T.writes(T_softmax_expsum_shared[T.int64(0), v0, v1])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[T.int64(0), v0, v1] = T.float32(0)
|
||||
T_softmax_expsum_shared[T.int64(0), v0, v1] = T_softmax_expsum_shared[T.int64(0), v0, v1] + T.exp(lv44[T.int64(0), v0, v1, v2] - T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
for ax2_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_0 in T.serial((m + T.int64(255)) // T.int64(256), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("compute"):
|
||||
v0 = T.axis.spatial(T.int64(32), ax0_ax1_fused // n)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused % n)
|
||||
v2 = T.axis.spatial(m, ax2_0 * T.int64(256) + ax2_1)
|
||||
T.where(ax2_0 * T.int64(256) + ax2_1 < m)
|
||||
T.reads(lv44[T.int64(0), v0, v1, v2], T_softmax_maxelem_shared[T.int64(0), v0, v1], T_softmax_expsum_shared[T.int64(0), v0, v1])
|
||||
T.writes(var_compute_intermediate[T.int64(0), v0, v1, v2])
|
||||
var_compute_intermediate[T.int64(0), v0, v1, v2] = T.Cast("float16", T.exp(lv44[T.int64(0), v0, v1, v2] - T_softmax_maxelem_shared[T.int64(0), v0, v1]) / T_softmax_expsum_shared[T.int64(0), v0, v1])
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_softmax_2():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((T.int64(1), T.int64(1), T.int64(32000)), "float32"), T_softmax_norm: T.Buffer((T.int64(1), T.int64(1), T.int64(32000)), "float32")):
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer((T.int64(1), T.int64(1)))
|
||||
T_softmax_exp = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(32000)))
|
||||
T_softmax_expsum = T.sblock_alloc_buffer((T.int64(1), T.int64(1)))
|
||||
for i0, i1, k in T.grid(T.int64(1), T.int64(1), T.int64(32000)):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
||||
T.reads(A[v_i0, v_i1, v_k])
|
||||
T.writes(T_softmax_maxelem[v_i0, v_i1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[v_i0, v_i1] = T.float32(-3.4028234663852886e+38)
|
||||
T_softmax_maxelem[v_i0, v_i1] = T.max(T_softmax_maxelem[v_i0, v_i1], A[v_i0, v_i1, v_k])
|
||||
for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(32000)):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(A[v_i0, v_i1, v_i2], T_softmax_maxelem[v_i0, v_i1])
|
||||
T.writes(T_softmax_exp[v_i0, v_i1, v_i2])
|
||||
T_softmax_exp[v_i0, v_i1, v_i2] = T.exp(A[v_i0, v_i1, v_i2] - T_softmax_maxelem[v_i0, v_i1])
|
||||
for i0, i1, k in T.grid(T.int64(1), T.int64(1), T.int64(32000)):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_k])
|
||||
T.writes(T_softmax_expsum[v_i0, v_i1])
|
||||
with T.init():
|
||||
T_softmax_expsum[v_i0, v_i1] = T.float32(0)
|
||||
T_softmax_expsum[v_i0, v_i1] = T_softmax_expsum[v_i0, v_i1] + T_softmax_exp[v_i0, v_i1, v_k]
|
||||
for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(32000)):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_i2], T_softmax_expsum[v_i0, v_i1])
|
||||
T.writes(T_softmax_norm[v_i0, v_i1, v_i2])
|
||||
T.sblock_attr({"axis": 2})
|
||||
T_softmax_norm[v_i0, v_i1, v_i2] = T_softmax_exp[v_i0, v_i1, v_i2] / T_softmax_expsum[v_i0, v_i1]
|
||||
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((T.int64(1), T.int64(1), T.int64(32000)), "float32"), T_softmax_norm: T.Buffer((T.int64(1), T.int64(1), T.int64(32000)), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), scope="shared")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), scope="shared")
|
||||
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
||||
for ax0 in T.serial(T.int64(0), T.int64(1)):
|
||||
for ax1_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_fused_0 in T.serial(T.int64(125), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.reduce(T.int64(32000), ax1_fused_0 * T.int64(256) + ax1_fused_1)
|
||||
T.reads(A[T.int64(0), T.int64(0), v1])
|
||||
T.writes(T_softmax_maxelem_shared[T.int64(0), T.int64(0)])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[T.int64(0), T.int64(0)] = T.float32(-3.4028234663852886e+38)
|
||||
T_softmax_maxelem_shared[T.int64(0), T.int64(0)] = T.max(T_softmax_maxelem_shared[T.int64(0), T.int64(0)], A[T.int64(0), T.int64(0), v1])
|
||||
for ax0 in T.serial(T.int64(0), T.int64(1)):
|
||||
for ax1_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_fused_0 in T.serial(T.int64(125), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.reduce(T.int64(32000), ax1_fused_0 * T.int64(256) + ax1_fused_1)
|
||||
T.reads(A[T.int64(0), T.int64(0), v1], T_softmax_maxelem_shared[T.int64(0), T.int64(0)])
|
||||
T.writes(T_softmax_expsum_shared[T.int64(0), T.int64(0)])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[T.int64(0), T.int64(0)] = T.float32(0)
|
||||
T_softmax_expsum_shared[T.int64(0), T.int64(0)] = T_softmax_expsum_shared[T.int64(0), T.int64(0)] + T.exp(A[T.int64(0), T.int64(0), v1] - T_softmax_maxelem_shared[T.int64(0), T.int64(0)])
|
||||
for ax1_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_0 in T.serial(T.int64(125), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(T.int64(32000), ax1_0 * T.int64(256) + ax1_1)
|
||||
T.reads(A[T.int64(0), T.int64(0), v1], T_softmax_maxelem_shared[T.int64(0), T.int64(0)], T_softmax_expsum_shared[T.int64(0), T.int64(0)])
|
||||
T.writes(T_softmax_norm[T.int64(0), T.int64(0), v1])
|
||||
T.sblock_attr({"axis": 2})
|
||||
T_softmax_norm[T.int64(0), T.int64(0), v1] = T.exp(A[T.int64(0), T.int64(0), v1] - T_softmax_maxelem_shared[T.int64(0), T.int64(0)]) / T_softmax_expsum_shared[T.int64(0), T.int64(0)]
|
||||
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_softmax_3():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(input: T.Buffer((T.int64(1), T.int64(4), T.int64(32), T.int64(8192)), "float32"), T_softmax_norm: T.Buffer((T.int64(1), T.int64(4), T.int64(32), T.int64(8192)), "float32")):
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer((T.int64(1), T.int64(4), T.int64(8192)))
|
||||
T_softmax_exp = T.sblock_alloc_buffer((T.int64(1), T.int64(4), T.int64(32), T.int64(8192)))
|
||||
T_softmax_expsum = T.sblock_alloc_buffer((T.int64(1), T.int64(4), T.int64(8192)))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(4), T.int64(8192), T.int64(32)):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(input[v_i0, v_i1, v_k, v_i2])
|
||||
T.writes(T_softmax_maxelem[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
T_softmax_maxelem[v_i0, v_i1, v_i2] = T.float32(-340282346638528859811704183484516925440.0)
|
||||
T_softmax_maxelem[v_i0, v_i1, v_i2] = T.max(T_softmax_maxelem[v_i0, v_i1, v_i2], input[v_i0, v_i1, v_k, v_i2])
|
||||
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(4), T.int64(32), T.int64(8192)):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(input[v_i0, v_i1, v_i2, v_i3], T_softmax_maxelem[v_i0, v_i1, v_i3])
|
||||
T.writes(T_softmax_exp[v_i0, v_i1, v_i2, v_i3])
|
||||
T_softmax_exp[v_i0, v_i1, v_i2, v_i3] = T.exp(input[v_i0, v_i1, v_i2, v_i3] - T_softmax_maxelem[v_i0, v_i1, v_i3])
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(4), T.int64(8192), T.int64(32)):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_k, v_i2])
|
||||
T.writes(T_softmax_expsum[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
T_softmax_expsum[v_i0, v_i1, v_i2] = T.float32(0.0)
|
||||
T_softmax_expsum[v_i0, v_i1, v_i2] = T_softmax_expsum[v_i0, v_i1, v_i2] + T_softmax_exp[v_i0, v_i1, v_k, v_i2]
|
||||
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(4), T.int64(32), T.int64(8192)):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(T_softmax_exp[v_i0, v_i1, v_i2, v_i3], T_softmax_expsum[v_i0, v_i1, v_i3])
|
||||
T.writes(T_softmax_norm[v_i0, v_i1, v_i2, v_i3])
|
||||
T.sblock_attr({"axis": 2})
|
||||
T_softmax_norm[v_i0, v_i1, v_i2, v_i3] = T_softmax_exp[v_i0, v_i1, v_i2, v_i3] / T_softmax_expsum[v_i0, v_i1, v_i3]
|
||||
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(input: T.Buffer((T.int64(1), T.int64(4), T.int64(32), T.int64(8192)), "float32"), T_softmax_norm: T.Buffer((T.int64(1), T.int64(4), T.int64(32), T.int64(8192)), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(4), T.int64(8192)), scope="shared")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(4), T.int64(8192)), scope="shared")
|
||||
for ax0_ax2_fused in T.thread_binding(T.int64(32768), thread="blockIdx.x"):
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
v0 = T.axis.spatial(T.int64(4), ax0_ax2_fused // T.int64(8192) + ax0)
|
||||
v1 = T.axis.spatial(T.int64(8192), ax0_ax2_fused % T.int64(8192) + ax1)
|
||||
v2 = T.axis.reduce(T.int64(32), ax2_fused_0 * T.int64(256) + ax2_fused_1)
|
||||
T.where(ax2_fused_0 * T.int64(256) + ax2_fused_1 < T.int64(32))
|
||||
T.reads(input[T.int64(0), v0, v2, v1])
|
||||
T.writes(T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[T.int64(0), v0, v1] = T.float32(-340282346638528859811704183484516925440.0)
|
||||
T_softmax_maxelem_shared[T.int64(0), v0, v1] = T.max(T_softmax_maxelem_shared[T.int64(0), v0, v1], input[T.int64(0), v0, v2, v1])
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
v0 = T.axis.spatial(T.int64(4), ax0_ax2_fused // T.int64(8192) + ax0)
|
||||
v1 = T.axis.spatial(T.int64(8192), ax0_ax2_fused % T.int64(8192) + ax1)
|
||||
v2 = T.axis.reduce(T.int64(32), ax2_fused_0 * T.int64(256) + ax2_fused_1)
|
||||
T.where(ax2_fused_0 * T.int64(256) + ax2_fused_1 < T.int64(32))
|
||||
T.reads(input[T.int64(0), v0, v2, v1], T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
T.writes(T_softmax_expsum_shared[T.int64(0), v0, v1])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[T.int64(0), v0, v1] = T.float32(0.0)
|
||||
T_softmax_expsum_shared[T.int64(0), v0, v1] = T_softmax_expsum_shared[T.int64(0), v0, v1] + T.exp(input[T.int64(0), v0, v2, v1] - T_softmax_maxelem_shared[T.int64(0), v0, v1])
|
||||
for ax1_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
v0 = T.axis.spatial(T.int64(4), ax0_ax2_fused // T.int64(8192))
|
||||
v1 = T.axis.spatial(T.int64(32), ax1_0 * T.int64(256) + ax1_1)
|
||||
v2 = T.axis.spatial(T.int64(8192), ax0_ax2_fused % T.int64(8192))
|
||||
T.where(ax1_0 * T.int64(256) + ax1_1 < T.int64(32))
|
||||
T.reads(input[T.int64(0), v0, v1, v2], T_softmax_maxelem_shared[T.int64(0), v0, v2], T_softmax_expsum_shared[T.int64(0), v0, v2])
|
||||
T.writes(T_softmax_norm[T.int64(0), v0, v1, v2])
|
||||
T.sblock_attr({"axis": 2})
|
||||
T_softmax_norm[T.int64(0), v0, v1, v2] = T.exp(input[T.int64(0), v0, v1, v2] - T_softmax_maxelem_shared[T.int64(0), v0, v2]) / T_softmax_expsum_shared[T.int64(0), v0, v2]
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_layer_norm():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p_lv6: T.handle, weight1: T.Buffer((T.int64(2560),), "float32"), bias: T.Buffer((T.int64(2560),), "float32"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv6 = T.match_buffer(p_lv6, (T.int64(1), n, T.int64(2560)))
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (T.int64(1), n, T.int64(2560)), "float16")
|
||||
# with T.sblock("root"):
|
||||
A_red_temp_v0 = T.sblock_alloc_buffer((T.int64(1), n))
|
||||
A_red_temp_v1 = T.sblock_alloc_buffer((T.int64(1), n))
|
||||
var_T_layer_norm_intermediate = T.sblock_alloc_buffer((T.int64(1), n, T.int64(2560)))
|
||||
for ax0, ax1, k2 in T.grid(T.int64(1), n, T.int64(2560)):
|
||||
with T.sblock("A_red_temp"):
|
||||
v_ax0, v_ax1, v_k2 = T.axis.remap("SSR", [ax0, ax1, k2])
|
||||
T.reads(lv6[v_ax0, v_ax1, v_k2])
|
||||
T.writes(A_red_temp_v0[v_ax0, v_ax1], A_red_temp_v1[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
A_red_temp_v0[v_ax0, v_ax1] = T.float32(0)
|
||||
A_red_temp_v1[v_ax0, v_ax1] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0[v_ax0, v_ax1] + lv6[v_ax0, v_ax1, v_k2]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1[v_ax0, v_ax1] + lv6[v_ax0, v_ax1, v_k2] * lv6[v_ax0, v_ax1, v_k2]
|
||||
A_red_temp_v0[v_ax0, v_ax1] = v_A_red_temp_v0
|
||||
A_red_temp_v1[v_ax0, v_ax1] = v_A_red_temp_v1
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), n, T.int64(2560)):
|
||||
with T.sblock("T_layer_norm"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(lv6[v_ax0, v_ax1, v_ax2], A_red_temp_v0[v_ax0, v_ax1], A_red_temp_v1[v_ax0, v_ax1], weight1[v_ax2], bias[v_ax2])
|
||||
T.writes(var_T_layer_norm_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
var_T_layer_norm_intermediate[v_ax0, v_ax1, v_ax2] = (lv6[v_ax0, v_ax1, v_ax2] - A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.00039062500000000002)) * T.rsqrt(A_red_temp_v1[v_ax0, v_ax1] * T.float32(0.00039062500000000002) - A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.00039062500000000002) * (A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.00039062500000000002)) + T.float32(1.0000000000000001e-05)) * weight1[v_ax2] + bias[v_ax2]
|
||||
for i0, i1, i2 in T.grid(T.int64(1), n, T.int64(2560)):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(var_T_layer_norm_intermediate[v_i0, v_i1, v_i2])
|
||||
T.writes(var_compute_intermediate[v_i0, v_i1, v_i2])
|
||||
var_compute_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_layer_norm_intermediate[v_i0, v_i1, v_i2])
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p_lv6: T.handle, weight1: T.Buffer((T.int64(2560),), "float32"), bias: T.Buffer((T.int64(2560),), "float32"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv6 = T.match_buffer(p_lv6, (T.int64(1), n, T.int64(2560)))
|
||||
var_compute_intermediate = T.match_buffer(p_output0, (T.int64(1), n, T.int64(2560)), "float16")
|
||||
# with T.sblock("root"):
|
||||
A_red_temp_v0_shared = T.sblock_alloc_buffer((T.int64(1), n), scope="shared")
|
||||
A_red_temp_v1_shared = T.sblock_alloc_buffer((T.int64(1), n), scope="shared")
|
||||
for ax0_fused in T.thread_binding(n, thread="blockIdx.x"):
|
||||
for ax0 in T.serial(T.int64(0), T.int64(1)):
|
||||
for ax1_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_fused_0 in T.serial(T.int64(10), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("A_red_temp"):
|
||||
v0 = T.axis.spatial(n, ax0_fused + ax0)
|
||||
v1 = T.axis.reduce(T.int64(2560), ax1_fused_0 * T.int64(256) + ax1_fused_1)
|
||||
T.reads(lv6[T.int64(0), v0, v1])
|
||||
T.writes(A_red_temp_v0_shared[T.int64(0), v0], A_red_temp_v1_shared[T.int64(0), v0])
|
||||
with T.init():
|
||||
A_red_temp_v0_shared[T.int64(0), v0] = T.float32(0)
|
||||
A_red_temp_v1_shared[T.int64(0), v0] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0_shared[T.int64(0), v0] + lv6[T.int64(0), v0, v1]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1_shared[T.int64(0), v0] + lv6[T.int64(0), v0, v1] * lv6[T.int64(0), v0, v1]
|
||||
A_red_temp_v0_shared[T.int64(0), v0] = v_A_red_temp_v0
|
||||
A_red_temp_v1_shared[T.int64(0), v0] = v_A_red_temp_v1
|
||||
for ax1_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_0 in T.serial(T.int64(10), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("compute"):
|
||||
v0 = T.axis.spatial(n, ax0_fused)
|
||||
v1 = T.axis.spatial(T.int64(2560), ax1_0 * T.int64(256) + ax1_1)
|
||||
T.reads(lv6[T.int64(0), v0, v1], A_red_temp_v0_shared[T.int64(0), v0], A_red_temp_v1_shared[T.int64(0), v0], weight1[v1], bias[v1])
|
||||
T.writes(var_compute_intermediate[T.int64(0), v0, v1])
|
||||
var_compute_intermediate[T.int64(0), v0, v1] = T.Cast("float16", (lv6[T.int64(0), v0, v1] - A_red_temp_v0_shared[T.int64(0), v0] * T.float32(0.00039062500000000002)) * T.rsqrt(A_red_temp_v1_shared[T.int64(0), v0] * T.float32(0.00039062500000000002) - A_red_temp_v0_shared[T.int64(0), v0] * T.float32(0.00039062500000000002) * (A_red_temp_v0_shared[T.int64(0), v0] * T.float32(0.00039062500000000002)) + T.float32(1.0000000000000001e-05)) * weight1[v1] + bias[v1])
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_rms_norm():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, B: T.Buffer((T.int64(4096),), "float16"), var_rms_norm: T.handle):
|
||||
T.func_attr({"op_pattern": 4, "tirx.noalias": True})
|
||||
n = T.int64()
|
||||
A = T.match_buffer(var_A, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
rms_norm_1 = T.match_buffer(var_rms_norm, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
Ared_temp = T.sblock_alloc_buffer((T.int64(1), n))
|
||||
for bsz, i, k in T.grid(T.int64(1), n, T.int64(4096)):
|
||||
with T.sblock("Ared_temp"):
|
||||
v_bsz, v_i, v_k = T.axis.remap("SSR", [bsz, i, k])
|
||||
T.reads(A[v_bsz, v_i, v_k])
|
||||
T.writes(Ared_temp[v_bsz, v_i])
|
||||
with T.init():
|
||||
Ared_temp[v_bsz, v_i] = T.float32(0)
|
||||
Ared_temp[v_bsz, v_i] = Ared_temp[v_bsz, v_i] + T.Cast("float32", A[v_bsz, v_i, v_k]) * T.Cast("float32", A[v_bsz, v_i, v_k])
|
||||
for bsz, i, k in T.grid(T.int64(1), n, T.int64(4096)):
|
||||
with T.sblock("rms_norm"):
|
||||
v_bsz, v_i, v_k = T.axis.remap("SSS", [bsz, i, k])
|
||||
T.reads(B[v_k], A[v_bsz, v_i, v_k], Ared_temp[v_bsz, v_i])
|
||||
T.writes(rms_norm_1[v_bsz, v_i, v_k])
|
||||
rms_norm_1[v_bsz, v_i, v_k] = T.Cast("float16", T.Cast("float32", B[v_k]) * (T.Cast("float32", A[v_bsz, v_i, v_k]) / T.sqrt(Ared_temp[v_bsz, v_i] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))))
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, B: T.Buffer((T.int64(4096),), "float16"), var_rms_norm: T.handle):
|
||||
T.func_attr({"op_pattern": 4, "tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int64()
|
||||
A = T.match_buffer(var_A, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
rms_norm_1 = T.match_buffer(var_rms_norm, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
Ared_temp_shared = T.sblock_alloc_buffer((T.int64(1), n), scope="shared")
|
||||
for ax0_fused in T.thread_binding(n, thread="blockIdx.x"):
|
||||
for ax0 in T.serial(T.int64(0), T.int64(1)):
|
||||
for ax1_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_fused_0 in T.serial(T.int64(16), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("Ared_temp"):
|
||||
v0 = T.axis.spatial(n, ax0_fused + ax0)
|
||||
v1 = T.axis.reduce(T.int64(4096), ax1_fused_0 * T.int64(256) + ax1_fused_1)
|
||||
T.reads(A[T.int64(0), v0, v1])
|
||||
T.writes(Ared_temp_shared[T.int64(0), v0])
|
||||
with T.init():
|
||||
Ared_temp_shared[T.int64(0), v0] = T.float32(0)
|
||||
Ared_temp_shared[T.int64(0), v0] = Ared_temp_shared[T.int64(0), v0] + T.Cast("float32", A[T.int64(0), v0, v1]) * T.Cast("float32", A[T.int64(0), v0, v1])
|
||||
for ax1_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax1_0 in T.serial(T.int64(16), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("rms_norm"):
|
||||
v0 = T.axis.spatial(n, ax0_fused)
|
||||
v1 = T.axis.spatial(T.int64(4096), ax1_0 * T.int64(256) + ax1_1)
|
||||
T.reads(B[v1], A[T.int64(0), v0, v1], Ared_temp_shared[T.int64(0), v0])
|
||||
T.writes(rms_norm_1[T.int64(0), v0, v1])
|
||||
rms_norm_1[T.int64(0), v0, v1] = T.Cast("float16", T.Cast("float32", B[v1]) * (T.Cast("float32", A[T.int64(0), v0, v1]) / T.sqrt(Ared_temp_shared[T.int64(0), v0] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))))
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_group_norm():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 2048), "float32"), B: T.Buffer((2048,), "float32"), C: T.Buffer((2048,), "float32"), T_reshape: T.Buffer((1, 2048), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
T_reshape_1 = T.sblock_alloc_buffer((1, 32, 64))
|
||||
A_red_temp_v0 = T.sblock_alloc_buffer((1, 32))
|
||||
A_red_temp_v1 = T.sblock_alloc_buffer((1, 32))
|
||||
T_reshape_2 = T.sblock_alloc_buffer((32, 64))
|
||||
T_reshape_3 = T.sblock_alloc_buffer((32, 64))
|
||||
T_group_norm = T.sblock_alloc_buffer((1, 32, 64))
|
||||
for ax0, ax1, ax2 in T.grid(1, 32, 64):
|
||||
with T.sblock("T_reshape"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(A[0, (v_ax1 * 64 + v_ax2) % 2048])
|
||||
T.writes(T_reshape_1[v_ax0, v_ax1, v_ax2])
|
||||
T_reshape_1[v_ax0, v_ax1, v_ax2] = A[0, (v_ax1 * 64 + v_ax2) % 2048]
|
||||
for ax0, ax1, k2 in T.grid(1, 32, 64):
|
||||
with T.sblock("A_red_temp"):
|
||||
v_ax0, v_ax1, v_k2 = T.axis.remap("SSR", [ax0, ax1, k2])
|
||||
T.reads(T_reshape_1[v_ax0, v_ax1, v_k2])
|
||||
T.writes(A_red_temp_v0[v_ax0, v_ax1], A_red_temp_v1[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
A_red_temp_v0[v_ax0, v_ax1] = T.float32(0)
|
||||
A_red_temp_v1[v_ax0, v_ax1] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0[v_ax0, v_ax1] + T_reshape_1[v_ax0, v_ax1, v_k2]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1[v_ax0, v_ax1] + T_reshape_1[v_ax0, v_ax1, v_k2] * T_reshape_1[v_ax0, v_ax1, v_k2]
|
||||
A_red_temp_v0[v_ax0, v_ax1] = v_A_red_temp_v0
|
||||
A_red_temp_v1[v_ax0, v_ax1] = v_A_red_temp_v1
|
||||
for ax0, ax1 in T.grid(32, 64):
|
||||
with T.sblock("T_reshape_1"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(B[(v_ax0 * 64 + v_ax1) % 2048])
|
||||
T.writes(T_reshape_2[v_ax0, v_ax1])
|
||||
T_reshape_2[v_ax0, v_ax1] = B[(v_ax0 * 64 + v_ax1) % 2048]
|
||||
for ax0, ax1 in T.grid(32, 64):
|
||||
with T.sblock("T_reshape_2"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(C[(v_ax0 * 64 + v_ax1) % 2048])
|
||||
T.writes(T_reshape_3[v_ax0, v_ax1])
|
||||
T_reshape_3[v_ax0, v_ax1] = C[(v_ax0 * 64 + v_ax1) % 2048]
|
||||
for ax0, ax1, ax2 in T.grid(1, 32, 64):
|
||||
with T.sblock("T_group_norm"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(T_reshape_1[v_ax0, v_ax1, v_ax2], A_red_temp_v0[v_ax0, v_ax1], A_red_temp_v1[v_ax0, v_ax1], T_reshape_2[v_ax1, v_ax2], T_reshape_3[v_ax1, v_ax2])
|
||||
T.writes(T_group_norm[v_ax0, v_ax1, v_ax2])
|
||||
T_group_norm[v_ax0, v_ax1, v_ax2] = (T_reshape_1[v_ax0, v_ax1, v_ax2] - A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.015625)) * T.rsqrt(A_red_temp_v1[v_ax0, v_ax1] * T.float32(0.015625) - A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.015625) * (A_red_temp_v0[v_ax0, v_ax1] * T.float32(0.015625)) + T.float32(1.0000000000000001e-05)) * T_reshape_2[v_ax1, v_ax2] + T_reshape_3[v_ax1, v_ax2]
|
||||
for ax0, ax1 in T.grid(1, 2048):
|
||||
with T.sblock("T_reshape_3"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(T_group_norm[0, v_ax1 % 2048 // 64, v_ax1 % 64])
|
||||
T.writes(T_reshape[v_ax0, v_ax1])
|
||||
T_reshape[v_ax0, v_ax1] = T_group_norm[0, v_ax1 % 2048 // 64, v_ax1 % 64]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 2048), "float32"), B: T.Buffer((2048,), "float32"), C: T.Buffer((2048,), "float32"), T_reshape: T.Buffer((1, 2048), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
A_red_temp_v0_shared = T.sblock_alloc_buffer((1, 32), scope="shared")
|
||||
A_red_temp_v1_shared = T.sblock_alloc_buffer((1, 32), scope="shared")
|
||||
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
||||
for ax0 in range(32):
|
||||
for ax1_fused_1 in T.thread_binding(256, thread="threadIdx.x"):
|
||||
for ax1_fused_0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("A_red_temp"):
|
||||
v0 = T.axis.spatial(32, ax0)
|
||||
v1 = T.axis.reduce(64, ax1_fused_0 * 256 + ax1_fused_1)
|
||||
T.where(ax1_fused_0 * 256 + ax1_fused_1 < 64)
|
||||
T.reads(A[0, v0 * 64 + v1])
|
||||
T.writes(A_red_temp_v0_shared[0, v0], A_red_temp_v1_shared[0, v0])
|
||||
with T.init():
|
||||
A_red_temp_v0_shared[0, v0] = T.float32(0)
|
||||
A_red_temp_v1_shared[0, v0] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0_shared[0, v0] + A[0, v0 * 64 + v1]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1_shared[0, v0] + A[0, v0 * 64 + v1] * A[0, v0 * 64 + v1]
|
||||
A_red_temp_v0_shared[0, v0] = v_A_red_temp_v0
|
||||
A_red_temp_v1_shared[0, v0] = v_A_red_temp_v1
|
||||
for ax1_1 in T.thread_binding(256, thread="threadIdx.x"):
|
||||
for ax1_0 in T.serial(8, annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
with T.sblock("T_reshape_3"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(2048, ax1_0 * 256 + ax1_1)
|
||||
T.reads(A[0, v1], A_red_temp_v0_shared[0, v1 // 64], A_red_temp_v1_shared[0, v1 // 64], B[v1], C[v1])
|
||||
T.writes(T_reshape[0, v1])
|
||||
T_reshape[0, v1] = (A[0, v1] - A_red_temp_v0_shared[0, v1 // 64] * T.float32(0.015625)) * T.rsqrt(A_red_temp_v1_shared[0, v1 // 64] * T.float32(0.015625) - A_red_temp_v0_shared[0, v1 // 64] * T.float32(0.015625) * (A_red_temp_v0_shared[0, v1 // 64] * T.float32(0.015625)) + T.float32(1.0000000000000001e-05)) * B[v1] + C[v1]
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_logsumexp():
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def compute_lse(var_A: T.handle, var_blocked_lse: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
batch_size = T.int64()
|
||||
vocab_size = T.int64()
|
||||
num_chunks = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
|
||||
blocked_lse = T.match_buffer(var_blocked_lse, (batch_size, num_chunks), dtype="float32")
|
||||
A_pad = T.sblock_alloc_buffer((batch_size, num_chunks, T.int64(4096)), dtype="float32")
|
||||
temp_max = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
|
||||
temp_sum = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
|
||||
|
||||
for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(4096)):
|
||||
with T.sblock("pad"):
|
||||
v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
|
||||
A_pad[v0, v1, v2] = T.if_then_else(
|
||||
v1 * T.int64(4096) + v2 < vocab_size,
|
||||
A[v0, v1 * T.int64(4096) + v2],
|
||||
T.min_value("float32"),
|
||||
)
|
||||
|
||||
for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(4096)):
|
||||
with T.sblock("max"):
|
||||
v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
|
||||
with T.init():
|
||||
temp_max[v0, v1] = T.min_value("float32")
|
||||
temp_max[v0, v1] = T.max(temp_max[v0, v1], A_pad[v0, v1, v2])
|
||||
|
||||
for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(4096)):
|
||||
with T.sblock("sum_exp"):
|
||||
v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
|
||||
with T.init():
|
||||
temp_sum[v0, v1] = T.float32(0)
|
||||
temp_sum[v0, v1] += T.if_then_else(
|
||||
v1 * T.int64(4096) + v2 < vocab_size,
|
||||
T.exp(A_pad[v0, v1, v2] - temp_max[v0, v1]),
|
||||
T.float32(0),
|
||||
)
|
||||
|
||||
for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(1)):
|
||||
with T.sblock("log"):
|
||||
v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
|
||||
blocked_lse[v0, v1] = T.log(temp_sum[v0, v1]) + temp_max[v0, v1]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def compute_lse(var_A: T.handle, var_blocked_lse: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
batch_size, vocab_size = T.int64(), T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, vocab_size))
|
||||
num_chunks = T.int64()
|
||||
blocked_lse = T.match_buffer(var_blocked_lse, (batch_size, num_chunks))
|
||||
temp_max_shared = T.sblock_alloc_buffer((batch_size, num_chunks), scope="shared")
|
||||
temp_sum_shared = T.sblock_alloc_buffer((batch_size, num_chunks), scope="shared")
|
||||
for ax0_ax1_fused in T.thread_binding(batch_size * num_chunks, thread="blockIdx.x"):
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial(
|
||||
T.int64(16),
|
||||
annotations={
|
||||
"pragma_auto_unroll_max_step": 256,
|
||||
"pragma_unroll_explicit": 1,
|
||||
},
|
||||
):
|
||||
with T.sblock("max"):
|
||||
v0 = T.axis.spatial(
|
||||
batch_size,
|
||||
ax0_ax1_fused // num_chunks + ax0,
|
||||
)
|
||||
v1 = T.axis.spatial(num_chunks, ax0_ax1_fused % num_chunks + ax1)
|
||||
v2 = T.axis.reduce(
|
||||
T.int64(4096), ax2_fused_0 * T.int64(256) + ax2_fused_1
|
||||
)
|
||||
T.reads(A[v0, v1 * T.int64(4096) + v2])
|
||||
T.writes(temp_max_shared[v0, v1])
|
||||
with T.init():
|
||||
temp_max_shared[v0, v1] = T.min_value("float32")
|
||||
temp_max_shared[v0, v1] = T.max(
|
||||
temp_max_shared[v0, v1],
|
||||
T.if_then_else(
|
||||
v1 * T.int64(4096) + v2 < vocab_size,
|
||||
A[v0, v1 * T.int64(4096) + v2],
|
||||
T.min_value("float32"),
|
||||
),
|
||||
)
|
||||
for ax0, ax1 in T.grid(T.int64(1), T.int64(1)):
|
||||
for ax2_fused_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_fused_0 in T.serial(
|
||||
T.int64(16),
|
||||
annotations={
|
||||
"pragma_auto_unroll_max_step": 256,
|
||||
"pragma_unroll_explicit": 1,
|
||||
},
|
||||
):
|
||||
with T.sblock("sum_exp"):
|
||||
v0 = T.axis.spatial(
|
||||
batch_size,
|
||||
ax0_ax1_fused // num_chunks + ax0,
|
||||
)
|
||||
v1 = T.axis.spatial(num_chunks, ax0_ax1_fused % num_chunks + ax1)
|
||||
v2 = T.axis.reduce(
|
||||
T.int64(4096), ax2_fused_0 * T.int64(256) + ax2_fused_1
|
||||
)
|
||||
T.reads(A[v0, v1 * T.int64(4096) + v2], temp_max_shared[v0, v1])
|
||||
T.writes(temp_sum_shared[v0, v1])
|
||||
with T.init():
|
||||
temp_sum_shared[v0, v1] = T.float32(0)
|
||||
temp_sum_shared[v0, v1] = temp_sum_shared[v0, v1] + T.if_then_else(
|
||||
v1 * T.int64(4096) + v2 < vocab_size,
|
||||
T.exp(
|
||||
T.if_then_else(
|
||||
v1 * T.int64(4096) + v2 < vocab_size,
|
||||
A[v0, v1 * T.int64(4096) + v2],
|
||||
T.min_value("float32"),
|
||||
)
|
||||
- temp_max_shared[v0, v1]
|
||||
),
|
||||
T.float32(0),
|
||||
)
|
||||
for ax2_1 in T.thread_binding(T.int64(256), thread="threadIdx.x"):
|
||||
for ax2_0 in T.serial(
|
||||
T.int64(1),
|
||||
annotations={
|
||||
"pragma_auto_unroll_max_step": 256,
|
||||
"pragma_unroll_explicit": 1,
|
||||
},
|
||||
):
|
||||
with T.sblock("log"):
|
||||
v0 = T.axis.spatial(batch_size, ax0_ax1_fused // num_chunks)
|
||||
v1 = T.axis.spatial(num_chunks, ax0_ax1_fused % num_chunks)
|
||||
v2 = T.axis.spatial(T.int64(1), ax2_0 * T.int64(256) + ax2_1)
|
||||
T.where(ax2_0 * T.int64(256) + ax2_1 < T.int64(1))
|
||||
T.reads(temp_sum_shared[v0, v1], temp_max_shared[v0, v1])
|
||||
T.writes(blocked_lse[v0, v1])
|
||||
blocked_lse[v0, v1] = (
|
||||
T.log(temp_sum_shared[v0, v1]) + temp_max_shared[v0, v1]
|
||||
)
|
||||
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,561 @@
|
||||
# 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
|
||||
|
||||
|
||||
import tvm.testing
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_batch_decode_gemv():
|
||||
# fmt: off
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv429: T.Buffer((T.int64(4096), T.int64(3584)), "uint32"), lv430: T.Buffer((T.int64(4096), T.int64(896)), "float16"), p_lv807: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True, "tirx.HoistIfThenElseExprWithBlock": 1})
|
||||
batch_size = T.int64()
|
||||
lv807 = T.match_buffer(p_lv807, (batch_size, T.int64(1), T.int64(28672)), "float16")
|
||||
NT_matmul_intermediate = T.match_buffer(p_output0, (batch_size, T.int64(1), T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
compute = T.sblock_alloc_buffer((T.int64(4096), T.int64(28672)), "float16")
|
||||
dequantize_intermediate_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(28672)), "float16")
|
||||
for i0, i1 in T.grid(T.int64(4096), T.int64(28672)):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(lv429[v_i0, v_i1 // T.int64(8)])
|
||||
T.writes(compute[v_i0, v_i1])
|
||||
compute[v_i0, v_i1] = T.Cast("float16", T.bitwise_and(T.shift_right(lv429[v_i0, v_i1 // T.int64(8)], T.Cast("uint32", v_i1 % T.int64(8) * T.int64(4))), T.uint32(15)))
|
||||
for i0, i1 in T.grid(T.int64(4096), T.int64(28672)):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(compute[v_i0, v_i1], lv430[v_i0, v_i1 // T.int64(32)])
|
||||
T.writes(dequantize_intermediate_intermediate[v_i0, v_i1])
|
||||
dequantize_intermediate_intermediate[v_i0, v_i1] = (compute[v_i0, v_i1] - T.float16(7)) * lv430[v_i0, v_i1 // T.int64(32)]
|
||||
for i0, i1, i2, k in T.grid(batch_size, T.int64(1), T.int64(4096), T.int64(28672)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv807[v_i0, v_i1, v_k], dequantize_intermediate_intermediate[v_i2, v_k])
|
||||
T.writes(NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
NT_matmul_intermediate[v_i0, v_i1, v_i2] = NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv807[v_i0, v_i1, v_k] * dequantize_intermediate_intermediate[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv429: T.Buffer((T.int64(4096), T.int64(3584)), "uint32"), lv430: T.Buffer((T.int64(4096), T.int64(896)), "float16"), p_lv807: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.HoistIfThenElseExprWithBlock": 1, "tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
batch_size = T.int64()
|
||||
lv807 = T.match_buffer(p_lv807, (batch_size, T.int64(1), T.int64(28672)), "float16")
|
||||
NT_matmul_intermediate = T.match_buffer(p_output0, (batch_size, T.int64(1), T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
dequantize_intermediate_intermediate_local = T.sblock_alloc_buffer((T.int64(4096), T.int64(28672)), "float16", scope="local")
|
||||
NT_matmul_intermediate_pad_local = T.sblock_alloc_buffer(((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
NT_matmul_intermediate_pad_rf_local = T.sblock_alloc_buffer((T.int64(128), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
NT_matmul_intermediate_pad_rf_local_1 = T.sblock_alloc_buffer((T.int64(32), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
for ax0_0 in T.thread_binding((batch_size + T.int64(3)) // T.int64(4), thread="blockIdx.y"):
|
||||
for u_fused_ax1_fused_fused_0 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
||||
for u_fused_ax1_fused_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
for ax0_1_init, u_fused_ax1_fused_fused_2_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1_init)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2_init)
|
||||
T.reads()
|
||||
T.writes(NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1])
|
||||
NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] = T.float16(0)
|
||||
for ax2_fused_u_fused_0 in T.serial(T.int64(112), annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
|
||||
for ax0_0_1, ax1 in T.grid(T.int64(2), T.int64(8)):
|
||||
for ax0_1 in T.vectorized(T.int64(1)):
|
||||
with T.sblock("dequantize"):
|
||||
v0 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1 * T.int64(2) + ax0_0_1 + ax0_1)
|
||||
v1 = T.axis.spatial(T.int64(28672), ax2_fused_u_fused_0 * T.int64(256) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(8) + ax1)
|
||||
T.reads(lv429[v0, v1 // T.int64(8)], lv430[v0, v1 // T.int64(32)])
|
||||
T.writes(dequantize_intermediate_intermediate_local[v0, v1])
|
||||
dequantize_intermediate_intermediate_local[v0, v1] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv429[v0, v1 // T.int64(8)], T.Cast("uint32", v1 % T.int64(8) * T.int64(4))), T.uint32(15))) - T.float16(7)) * lv430[v0, v1 // T.int64(32)]
|
||||
for ax0_1, u_fused_ax1_fused_fused_2, ax2_fused_u_fused_2 in T.grid(T.int64(4), T.int64(2), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2)
|
||||
vax2_fused_u_fused_0, vax2_fused_u_fused_2 = T.axis.remap("RR", [ax2_fused_u_fused_0, ax2_fused_u_fused_2])
|
||||
T.reads(NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1], lv807[v0, T.int64(0), vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], dequantize_intermediate_intermediate_local[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)])
|
||||
T.writes(NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1])
|
||||
NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] = NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] + T.if_then_else(v0 < batch_size, lv807[v0, T.int64(0), vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], T.float16(0)) * dequantize_intermediate_intermediate_local[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)]
|
||||
for ax3_fused_0_ax3_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
for ax3_fused_2_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
|
||||
for ax2 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.spatial(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.reads()
|
||||
T.writes(NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] = T.float16(0)
|
||||
for ax1 in T.serial(T.int64(0), T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.reads(NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1], NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] = NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] + NT_matmul_intermediate_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, T.int64(0), v1]
|
||||
for ax2_fused_2, ax1 in T.grid(T.int64(2), T.int64(4)):
|
||||
for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
with T.sblock("NT_matmul"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.reduce(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax1)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax2_fused_0_ax2_fused_1_fused * T.int64(2) + ax2_fused_2)
|
||||
T.reads(NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul_intermediate_pad_local[v0, T.int64(0), v1])
|
||||
with T.init():
|
||||
NT_matmul_intermediate_pad_local[v0, T.int64(0), v1] = T.float16(0)
|
||||
NT_matmul_intermediate_pad_local[v0, T.int64(0), v1] = NT_matmul_intermediate_pad_local[v0, T.int64(0), v1] + NT_matmul_intermediate_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1]
|
||||
for ax0 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax1_fused_0_ax1_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax1_fused_2 in T.serial(T.int64(0), T.int64(2)):
|
||||
with T.sblock("NT_matmul_intermediate_pad"):
|
||||
v0 = T.axis.spatial(batch_size, ax0_0 * T.int64(4) + ax0)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax1_fused_0_ax1_fused_1_fused * T.int64(2) + ax1_fused_2)
|
||||
T.where((ax0_0 - (batch_size + T.int64(3)) // T.int64(4) < T.int64(0) or ax0_0 * T.int64(4) + ax0 == T.int64(0)) and ax0_0 * T.int64(4) + ax0 < batch_size)
|
||||
T.reads(NT_matmul_intermediate_pad_local[v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul_intermediate[v0, T.int64(0), v1])
|
||||
NT_matmul_intermediate[v0, T.int64(0), v1] = NT_matmul_intermediate_pad_local[v0, T.int64(0), v1]
|
||||
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_batch_gemv():
|
||||
N = 4096
|
||||
K = 4096
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_A: T.handle, B: T.Buffer((T.int64(N), T.int64(K)), "float16"), var_NT_matmul: T.handle):
|
||||
T.func_attr({"tirx.noalias": True, "tirx.HoistIfThenElseExprWithBlock": 1})
|
||||
batch_size = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, T.int64(1), T.int64(K)), "float16")
|
||||
NT_matmul = T.match_buffer(var_NT_matmul, (batch_size, T.int64(1), T.int64(N)), "float16")
|
||||
# with T.sblock("root"):
|
||||
for i0, i1, i2, k in T.grid(batch_size, T.int64(1), T.int64(N), T.int64(K)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(A[v_i0, v_i1, v_k], B[v_i2, v_k])
|
||||
T.writes(NT_matmul[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
NT_matmul[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
NT_matmul[v_i0, v_i1, v_i2] = NT_matmul[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_A: T.handle, B: T.Buffer((T.int64(4096), T.int64(4096)), "float16"), var_NT_matmul: T.handle):
|
||||
T.func_attr({"tirx.HoistIfThenElseExprWithBlock": 1, "tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
batch_size = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, T.int64(1), T.int64(4096)), "float16")
|
||||
NT_matmul = T.match_buffer(var_NT_matmul, (batch_size, T.int64(1), T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
NT_matmul_pad_local = T.sblock_alloc_buffer(((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
NT_matmul_pad_rf_local = T.sblock_alloc_buffer((T.int64(128), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
NT_matmul_pad_rf_local_1 = T.sblock_alloc_buffer((T.int64(32), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(1), T.int64(4096)), "float16", scope="local")
|
||||
for ax0_0 in T.thread_binding((batch_size + T.int64(3)) // T.int64(4), thread="blockIdx.y"):
|
||||
for u_fused_ax1_fused_fused_0 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
||||
for u_fused_ax1_fused_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
for ax0_1_init, u_fused_ax1_fused_fused_2_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1_init)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2_init)
|
||||
T.reads()
|
||||
T.writes(NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1])
|
||||
NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] = T.float16(0)
|
||||
for ax2_fused_u_fused_0 in T.serial(T.int64(16), annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
|
||||
for ax0_1, u_fused_ax1_fused_fused_2, ax2_fused_u_fused_2 in T.grid(T.int64(4), T.int64(2), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2)
|
||||
vax2_fused_u_fused_0, vax2_fused_u_fused_2 = T.axis.remap("RR", [ax2_fused_u_fused_0, ax2_fused_u_fused_2])
|
||||
T.reads(NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1], A[v0, T.int64(0), vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], B[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)])
|
||||
T.writes(NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1])
|
||||
NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] = NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, T.int64(0), v1] + T.if_then_else(v0 < batch_size, A[v0, T.int64(0), vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], T.float16(0)) * B[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)]
|
||||
for ax3_fused_0_ax3_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
for ax3_fused_2_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 8, "pragma_unroll_explicit": 1}):
|
||||
for ax2 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.spatial(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.reads()
|
||||
T.writes(NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] = T.float16(0)
|
||||
for ax1 in T.serial(T.int64(0), T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.reads(NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1], NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] = NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1] + NT_matmul_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, T.int64(0), v1]
|
||||
for ax2_fused_2, ax1 in T.grid(T.int64(2), T.int64(4)):
|
||||
for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.y"):
|
||||
with T.sblock("NT_matmul"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.reduce(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax1)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax2_fused_0_ax2_fused_1_fused * T.int64(2) + ax2_fused_2)
|
||||
T.reads(NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul_pad_local[v0, T.int64(0), v1])
|
||||
with T.init():
|
||||
NT_matmul_pad_local[v0, T.int64(0), v1] = T.float16(0)
|
||||
NT_matmul_pad_local[v0, T.int64(0), v1] = NT_matmul_pad_local[v0, T.int64(0), v1] + NT_matmul_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, T.int64(0), v1]
|
||||
for ax0 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax1_fused_0_ax1_fused_1_fused in T.thread_binding(T.int64(8), thread="threadIdx.x"):
|
||||
for ax1_fused_2 in T.serial(T.int64(0), T.int64(2)):
|
||||
with T.sblock("NT_matmul_pad"):
|
||||
v0 = T.axis.spatial(batch_size, ax0_0 * T.int64(4) + ax0)
|
||||
v1 = T.axis.spatial(T.int64(4096), u_fused_ax1_fused_fused_0 * T.int64(16) + ax1_fused_0_ax1_fused_1_fused * T.int64(2) + ax1_fused_2)
|
||||
T.where((ax0_0 - (batch_size + T.int64(3)) // T.int64(4) < T.int64(0) or ax0_0 * T.int64(4) + ax0 == T.int64(0)) and ax0_0 * T.int64(4) + ax0 < batch_size)
|
||||
T.reads(NT_matmul_pad_local[v0, T.int64(0), v1])
|
||||
T.writes(NT_matmul[v0, T.int64(0), v1])
|
||||
NT_matmul[v0, T.int64(0), v1] = NT_matmul_pad_local[v0, T.int64(0), v1]
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_reduction_symbolic_var():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
kv_seq_len = T.int64()
|
||||
A = T.match_buffer(var_A, (T.int64(1), T.int64(32), T.int64(1), kv_seq_len))
|
||||
B = T.match_buffer(var_B, (T.int64(1), T.int64(32), kv_seq_len, T.int64(128)))
|
||||
# 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), kv_seq_len):
|
||||
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.float32(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]
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], before)
|
||||
|
||||
|
||||
def test_small_spatial_axis():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(var_A: T.handle, B: T.Buffer((T.int64(8), T.int64(4096)), "float16"), var_C: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
batch_size = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, T.int64(4096)), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, T.int64(8)), "float16")
|
||||
for i0, i1, k in T.grid(batch_size, T.int64(8), T.int64(4096)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
||||
T.reads(A[v_i0, v_k], B[v_i1, v_k])
|
||||
T.writes(C[v_i0, v_i1])
|
||||
with T.init():
|
||||
C[v_i0, v_i1] = T.float16(0)
|
||||
C[v_i0, v_i1] = C[v_i0, v_i1] + A[v_i0, v_k] * B[v_i1, v_k]
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_A: T.handle, B: T.Buffer((T.int64(8), T.int64(4096)), "float16"), var_C: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
batch_size = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, T.int64(4096)), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, T.int64(8)), "float16")
|
||||
# with T.sblock("root"):
|
||||
C_pad_local = T.sblock_alloc_buffer(((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(8)), "float16", scope="local")
|
||||
C_pad_rf_local = T.sblock_alloc_buffer((T.int64(128), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(8)), "float16", scope="local")
|
||||
C_pad_rf_local_1 = T.sblock_alloc_buffer((T.int64(32), (batch_size + T.int64(3)) // T.int64(4) * T.int64(4), T.int64(8)), "float16", scope="local")
|
||||
for ax0_0 in T.thread_binding((batch_size + T.int64(3)) // T.int64(4), thread="blockIdx.y"):
|
||||
for u_fused_ax1_fused_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
||||
for u_fused_ax1_fused_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for ax0_1_init, u_fused_ax1_fused_fused_2_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1_init)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1_init)
|
||||
v1 = T.axis.spatial(T.int64(8), u_fused_ax1_fused_fused_0 * T.int64(32) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2_init)
|
||||
T.where((u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1) * T.int64(2) + u_fused_ax1_fused_fused_2_init < T.int64(8))
|
||||
T.reads()
|
||||
T.writes(C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1])
|
||||
C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] = T.float16(0)
|
||||
for ax2_fused_u_fused_0 in T.serial(T.int64(16), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax0_1, u_fused_ax1_fused_fused_2, ax2_fused_u_fused_2 in T.grid(T.int64(4), T.int64(2), T.int64(2)):
|
||||
for ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 in T.vectorized(T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused = T.axis.spatial(T.int64(128), ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + ax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax0_1)
|
||||
v1 = T.axis.spatial(T.int64(8), u_fused_ax1_fused_fused_0 * T.int64(32) + u_fused_ax1_fused_fused_1 * T.int64(2) + u_fused_ax1_fused_fused_2)
|
||||
vax2_fused_u_fused_0, vax2_fused_u_fused_2 = T.axis.remap("RR", [ax2_fused_u_fused_0, ax2_fused_u_fused_2])
|
||||
T.where((u_fused_ax1_fused_fused_0 * T.int64(16) + u_fused_ax1_fused_fused_1) * T.int64(2) + u_fused_ax1_fused_fused_2 < T.int64(8))
|
||||
T.reads(C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1], A[v0, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], B[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)])
|
||||
T.writes(C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1])
|
||||
C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] = C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused, v0, v1] + T.if_then_else(v0 < batch_size, A[v0, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)], T.float16(0)) * B[v1, vax2_fused_u_fused_0 * T.int64(256) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused // T.int64(4) * T.int64(8) + vax2_fused_u_fused_2 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused % T.int64(4)]
|
||||
for ax3_fused_0_ax3_fused_1_fused in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for ax3_fused_2_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax2 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax3_fused_2_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("NT_matmul_rf_init"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.spatial(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(8), ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.where((T.Mul(T.int64(0), T.int64(16)) + ax3_fused_0_ax3_fused_1_fused % T.int64(16)) * T.int64(2) + (ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1) < T.int64(8))
|
||||
T.reads()
|
||||
T.writes(C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1])
|
||||
C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1] = T.float16(0)
|
||||
for ax1 in T.serial(T.int64(0), T.int64(4)):
|
||||
with T.sblock("NT_matmul_rf_update"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax2)
|
||||
v1 = T.axis.spatial(T.int64(8), ax3_fused_0_ax3_fused_1_fused * T.int64(2) + ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1)
|
||||
T.where((T.Mul(T.int64(0), T.int64(16)) + ax3_fused_0_ax3_fused_1_fused % T.int64(16)) * T.int64(2) + (ax3_fused_2_0 * T.int64(2) + ax3_fused_2_1) < T.int64(8))
|
||||
T.reads(C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1], C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, v1])
|
||||
T.writes(C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1])
|
||||
C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1] = C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1] + C_pad_rf_local[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 * T.int64(4) + vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_1, v0, v1]
|
||||
for ax2_fused_2, ax1 in T.grid(T.int64(2), T.int64(4)):
|
||||
for ax2_fused_0_ax2_fused_1_fused in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
||||
for ax0 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
with T.sblock("NT_matmul"):
|
||||
vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0 = T.axis.reduce(T.int64(32), ax0)
|
||||
v0 = T.axis.spatial((batch_size + T.int64(3)) // T.int64(4) * T.int64(4), ax0_0 * T.int64(4) + ax1)
|
||||
v1 = T.axis.spatial(T.int64(8), ax2_fused_0_ax2_fused_1_fused * T.int64(2) + ax2_fused_2)
|
||||
T.where((T.Mul(T.int64(0), T.int64(16)) + ax2_fused_0_ax2_fused_1_fused % T.int64(16)) * T.int64(2) + ax2_fused_2 < T.int64(8))
|
||||
T.reads(C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1])
|
||||
T.writes(C_pad_local[v0, v1])
|
||||
with T.init():
|
||||
C_pad_local[v0, v1] = T.float16(0)
|
||||
C_pad_local[v0, v1] = C_pad_local[v0, v1] + C_pad_rf_local_1[vax2_fused_u_fused_1_ax2_fused_u_fused_3_fused_0, v0, v1]
|
||||
for ax0 in T.serial(T.int64(0), T.int64(4)):
|
||||
for ax1_fused_0_ax1_fused_1_fused in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
||||
for ax1_fused_2 in T.serial(T.int64(0), T.int64(2)):
|
||||
with T.sblock("C_pad"):
|
||||
v0 = T.axis.spatial(batch_size, ax0_0 * T.int64(4) + ax0)
|
||||
v1 = T.axis.spatial(T.int64(8), ax1_fused_0_ax1_fused_1_fused * T.int64(2) + ax1_fused_2)
|
||||
T.where((ax0_0 - (batch_size + T.int64(3)) // T.int64(4) < T.int64(0) or ax0_0 * T.int64(4) + ax0 == T.int64(0)) and ax0_0 * T.int64(4) + ax0 < batch_size and (T.Mul(T.int64(0), T.int64(16)) + ax1_fused_0_ax1_fused_1_fused % T.int64(16)) * T.int64(2) + ax1_fused_2 < T.int64(8))
|
||||
T.reads(C_pad_local[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_pad_local[v0, v1]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": func})
|
||||
with Target("cuda"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_outer_reduction():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
B0: T.Buffer((512, 6144), "uint32"),
|
||||
B1: T.Buffer((128, 6144), "float16"),
|
||||
var_A: T.handle,
|
||||
var_C: T.handle
|
||||
):
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 6144), "float16")
|
||||
compute = T.sblock_alloc_buffer((4096, 6144), "float16")
|
||||
B = T.sblock_alloc_buffer((4096, 6144), "float16")
|
||||
for i0, i1 in T.grid(4096, 6144):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
compute[v_i0, v_i1] = T.Cast("float16", T.bitwise_and(T.shift_right(B0[v_i0 // 8, v_i1], T.Cast("uint32", v_i0 % 8 * 4)), T.uint32(15)))
|
||||
for i0, i1 in T.grid(4096, 6144):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
B[v_i0, v_i1] = (compute[v_i0, v_i1] - T.float16(7)) * B1[v_i0 // 32, v_i1]
|
||||
for i0, i1, i2, k in T.grid(batch_size, 1, 6144, 4096):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
C[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * B[v_k, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(B0: T.Buffer((512, 6144), "uint32"), B1: T.Buffer((128, 6144), "float16"), var_A: T.handle, var_C: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 6144), "float16")
|
||||
# with T.sblock("root"):
|
||||
B_local = T.sblock_alloc_buffer((4096, 6144), "float16", scope="local")
|
||||
A_pad_shared = T.sblock_alloc_buffer(((batch_size + 3) // 4 * 4, 1, 4096), "float16", scope="shared")
|
||||
C_pad_local = T.sblock_alloc_buffer(((batch_size + 3) // 4 * 4, 1, 6144), "float16", scope="local")
|
||||
C_pad_rf_local = T.sblock_alloc_buffer((32, (batch_size + 3) // 4 * 4, 1, 6144), "float16", scope="local")
|
||||
C_pad_rf_local_1 = T.sblock_alloc_buffer((4, (batch_size + 3) // 4 * 4, 1, 6144), "float16", scope="local")
|
||||
B0_local = T.sblock_alloc_buffer((512, 6144), "uint32", scope="local")
|
||||
B1_local = T.sblock_alloc_buffer((128, 6144), "float16", scope="local")
|
||||
for ax0_0 in T.thread_binding((batch_size + 3) // 4, thread="blockIdx.y"):
|
||||
for ax1_fused_0 in T.thread_binding(96, thread="blockIdx.x"):
|
||||
for ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for ax2_fused_1_ax2_fused_3_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
|
||||
for ax0_1_init, ax2_fused_1_ax2_fused_3_fused_1_0_init in T.grid(4, 2):
|
||||
for ax2_fused_1_ax2_fused_3_fused_1_1_init in T.vectorized(4):
|
||||
with T.sblock("matmul_rf_init"):
|
||||
vax2_fused_1_ax2_fused_3_fused = T.axis.spatial(32, ax2_fused_1_ax2_fused_3_fused_0 * 8 + ax2_fused_1_ax2_fused_3_fused_1_0_init * 4 + ax2_fused_1_ax2_fused_3_fused_1_1_init)
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + ax0_1_init)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
T.reads()
|
||||
T.writes(C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1])
|
||||
C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1] = T.float16(0)
|
||||
for ax2_fused_0 in range(32):
|
||||
for ax0_ax1_fused in T.vectorized(4):
|
||||
with T.sblock("B0_local"):
|
||||
v0 = T.axis.spatial(512, ax2_fused_0 * 16 + ax2_fused_1_ax2_fused_3_fused_0 * 4 + ax0_ax1_fused)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
T.reads(B0[v0, v1])
|
||||
T.writes(B0_local[v0, v1])
|
||||
B0_local[v0, v1] = B0[v0, v1]
|
||||
for ax0_ax1_fused in T.vectorized(1):
|
||||
with T.sblock("B1_local"):
|
||||
v0 = T.axis.spatial(128, ax2_fused_0 * 4 + ax2_fused_1_ax2_fused_3_fused_0)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
T.reads(B1[v0, v1])
|
||||
T.writes(B1_local[v0, v1])
|
||||
B1_local[v0, v1] = B1[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.thread_binding(4, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(2):
|
||||
with T.sblock("A_pad"):
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) // 128)
|
||||
v1 = T.axis.spatial(4096, ax2_fused_0 * 128 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) % 128)
|
||||
T.reads(A[v0, 0, v1])
|
||||
T.writes(A_pad_shared[v0, 0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 1]]})
|
||||
A_pad_shared[v0, 0, v1] = T.if_then_else(v0 < batch_size, A[v0, 0, v1], T.float16(0))
|
||||
for ax2_fused_2 in range(4):
|
||||
for ax0_ax1_fused_0 in range(2):
|
||||
for ax0_ax1_fused_1 in T.vectorized(4):
|
||||
with T.sblock("dequantize"):
|
||||
v0 = T.axis.spatial(4096, ax2_fused_0 * 128 + ax2_fused_1_ax2_fused_3_fused_0 * 32 + ax2_fused_2 * 8 + ax0_ax1_fused_0 * 4 + ax0_ax1_fused_1)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
T.reads(B0_local[v0 // 8, v1], B1_local[v0 // 32, v1])
|
||||
T.writes(B_local[v0, v1])
|
||||
B_local[v0, v1] = (T.Cast("float16", T.bitwise_and(T.shift_right(B0_local[v0 // 8, v1], T.Cast("uint32", v0 % 8 * 4)), T.uint32(15))) - T.float16(7)) * B1_local[v0 // 32, v1]
|
||||
for ax0_1, ax2_fused_1_ax2_fused_3_fused_1_0 in T.grid(4, 2):
|
||||
for ax2_fused_1_ax2_fused_3_fused_1_1 in T.vectorized(4):
|
||||
with T.sblock("matmul_rf_update"):
|
||||
vax2_fused_1_ax2_fused_3_fused = T.axis.spatial(32, ax2_fused_1_ax2_fused_3_fused_0 * 8 + ax2_fused_1_ax2_fused_3_fused_1_0 * 4 + ax2_fused_1_ax2_fused_3_fused_1_1)
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + ax0_1)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1_fused_1)
|
||||
vax2_fused_0, vax2_fused_2 = T.axis.remap("RR", [ax2_fused_0, ax2_fused_2])
|
||||
T.reads(C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1], A_pad_shared[v0, 0, vax2_fused_0 * 128 + vax2_fused_1_ax2_fused_3_fused // 8 * 32 + vax2_fused_2 * 8 + vax2_fused_1_ax2_fused_3_fused % 8], B_local[vax2_fused_0 * 128 + vax2_fused_1_ax2_fused_3_fused // 8 * 32 + vax2_fused_2 * 8 + vax2_fused_1_ax2_fused_3_fused % 8, v1])
|
||||
T.writes(C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1])
|
||||
C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1] = C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused, v0, 0, v1] + A_pad_shared[v0, 0, vax2_fused_0 * 128 + vax2_fused_1_ax2_fused_3_fused // 8 * 32 + vax2_fused_2 * 8 + vax2_fused_1_ax2_fused_3_fused % 8] * B_local[vax2_fused_0 * 128 + vax2_fused_1_ax2_fused_3_fused // 8 * 32 + vax2_fused_2 * 8 + vax2_fused_1_ax2_fused_3_fused % 8, v1]
|
||||
for ax3 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(4, thread="threadIdx.y"):
|
||||
for ax2_init in range(4):
|
||||
with T.sblock("matmul_rf_init"):
|
||||
vax2_fused_1_ax2_fused_3_fused_0 = T.axis.spatial(4, ax0)
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + ax2_init)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax3)
|
||||
T.reads()
|
||||
T.writes(C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1])
|
||||
C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1] = T.float16(0)
|
||||
for ax2, ax1 in T.grid(4, 8):
|
||||
with T.sblock("matmul_rf_update"):
|
||||
vax2_fused_1_ax2_fused_3_fused_0, vax2_fused_1_ax2_fused_3_fused_1 = T.axis.remap("SR", [ax0, ax1])
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + ax2)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax3)
|
||||
T.reads(C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1], C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused_0 * 8 + vax2_fused_1_ax2_fused_3_fused_1, v0, 0, v1])
|
||||
T.writes(C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1])
|
||||
C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1] = C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1] + C_pad_rf_local[vax2_fused_1_ax2_fused_3_fused_0 * 8 + vax2_fused_1_ax2_fused_3_fused_1, v0, 0, v1]
|
||||
for ax1 in range(4):
|
||||
for ax2 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for ax0 in T.thread_binding(4, thread="threadIdx.y"):
|
||||
with T.sblock("matmul"):
|
||||
vax2_fused_1_ax2_fused_3_fused_0 = T.axis.reduce(4, ax0)
|
||||
v0 = T.axis.spatial((batch_size + 3) // 4 * 4, ax0_0 * 4 + ax1)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax2)
|
||||
T.reads(C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1])
|
||||
T.writes(C_pad_local[v0, 0, v1])
|
||||
with T.init():
|
||||
C_pad_local[v0, 0, v1] = T.float16(0)
|
||||
C_pad_local[v0, 0, v1] = C_pad_local[v0, 0, v1] + C_pad_rf_local_1[vax2_fused_1_ax2_fused_3_fused_0, v0, 0, v1]
|
||||
for ax0 in range(4):
|
||||
for ax1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
with T.sblock("C_pad"):
|
||||
v0 = T.axis.spatial(batch_size, ax0_0 * 4 + ax0)
|
||||
v1 = T.axis.spatial(6144, ax1_fused_0 * 64 + ax1)
|
||||
T.where((ax0_0 - (batch_size + 3) // 4 < 0 or ax0_0 * 4 + ax0 == 0) and ax0_0 * 4 + ax0 < batch_size)
|
||||
T.reads(C_pad_local[v0, 0, v1])
|
||||
T.writes(C[v0, 0, v1])
|
||||
C[v0, 0, v1] = C_pad_local[v0, 0, v1]
|
||||
# fmt: on
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod) # pylint: disable=not-callable
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_low_batch_gemv_cuda_target_without_max_shared_memory_per_block():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_A: T.handle, B: T.Buffer((T.int64(128), T.int64(128)), "float16"), var_C: T.handle):
|
||||
T.func_attr({"tir.noalias": True})
|
||||
batch_size = T.int64()
|
||||
A = T.match_buffer(var_A, (batch_size, T.int64(1), T.int64(128)), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, T.int64(1), T.int64(128)), "float16")
|
||||
for i0, i1, i2, k in T.grid(batch_size, T.int64(1), T.int64(128), T.int64(128)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(A[v_i0, v_i1, v_k], B[v_i2, v_k])
|
||||
T.writes(C[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
C[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
||||
# fmt: on
|
||||
|
||||
target = Target({"kind": "cuda", "max_num_threads": 1024})
|
||||
assert target.attrs.get("max_shared_memory_per_block", None) is None
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.LowBatchGEMV(4))(mod)
|
||||
assert mod["main"].attrs["tirx.is_scheduled"] == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,858 @@
|
||||
# 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
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_matmul():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_inp0: T.handle, inp1: T.Buffer((T.int64(4096), T.int64(4096)), "float32"), var_matmul: T.handle):
|
||||
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)))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), m, T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
matmul[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
matmul[v_i0, v_i1, v_i2] = matmul[v_i0, v_i1, v_i2] + inp0[v_i0, v_i1, v_k] * inp1[v_k, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(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_ax2_0_fused in T.thread_binding(T.int64(64), thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding((m + T.int64(31)) // T.int64(32), thread="blockIdx.x"):
|
||||
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 ax1_3_init, ax2_3_0_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_3_1_init in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_init"):
|
||||
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) + ax1_1 * T.int64(32) + ax1_2 * T.int64(4) + ax1_3_init)
|
||||
v2 = T.axis.spatial(T.int64(4096), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0_init * T.int64(2) + ax2_3_1_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), ax0_ax2_0_fused * 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, ax1_3, ax2_3_0 in T.grid(T.int64(16), T.int64(4), T.int64(2)):
|
||||
for ax2_3_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_update"):
|
||||
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) + ax1_1 * T.int64(32) + ax1_2 * T.int64(4) + ax1_3)
|
||||
v2 = T.axis.spatial(T.int64(4096), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0 * T.int64(2) + ax2_3_1)
|
||||
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, ax1, ax2_0 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)
|
||||
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), ax0_ax2_0_fused * T.int64(64) + ax2_2 * T.int64(4) + ax2_0 * T.int64(2) + ax2_1_1)
|
||||
T.where(ax1_0 * T.int64(32) + ax1_2 * T.int64(4) + ax1 < m)
|
||||
T.reads(matmul_reindex_pad_local[v0, v1, v2])
|
||||
T.writes(matmul[T.int64(0), v1, v2])
|
||||
matmul[T.int64(0), v1, v2] = matmul_reindex_pad_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_int32():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def func(var_inp0: T.handle, inp1: T.Buffer((4096, 4096), "float32"), var_matmul: T.handle):
|
||||
m = T.int32()
|
||||
inp0 = T.match_buffer(var_inp0, (1, m, 4096))
|
||||
matmul = T.match_buffer(var_matmul, (1, m, 4096))
|
||||
for i0, i1, i2, k in T.grid(1, m, 4096, 4096):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
matmul[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
matmul[v_i0, v_i1, v_i2] = matmul[v_i0, v_i1, v_i2] + inp0[v_i0, v_i1, v_k] * inp1[v_k, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_inp0: T.handle, inp1: T.Buffer((4096, 4096), "float32"), var_matmul: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
m = T.int32()
|
||||
inp0 = T.match_buffer(var_inp0, (1, m, 4096))
|
||||
matmul = T.match_buffer(var_matmul, (1, m, 4096))
|
||||
# with T.sblock("root"):
|
||||
matmul_reindex_pad_local = T.sblock_alloc_buffer((1, (m + 31) // 32 * 32, 4096), scope="local")
|
||||
inp0_reindex_pad_shared = T.sblock_alloc_buffer((1, (m + 31) // 32 * 32, 4096), scope="shared")
|
||||
inp1_reindex_shared = T.sblock_alloc_buffer((1, 4096, 4096), scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(64, thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding((m + 31) // 32, thread="blockIdx.x"):
|
||||
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
|
||||
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
|
||||
for ax2_2 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_2 in T.thread_binding(8, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
|
||||
for ax2_3_1_init in T.vectorized(2):
|
||||
with T.sblock("matmul_init"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3_init)
|
||||
v2 = T.axis.spatial(4096, ax0_ax2_0_fused * 64 + ax2_1 * 64 + ax2_2 * 4 + ax2_3_0_init * 2 + ax2_3_1_init)
|
||||
T.reads()
|
||||
T.writes(matmul_reindex_pad_local[0, v1, v2])
|
||||
matmul_reindex_pad_local[0, v1, v2] = T.float32(0)
|
||||
for ax3_0 in range(256):
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(2):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("inp0_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 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(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(4):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("inp1_reindex_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(4096, ax0_ax2_0_fused * 64 + (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 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, ax1_3, ax2_3_0 in T.grid(16, 4, 2):
|
||||
for ax2_3_1 in T.vectorized(2):
|
||||
with T.sblock("matmul_update"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3)
|
||||
v2 = T.axis.spatial(4096, ax0_ax2_0_fused * 64 + ax2_1 * 64 + ax2_2 * 4 + ax2_3_0 * 2 + ax2_3_1)
|
||||
v3 = T.axis.reduce(4096, ax3_0 * 16 + ax3_1)
|
||||
T.reads(matmul_reindex_pad_local[0, v1, v2], inp0_reindex_pad_shared[0, v1, v3], inp1_reindex_shared[0, v2, v3])
|
||||
T.writes(matmul_reindex_pad_local[0, v1, v2])
|
||||
matmul_reindex_pad_local[0, v1, v2] = matmul_reindex_pad_local[0, v1, v2] + inp0_reindex_pad_shared[0, v1, v3] * inp1_reindex_shared[0, v2, v3]
|
||||
for ax0, ax1, ax2_0 in T.grid(1, 4, 2):
|
||||
for ax2_1_1 in T.vectorized(2):
|
||||
with T.sblock("matmul_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_2 * 4 + ax1)
|
||||
v2 = T.axis.spatial(4096, ax0_ax2_0_fused * 64 + ax2_2 * 4 + ax2_0 * 2 + ax2_1_1)
|
||||
T.where(ax1_0 * 32 + ax1_2 * 4 + ax1 < m)
|
||||
T.reads(matmul_reindex_pad_local[v0, v1, v2])
|
||||
T.writes(matmul[0, v1, v2])
|
||||
matmul[0, v1, v2] = matmul_reindex_pad_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": func})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_fused_matmul():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(W: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), S: T.Buffer((T.int64(128), T.int64(4096)), "uint32"), A: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32"), C: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32"), Out: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32")):
|
||||
var_decode_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)))
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(32), T.int64(4096)))
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(W[v_i // T.int64(8), v_j], S[v_i // T.int64(32), v_j])
|
||||
T.writes(var_decode_intermediate[v_i, v_j])
|
||||
var_decode_intermediate[v_i, v_j] = T.Cast("float32", T.bitwise_and(T.shift_right(W[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8) * T.int64(4))), T.uint32(15))) * T.reinterpret("float32", T.shift_left(T.bitwise_and(S[v_i // T.int64(32), v_j], T.uint32(65535)), T.uint32(16))) + T.reinterpret("float32", T.shift_left(T.bitwise_and(T.shift_right(S[v_i // T.int64(32), v_j], T.uint32(16)), T.uint32(65535)), T.uint32(16)))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(32), T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(A[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(32), T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(C[v_ax0, v_ax1, v_ax2], var_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(Out[v_ax0, v_ax1, v_ax2])
|
||||
Out[v_ax0, v_ax1, v_ax2] = C[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(W: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), S: T.Buffer((T.int64(128), T.int64(4096)), "uint32"), A: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32"), C: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32"), Out: T.Buffer((T.int64(1), T.int64(32), T.int64(4096)), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
var_matmul_intermediate_reindex_local = T.sblock_alloc_buffer((T.int64(1), T.int64(32), T.int64(4096)), scope="local")
|
||||
A_reindex_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(32), T.int64(4096)), scope="shared")
|
||||
var_decode_intermediate_reindex_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(4096), T.int64(4096)), scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(T.int64(64), thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
||||
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 ax1_3_init, ax2_3_0_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_3_1_init in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_init"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(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), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0_init * T.int64(2) + ax2_3_1_init)
|
||||
T.reads()
|
||||
T.writes(var_matmul_intermediate_reindex_local[T.int64(0), v1, v2])
|
||||
var_matmul_intermediate_reindex_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("A_reindex_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(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(A[v0, v1, v2])
|
||||
T.writes(A_reindex_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
A_reindex_shared[v0, v1, v2] = A[v0, v1, v2]
|
||||
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("var_decode_intermediate_reindex_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax0_ax2_0_fused * 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(W[v2 // T.int64(8), v1], S[v2 // T.int64(32), v1])
|
||||
T.writes(var_decode_intermediate_reindex_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
var_decode_intermediate_reindex_shared[v0, v1, v2] = T.Cast("float32", T.bitwise_and(T.shift_right(W[v2 // T.int64(8), v1], T.Cast("uint32", v2 % T.int64(8) * T.int64(4))), T.uint32(15))) * T.reinterpret("float32", T.shift_left(T.bitwise_and(S[v2 // T.int64(32), v1], T.uint32(65535)), T.uint32(16))) + T.reinterpret("float32", T.shift_left(T.bitwise_and(T.shift_right(S[v2 // T.int64(32), v1], T.uint32(16)), T.uint32(65535)), T.uint32(16)))
|
||||
for ax3_1, ax1_3, ax2_3_0 in T.grid(T.int64(16), T.int64(4), T.int64(2)):
|
||||
for ax2_3_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_update"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(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), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0 * T.int64(2) + ax2_3_1)
|
||||
v3 = T.axis.reduce(T.int64(4096), ax3_0 * T.int64(16) + ax3_1)
|
||||
T.reads(var_matmul_intermediate_reindex_local[T.int64(0), v1, v2], A_reindex_shared[T.int64(0), v1, v3], var_decode_intermediate_reindex_shared[T.int64(0), v2, v3])
|
||||
T.writes(var_matmul_intermediate_reindex_local[T.int64(0), v1, v2])
|
||||
var_matmul_intermediate_reindex_local[T.int64(0), v1, v2] = var_matmul_intermediate_reindex_local[T.int64(0), v1, v2] + A_reindex_shared[T.int64(0), v1, v3] * var_decode_intermediate_reindex_shared[T.int64(0), v2, v3]
|
||||
for ax0, ax1, ax2_0 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("var_matmul_intermediate_reindex_local"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.spatial(T.int64(32), ax1_2 * T.int64(4) + ax1)
|
||||
v2 = T.axis.spatial(T.int64(4096), ax0_ax2_0_fused * T.int64(64) + ax2_2 * T.int64(4) + ax2_0 * T.int64(2) + ax2_1_1)
|
||||
T.reads(C[T.int64(0), v1, v2], var_matmul_intermediate_reindex_local[v0, v1, v2])
|
||||
T.writes(Out[T.int64(0), v1, v2])
|
||||
Out[T.int64(0), v1, v2] = C[T.int64(0), v1, v2] + var_matmul_intermediate_reindex_local[v0, v1, v2]
|
||||
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_skip_gemv():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(W: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), S: T.Buffer((T.int64(128), T.int64(4096)), "uint32"), A: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float32"), C: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float32"), Out: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
var_decode_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)))
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(4096)))
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(W[v_i // T.int64(8), v_j], S[v_i // T.int64(32), v_j])
|
||||
T.writes(var_decode_intermediate[v_i, v_j])
|
||||
var_decode_intermediate[v_i, v_j] = T.Cast("float32", T.bitwise_and(T.shift_right(W[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8) * T.int64(4))), T.uint32(15))) * T.reinterpret("float32", T.shift_left(T.bitwise_and(S[v_i // T.int64(32), v_j], T.uint32(65535)), T.uint32(16))) + T.reinterpret("float32", T.shift_left(T.bitwise_and(T.shift_right(S[v_i // T.int64(32), v_j], T.uint32(16)), T.uint32(65535)), T.uint32(16)))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(A[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(C[v_ax0, v_ax1, v_ax2], var_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(Out[v_ax0, v_ax1, v_ax2])
|
||||
Out[v_ax0, v_ax1, v_ax2] = C[v_ax0, v_ax1, v_ax2] + var_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
# fmt: on
|
||||
|
||||
expected = before
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_output_fp32():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv13: T.Buffer((T.int64(4096), T.int64(512)), "uint32"), lv14: T.Buffer((T.int64(4096), T.int64(128)), "float16"), p_lv48: T.handle, lv13_1: T.Buffer((T.int64(4096),), "float16"), p_lv3: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv48 = T.match_buffer(p_lv48, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
lv3 = T.match_buffer(p_lv3, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
|
||||
var_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), n, T.int64(4096)))
|
||||
var_compute_intermediate = T.sblock_alloc_buffer((T.int64(4096),))
|
||||
var_T_add_intermediate = T.sblock_alloc_buffer((T.int64(1), n, T.int64(4096)))
|
||||
var_compute_intermediate_1 = T.sblock_alloc_buffer((T.int64(1), n, T.int64(4096)), "float16")
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv13[v_i, v_j // T.int64(8)], lv14[v_i, v_j // T.int64(32)])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv13[v_i, v_j // T.int64(8)], T.Cast("uint32", v_j % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv14[v_i, v_j // T.int64(32)]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), n, T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv48[v_i0, v_i1, v_k], p_output0_intermediate_1[v_k, v_i2])
|
||||
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv48[v_i0, v_i1, v_k]) * T.Cast("float32", p_output0_intermediate_1[v_k, v_i2])
|
||||
for i0 in range(T.int64(4096)):
|
||||
with T.sblock("compute"):
|
||||
v_i0 = T.axis.spatial(T.int64(4096), i0)
|
||||
T.reads(lv13_1[v_i0])
|
||||
T.writes(var_compute_intermediate[v_i0])
|
||||
var_compute_intermediate[v_i0] = T.Cast("float32", lv13_1[v_i0])
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), n, T.int64(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], var_compute_intermediate[v_ax2])
|
||||
T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + var_compute_intermediate[v_ax2]
|
||||
for i0, i1, i2 in T.grid(T.int64(1), n, T.int64(4096)):
|
||||
with T.sblock("compute_1"):
|
||||
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(var_T_add_intermediate[v_i0, v_i1, v_i2])
|
||||
T.writes(var_compute_intermediate_1[v_i0, v_i1, v_i2])
|
||||
var_compute_intermediate_1[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_add_intermediate[v_i0, v_i1, v_i2])
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), n, T.int64(4096)):
|
||||
with T.sblock("T_add_1"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(var_compute_intermediate_1[v_ax0, v_ax1, v_ax2], lv3[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_compute_intermediate_1[v_ax0, v_ax1, v_ax2] + lv3[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv13: T.Buffer((T.int64(4096), T.int64(512)), "uint32"), lv14: T.Buffer((T.int64(4096), T.int64(128)), "float16"), p_lv48: T.handle, lv13_1: T.Buffer((T.int64(4096),), "float16"), p_lv3: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv48 = T.match_buffer(p_lv48, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
lv3 = T.match_buffer(p_lv3, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (T.int64(1), n, T.int64(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
var_matmul_intermediate_reindex_pad_local = T.sblock_alloc_buffer((T.int64(1), (n + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="local")
|
||||
lv48_reindex_pad_shared = T.sblock_alloc_buffer((T.int64(1), (n + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), "float16", scope="shared")
|
||||
p_output0_intermediate_1_reindex_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(4096), T.int64(4096)), "float16", scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(T.int64(64), thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding((n + T.int64(31)) // T.int64(32), thread="blockIdx.x"):
|
||||
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 ax1_3_init, ax2_3_0_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_3_1_init in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_init"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0_init * T.int64(2) + ax2_3_1_init)
|
||||
T.reads()
|
||||
T.writes(var_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2])
|
||||
var_matmul_intermediate_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("lv48_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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(lv48[v0, v1, v2])
|
||||
T.writes(lv48_reindex_pad_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
lv48_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < n, lv48[v0, v1, v2], T.float16(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("p_output0_intermediate_1_reindex_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax0_ax2_0_fused * 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(lv13[v2, v1 // T.int64(8)], lv14[v2, v1 // T.int64(32)])
|
||||
T.writes(p_output0_intermediate_1_reindex_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
p_output0_intermediate_1_reindex_shared[v0, v1, v2] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv13[v2, v1 // T.int64(8)], T.Cast("uint32", v1 % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv14[v2, v1 // T.int64(32)]
|
||||
for ax3_1, ax1_3, ax2_3_0 in T.grid(T.int64(16), T.int64(4), T.int64(2)):
|
||||
for ax2_3_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("matmul_update"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0 * T.int64(2) + ax2_3_1)
|
||||
v3 = T.axis.reduce(T.int64(4096), ax3_0 * T.int64(16) + ax3_1)
|
||||
T.reads(var_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2], lv48_reindex_pad_shared[T.int64(0), v1, v3], p_output0_intermediate_1_reindex_shared[T.int64(0), v2, v3])
|
||||
T.writes(var_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2])
|
||||
var_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2] = var_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2] + T.Cast("float32", lv48_reindex_pad_shared[T.int64(0), v1, v3]) * T.Cast("float32", p_output0_intermediate_1_reindex_shared[T.int64(0), v2, v3])
|
||||
for ax0, ax1, ax2_0 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("var_matmul_intermediate_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.spatial((n + 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), ax0_ax2_0_fused * T.int64(64) + ax2_2 * T.int64(4) + ax2_0 * T.int64(2) + ax2_1_1)
|
||||
T.where(ax1_0 * T.int64(32) + ax1_2 * T.int64(4) + ax1 < n)
|
||||
T.reads(var_matmul_intermediate_reindex_pad_local[v0, v1, v2], lv13_1[v2], lv3[T.int64(0), v1, v2])
|
||||
T.writes(p_output0_intermediate[T.int64(0), v1, v2])
|
||||
p_output0_intermediate[T.int64(0), v1, v2] = T.Cast("float16", var_matmul_intermediate_reindex_pad_local[v0, v1, v2] + T.Cast("float32", lv13_1[v2])) + lv3[T.int64(0), v1, v2]
|
||||
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_inline_consumer_chain():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(p_lv26: T.handle, lv9: T.Buffer((T.int64(2048), T.int64(2048)), "float16"), p_lv52: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv26 = T.match_buffer(p_lv26, (n, T.int64(2048)), "float16")
|
||||
lv52 = T.match_buffer(p_lv52, (T.int64(1), n, T.int64(2048)))
|
||||
var_T_multiply_intermediate = T.match_buffer(p_output0, (n, T.int64(2048)), "float16")
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((n, T.int64(2048)), "float16")
|
||||
compute = T.sblock_alloc_buffer((n, T.int64(2048)), "float16")
|
||||
var_T_multiply_intermediate_1 = T.sblock_alloc_buffer((n, T.int64(2048)), "float16")
|
||||
var_T_squeeze_intermediate = T.sblock_alloc_buffer((n, T.int64(2048)))
|
||||
var_compute_intermediate = T.sblock_alloc_buffer((n, T.int64(2048)), "float16")
|
||||
for i0, i1, k in T.grid(n, T.int64(2048), T.int64(2048)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
||||
T.reads(lv26[v_i0, v_k], lv9[v_i1, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1] = var_NT_matmul_intermediate[v_i0, v_i1] + lv26[v_i0, v_k] * lv9[v_i1, v_k]
|
||||
for i0, i1 in T.grid(n, T.int64(2048)):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(var_NT_matmul_intermediate[v_i0, v_i1])
|
||||
T.writes(compute[v_i0, v_i1])
|
||||
compute[v_i0, v_i1] = T.sigmoid(var_NT_matmul_intermediate[v_i0, v_i1])
|
||||
for ax0, ax1 in T.grid(n, T.int64(2048)):
|
||||
with T.sblock("T_multiply"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(var_NT_matmul_intermediate[v_ax0, v_ax1], compute[v_ax0, v_ax1])
|
||||
T.writes(var_T_multiply_intermediate_1[v_ax0, v_ax1])
|
||||
var_T_multiply_intermediate_1[v_ax0, v_ax1] = var_NT_matmul_intermediate[v_ax0, v_ax1] * compute[v_ax0, v_ax1]
|
||||
for ax0, ax1 in T.grid(n, T.int64(2048)):
|
||||
with T.sblock("T_squeeze"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(lv52[T.int64(0), v_ax0, v_ax1])
|
||||
T.writes(var_T_squeeze_intermediate[v_ax0, v_ax1])
|
||||
var_T_squeeze_intermediate[v_ax0, v_ax1] = lv52[T.int64(0), v_ax0, v_ax1]
|
||||
for i0, i1 in T.grid(n, T.int64(2048)):
|
||||
with T.sblock("compute_1"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(var_T_squeeze_intermediate[v_i0, v_i1])
|
||||
T.writes(var_compute_intermediate[v_i0, v_i1])
|
||||
var_compute_intermediate[v_i0, v_i1] = T.Cast("float16", var_T_squeeze_intermediate[v_i0, v_i1])
|
||||
for ax0, ax1 in T.grid(n, T.int64(2048)):
|
||||
with T.sblock("T_multiply_1"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(var_compute_intermediate[v_ax0, v_ax1], var_T_multiply_intermediate_1[v_ax0, v_ax1])
|
||||
T.writes(var_T_multiply_intermediate[v_ax0, v_ax1])
|
||||
var_T_multiply_intermediate[v_ax0, v_ax1] = var_compute_intermediate[v_ax0, v_ax1] * var_T_multiply_intermediate_1[v_ax0, v_ax1]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(p_lv26: T.handle, lv9: T.Buffer((T.int64(2048), T.int64(2048)), "float16"), p_lv52: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int64()
|
||||
lv26 = T.match_buffer(p_lv26, (n, T.int64(2048)), "float16")
|
||||
lv52 = T.match_buffer(p_lv52, (T.int64(1), n, T.int64(2048)))
|
||||
var_T_multiply_intermediate = T.match_buffer(p_output0, (n, T.int64(2048)), "float16")
|
||||
# with T.sblock("root"):
|
||||
var_NT_matmul_intermediate_reindex_pad_local = T.sblock_alloc_buffer((T.int64(1), (n + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(2048)), "float16", scope="local")
|
||||
lv26_reindex_pad_shared = T.sblock_alloc_buffer((T.int64(1), (n + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(2048)), "float16", scope="shared")
|
||||
lv9_reindex_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(2048), T.int64(2048)), "float16", scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(T.int64(32), thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding((n + T.int64(31)) // T.int64(32), thread="blockIdx.x"):
|
||||
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 ax1_3_init, ax2_3_0_init in T.grid(T.int64(4), T.int64(2)):
|
||||
for ax2_3_1_init in T.vectorized(T.int64(2)):
|
||||
with T.sblock("NT_matmul_init"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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(2048), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0_init * T.int64(2) + ax2_3_1_init)
|
||||
T.reads()
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2])
|
||||
var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2] = T.float16(0)
|
||||
for ax3_0 in range(T.int64(128)):
|
||||
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("lv26_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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(2048), 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(lv26[v1, v2])
|
||||
T.writes(lv26_reindex_pad_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
lv26_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < n, lv26[v1, v2], T.float16(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("lv9_reindex_shared"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(T.int64(2048), ax0_ax2_0_fused * 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(2048), 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(lv9[v1, v2])
|
||||
T.writes(lv9_reindex_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
lv9_reindex_shared[v0, v1, v2] = lv9[v1, v2]
|
||||
for ax3_1, ax1_3, ax2_3_0 in T.grid(T.int64(16), T.int64(4), T.int64(2)):
|
||||
for ax2_3_1 in T.vectorized(T.int64(2)):
|
||||
with T.sblock("NT_matmul_update"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial((n + 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(2048), ax0_ax2_0_fused * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(4) + ax2_3_0 * T.int64(2) + ax2_3_1)
|
||||
v3 = T.axis.reduce(T.int64(2048), ax3_0 * T.int64(16) + ax3_1)
|
||||
T.reads(var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2], lv26_reindex_pad_shared[T.int64(0), v1, v3], lv9_reindex_shared[T.int64(0), v2, v3])
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2])
|
||||
var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2] = var_NT_matmul_intermediate_reindex_pad_local[T.int64(0), v1, v2] + lv26_reindex_pad_shared[T.int64(0), v1, v3] * lv9_reindex_shared[T.int64(0), v2, v3]
|
||||
for ax0, ax1, ax2_0 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("var_NT_matmul_intermediate_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.spatial((n + 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(2048), ax0_ax2_0_fused * T.int64(64) + ax2_2 * T.int64(4) + ax2_0 * T.int64(2) + ax2_1_1)
|
||||
T.reads(lv52[T.int64(0), v1, v2], var_NT_matmul_intermediate_reindex_pad_local[v0, v1, v2])
|
||||
T.where(ax1_0 * T.int64(32) + ax1_2 * T.int64(4) + ax1 < n)
|
||||
T.writes(var_T_multiply_intermediate[v1, v2])
|
||||
var_T_multiply_intermediate[v1, v2] = T.Cast("float16", lv52[T.int64(0), v1, v2]) * (var_NT_matmul_intermediate_reindex_pad_local[v0, v1, v2] * T.sigmoid(var_NT_matmul_intermediate_reindex_pad_local[v0, v1, v2]))
|
||||
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-gtx-1080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_android():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_inp0: T.handle, inp1: T.Buffer((T.int64(4096), T.int64(4096)), "float32"), var_matmul: T.handle):
|
||||
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)))
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), m, T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
matmul[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
matmul[v_i0, v_i1, v_i2] = matmul[v_i0, v_i1, v_i2] + inp0[v_i0, v_i1, v_k] * inp1[v_k, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(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"):
|
||||
inp0_reindex_pad = T.sblock_alloc_buffer((T.int64(1), (m + T.int64(15)) // T.int64(16), T.int64(4096), T.int64(16)))
|
||||
matmul_pad_local = T.sblock_alloc_buffer((T.int64(1), (m + T.int64(15)) // T.int64(16) * T.int64(16), T.int64(4096)), scope="local")
|
||||
inp0_reindex_pad_local = T.sblock_alloc_buffer((T.int64(1), (m + T.int64(15)) // T.int64(16), T.int64(4096), T.int64(16)), scope="local")
|
||||
for i0 in T.thread_binding(T.int64(1), thread="blockIdx.z"):
|
||||
for i1_0 in T.thread_binding(((m + T.int64(15)) // T.int64(16) * T.int64(16) + T.int64(63)) // T.int64(64), thread="blockIdx.y"):
|
||||
for i2_0 in T.thread_binding(T.int64(128), thread="blockIdx.x"):
|
||||
for i1_1 in T.thread_binding(T.int64(4), thread="threadIdx.y"):
|
||||
for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for i1_2 in T.vectorized(T.int64(16)):
|
||||
with T.sblock("inp0_reindex_pad"):
|
||||
v0 = T.axis.spatial(T.int64(1), i0)
|
||||
v1 = T.axis.spatial((m + T.int64(15)) // T.int64(16) * T.int64(16), i1_0 * T.int64(64) + i1_1 * T.int64(16) + i1_2)
|
||||
v2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(32) + i2_1)
|
||||
T.where((i1_0 * T.int64(4) + i1_1) * T.int64(16) + i1_2 < (m + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads(inp0[v0, v1, v2])
|
||||
T.writes(inp0_reindex_pad[v0, v1 // T.int64(16), v2, v1 % T.int64(16)])
|
||||
inp0_reindex_pad[v0, v1 // T.int64(16), v2, v1 % T.int64(16)] = T.if_then_else(v1 < m, inp0[v0, v1, v2], T.float32(0))
|
||||
for i2_0 in T.thread_binding(T.int64(16), thread="blockIdx.x"):
|
||||
for i0_i1_fused_0 in T.thread_binding(((m + T.int64(15)) // T.int64(16) * T.int64(16) + T.int64(63)) // T.int64(64), thread="blockIdx.y"):
|
||||
for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for i0_i1_fused_1 in T.thread_binding(T.int64(4), thread="threadIdx.y"):
|
||||
for i0_i1_fused_2_init in T.unroll(T.int64(16)):
|
||||
for i2_2_init in T.vectorized(T.int64(8)):
|
||||
with T.sblock("matmul_init"):
|
||||
v_i0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_i1 = T.axis.spatial((m + T.int64(15)) // T.int64(16) * T.int64(16), i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + i0_i1_fused_2_init)
|
||||
v_i2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2_init)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1) * T.int64(16) + i0_i1_fused_2_init < (m + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads()
|
||||
T.writes(matmul_pad_local[v_i0, v_i1, v_i2])
|
||||
matmul_pad_local[v_i0, v_i1, v_i2] = T.float32(0)
|
||||
for k_0, k_1 in T.grid(T.int64(128), T.int64(4)):
|
||||
for k_2 in T.unroll(T.int64(8)):
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(1)):
|
||||
for ax3 in T.vectorized(T.int64(16)):
|
||||
with T.sblock("inp0_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.spatial((m + T.int64(15)) // T.int64(16), i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 + ax1)
|
||||
v2 = T.axis.spatial(T.int64(4096), k_0 * T.int64(32) + k_1 * T.int64(8) + k_2 + ax2)
|
||||
v3 = T.axis.spatial(T.int64(16), ax3)
|
||||
T.where(i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 < (m + T.int64(15)) // T.int64(16))
|
||||
T.reads(inp0_reindex_pad[v0, v1, v2, v3])
|
||||
T.writes(inp0_reindex_pad_local[v0, v1, v2, v3])
|
||||
inp0_reindex_pad_local[v0, v1, v2, v3] = inp0_reindex_pad[v0, v1, v2, v3]
|
||||
for i0_i1_fused_2 in T.unroll(T.int64(16)):
|
||||
for i2_2 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("matmul_update"):
|
||||
v_i0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_i1 = T.axis.spatial((m + T.int64(15)) // T.int64(16) * T.int64(16), i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + i0_i1_fused_2)
|
||||
v_i2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2)
|
||||
v_k = T.axis.reduce(T.int64(4096), k_0 * T.int64(32) + k_1 * T.int64(8) + k_2)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1) * T.int64(16) + i0_i1_fused_2 < (m + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads(matmul_pad_local[v_i0, v_i1, v_i2], inp0_reindex_pad_local[v_i0, v_i1 // T.int64(16), v_k, v_i1 % T.int64(16)], inp1[v_k, v_i2])
|
||||
T.writes(matmul_pad_local[v_i0, v_i1, v_i2])
|
||||
matmul_pad_local[v_i0, v_i1, v_i2] = matmul_pad_local[v_i0, v_i1, v_i2] + inp0_reindex_pad_local[v_i0, v_i1 // T.int64(16), v_k, v_i1 % T.int64(16)] * inp1[v_k, v_i2]
|
||||
for ax0 in T.unroll(T.int64(16)):
|
||||
for ax1 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("matmul_pad"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(m, i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0)
|
||||
v2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 - (m + T.int64(15)) // T.int64(16) < T.int64(0) or i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0 == T.int64(0)) and i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0 < m)
|
||||
T.reads(matmul_pad_local[v0, v1, v2])
|
||||
T.writes(matmul[v0, v1, v2])
|
||||
matmul[v0, v1, v2] = matmul_pad_local[v0, v1, v2]
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-android"}):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_fused_dequant_matmul_android():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv452: T.Buffer((T.int64(512), T.int64(12288)), "uint32"), lv453: T.Buffer((T.int64(128), T.int64(12288)), "float16"), p_rms_norm130: T.handle, transformer_h_0_attn_c_attn_bias3: T.Buffer((T.int64(12288),), "float16"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
seq_len = T.int64()
|
||||
rms_norm130 = T.match_buffer(p_rms_norm130, (T.int64(1), seq_len, T.int64(4096)), "float16")
|
||||
T_add_intermediate_intermediate = T.match_buffer(p_output0, (T.int64(1), seq_len, T.int64(12288)), "float16")
|
||||
# with T.sblock("root"):
|
||||
compute = T.sblock_alloc_buffer((T.int64(4096), T.int64(12288)), "float16")
|
||||
dequantize_intermediate_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(12288)), "float16")
|
||||
matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), seq_len, T.int64(12288)), "float16")
|
||||
for i0, i1 in T.grid(T.int64(4096), T.int64(12288)):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(lv452[v_i0 // T.int64(8), v_i1])
|
||||
T.writes(compute[v_i0, v_i1])
|
||||
compute[v_i0, v_i1] = T.Cast("float16", T.bitwise_and(T.shift_right(lv452[v_i0 // T.int64(8), v_i1], T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4))), T.uint32(15)))
|
||||
for i0, i1 in T.grid(T.int64(4096), T.int64(12288)):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(compute[v_i0, v_i1], lv453[v_i0 // T.int64(32), v_i1])
|
||||
T.writes(dequantize_intermediate_intermediate[v_i0, v_i1])
|
||||
dequantize_intermediate_intermediate[v_i0, v_i1] = (compute[v_i0, v_i1] - T.float16(7)) * lv453[v_i0 // T.int64(32), v_i1]
|
||||
for i0, i1, i2, k in T.grid(T.int64(1), seq_len, T.int64(12288), T.int64(4096)):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(rms_norm130[v_i0, v_i1, v_k], dequantize_intermediate_intermediate[v_k, v_i2])
|
||||
T.writes(matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
matmul_intermediate[v_i0, v_i1, v_i2] = matmul_intermediate[v_i0, v_i1, v_i2] + rms_norm130[v_i0, v_i1, v_k] * dequantize_intermediate_intermediate[v_k, v_i2]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), seq_len, T.int64(12288)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(matmul_intermediate[v_ax0, v_ax1, v_ax2], transformer_h_0_attn_c_attn_bias3[v_ax2])
|
||||
T.writes(T_add_intermediate_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T_add_intermediate_intermediate[v_ax0, v_ax1, v_ax2] = matmul_intermediate[v_ax0, v_ax1, v_ax2] + transformer_h_0_attn_c_attn_bias3[v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv452: T.Buffer((T.int64(512), T.int64(12288)), "uint32"), lv453: T.Buffer((T.int64(128), T.int64(12288)), "float16"), p_rms_norm130: T.handle, transformer_h_0_attn_c_attn_bias3: T.Buffer((T.int64(12288),), "float16"), p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
seq_len = T.int64()
|
||||
rms_norm130 = T.match_buffer(p_rms_norm130, (T.int64(1), seq_len, T.int64(4096)), "float16")
|
||||
T_add_intermediate_intermediate = T.match_buffer(p_output0, (T.int64(1), seq_len, T.int64(12288)), "float16")
|
||||
# with T.sblock("root"):
|
||||
dequantize_intermediate_intermediate_local = T.sblock_alloc_buffer((T.int64(4096), T.int64(12288)), "float16", scope="local")
|
||||
rms_norm130_reindex_pad = T.sblock_alloc_buffer((T.int64(1), (seq_len + T.int64(15)) // T.int64(16), T.int64(4096), T.int64(16)), "float16")
|
||||
matmul_intermediate_pad_local = T.sblock_alloc_buffer((T.int64(1), (seq_len + T.int64(15)) // T.int64(16) * T.int64(16), T.int64(12288)), "float16", scope="local")
|
||||
rms_norm130_reindex_pad_local = T.sblock_alloc_buffer((T.int64(1), (seq_len + T.int64(15)) // T.int64(16), T.int64(4096), T.int64(16)), "float16", scope="local")
|
||||
lv452_local = T.sblock_alloc_buffer((T.int64(512), T.int64(12288)), "uint32", scope="local")
|
||||
lv453_local = T.sblock_alloc_buffer((T.int64(128), T.int64(12288)), "float16", scope="local")
|
||||
for i0 in T.thread_binding(T.int64(1), thread="blockIdx.z"):
|
||||
for i1_0 in T.thread_binding(((seq_len + T.int64(15)) // T.int64(16) * T.int64(16) + T.int64(63)) // T.int64(64), thread="blockIdx.y"):
|
||||
for i2_0 in T.thread_binding(T.int64(128), thread="blockIdx.x"):
|
||||
for i1_1 in T.thread_binding(T.int64(4), thread="threadIdx.y"):
|
||||
for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for i1_2 in T.vectorized(T.int64(16)):
|
||||
with T.sblock("rms_norm130_reindex_pad"):
|
||||
v0 = T.axis.spatial(T.int64(1), i0)
|
||||
v1 = T.axis.spatial((seq_len + T.int64(15)) // T.int64(16) * T.int64(16), i1_0 * T.int64(64) + i1_1 * T.int64(16) + i1_2)
|
||||
v2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(32) + i2_1)
|
||||
T.where((i1_0 * T.int64(4) + i1_1) * T.int64(16) + i1_2 < (seq_len + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads(rms_norm130[v0, v1, v2])
|
||||
T.writes(rms_norm130_reindex_pad[v0, v1 // T.int64(16), v2, v1 % T.int64(16)])
|
||||
rms_norm130_reindex_pad[v0, v1 // T.int64(16), v2, v1 % T.int64(16)] = T.if_then_else(v1 < seq_len, rms_norm130[v0, v1, v2], T.float16(0))
|
||||
for i2_0 in T.thread_binding(T.int64(48), thread="blockIdx.x"):
|
||||
for i0_i1_fused_0 in T.thread_binding(((seq_len + T.int64(15)) // T.int64(16) * T.int64(16) + T.int64(63)) // T.int64(64), thread="blockIdx.y"):
|
||||
for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"):
|
||||
for i0_i1_fused_1 in T.thread_binding(T.int64(4), thread="threadIdx.y"):
|
||||
for i0_i1_fused_2_init in T.unroll(T.int64(16)):
|
||||
for i2_2_init in T.vectorized(T.int64(8)):
|
||||
with T.sblock("matmul_init"):
|
||||
v_i0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_i1 = T.axis.spatial((seq_len + T.int64(15)) // T.int64(16) * T.int64(16), i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + i0_i1_fused_2_init)
|
||||
v_i2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2_init)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1) * T.int64(16) + i0_i1_fused_2_init < (seq_len + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads()
|
||||
T.writes(matmul_intermediate_pad_local[v_i0, v_i1, v_i2])
|
||||
matmul_intermediate_pad_local[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
for k_0 in range(T.int64(128)):
|
||||
for ax0 in range(T.int64(1)):
|
||||
for ax1 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("lv453_local"):
|
||||
v0 = T.axis.spatial(T.int64(128), k_0 + ax0)
|
||||
v1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1)
|
||||
T.reads(lv453[v0, v1])
|
||||
T.writes(lv453_local[v0, v1])
|
||||
lv453_local[v0, v1] = lv453[v0, v1]
|
||||
for k_1 in range(T.int64(4)):
|
||||
for ax0 in range(T.int64(1)):
|
||||
for ax1 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("lv452_local"):
|
||||
v0 = T.axis.spatial(T.int64(512), k_0 * T.int64(4) + k_1 + ax0)
|
||||
v1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1)
|
||||
T.reads(lv452[v0, v1])
|
||||
T.writes(lv452_local[v0, v1])
|
||||
lv452_local[v0, v1] = lv452[v0, v1]
|
||||
for k_2 in T.unroll(T.int64(8)):
|
||||
for ax0 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0 = T.axis.spatial(T.int64(4096), k_0 * T.int64(32) + k_1 * T.int64(8) + k_2)
|
||||
v_i1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax0)
|
||||
T.reads(lv452_local[v_i0 // T.int64(8), v_i1], lv453_local[v_i0 // T.int64(32), v_i1])
|
||||
T.writes(dequantize_intermediate_intermediate_local[v_i0, v_i1])
|
||||
dequantize_intermediate_intermediate_local[v_i0, v_i1] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv452_local[v_i0 // T.int64(8), v_i1], T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4))), T.uint32(15))) - T.float16(7)) * lv453_local[v_i0 // T.int64(32), v_i1]
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(1)):
|
||||
for ax3 in T.vectorized(T.int64(16)):
|
||||
with T.sblock("rms_norm130_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(T.int64(1), ax0)
|
||||
v1 = T.axis.spatial((seq_len + T.int64(15)) // T.int64(16), i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 + ax1)
|
||||
v2 = T.axis.spatial(T.int64(4096), k_0 * T.int64(32) + k_1 * T.int64(8) + k_2 + ax2)
|
||||
v3 = T.axis.spatial(T.int64(16), ax3)
|
||||
T.where(i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 < (seq_len + T.int64(15)) // T.int64(16))
|
||||
T.reads(rms_norm130_reindex_pad[v0, v1, v2, v3])
|
||||
T.writes(rms_norm130_reindex_pad_local[v0, v1, v2, v3])
|
||||
rms_norm130_reindex_pad_local[v0, v1, v2, v3] = rms_norm130_reindex_pad[v0, v1, v2, v3]
|
||||
for i0_i1_fused_2 in T.unroll(T.int64(16)):
|
||||
for i2_2 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("matmul_update"):
|
||||
v_i0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_i1 = T.axis.spatial((seq_len + T.int64(15)) // T.int64(16) * T.int64(16), i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + i0_i1_fused_2)
|
||||
v_i2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2)
|
||||
v_k = T.axis.reduce(T.int64(4096), k_0 * T.int64(32) + k_1 * T.int64(8) + k_2)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1) * T.int64(16) + i0_i1_fused_2 < (seq_len + T.int64(15)) // T.int64(16) * T.int64(16))
|
||||
T.reads(matmul_intermediate_pad_local[v_i0, v_i1, v_i2], rms_norm130_reindex_pad_local[v_i0, v_i1 // T.int64(16), v_k, v_i1 % T.int64(16)], dequantize_intermediate_intermediate_local[v_k, v_i2])
|
||||
T.writes(matmul_intermediate_pad_local[v_i0, v_i1, v_i2])
|
||||
matmul_intermediate_pad_local[v_i0, v_i1, v_i2] = matmul_intermediate_pad_local[v_i0, v_i1, v_i2] + rms_norm130_reindex_pad_local[v_i0, v_i1 // T.int64(16), v_k, v_i1 % T.int64(16)] * dequantize_intermediate_intermediate_local[v_k, v_i2]
|
||||
for ax0 in T.unroll(T.int64(16)):
|
||||
for ax1 in T.vectorized(T.int64(8)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_ax1 = T.axis.spatial(seq_len, i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0)
|
||||
v_ax2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1)
|
||||
T.where((i0_i1_fused_0 * T.int64(4) + i0_i1_fused_1 - (seq_len + T.int64(15)) // T.int64(16) < T.int64(0) or i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0 == T.int64(0)) and i0_i1_fused_0 * T.int64(64) + i0_i1_fused_1 * T.int64(16) + ax0 < seq_len)
|
||||
T.reads(matmul_intermediate_pad_local[v_ax0, v_ax1, v_ax2], transformer_h_0_attn_c_attn_bias3[v_ax2])
|
||||
T.writes(T_add_intermediate_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T_add_intermediate_intermediate[v_ax0, v_ax1, v_ax2] = matmul_intermediate_pad_local[v_ax0, v_ax1, v_ax2] + transformer_h_0_attn_c_attn_bias3[v_ax2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("opencl", host={"kind": "llvm", "mtriple": "aarch64-linux-android"}):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,994 @@
|
||||
# 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, unused-variable, invalid-name
|
||||
# ruff: noqa: E501, F841
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_matmul_tensorize():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(X: T.Buffer((256, 256), "float16"), W: T.Buffer((256, 256), "float16"), compute: T.Buffer((256, 256), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
for i, j, k in T.grid(256, 256, 256):
|
||||
with T.sblock("compute"):
|
||||
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
|
||||
T.reads(X[v_i, v_k], W[v_j, v_k])
|
||||
T.writes(compute[v_i, v_j])
|
||||
with T.init():
|
||||
compute[v_i, v_j] = T.float16(0)
|
||||
compute[v_i, v_j] = compute[v_i, v_j] + X[v_i, v_k] * W[v_j, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(X: T.Buffer((256, 256), "float16"), W: T.Buffer((256, 256), "float16"), compute: T.Buffer((256, 256), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
X_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
|
||||
W_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
|
||||
X_reindex_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.matrix_a")
|
||||
W_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.matrix_b")
|
||||
compute_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="shared.dyn")
|
||||
compute_reindex_shared_dyn_wmma_accumulator = T.sblock_alloc_buffer((1, 256, 256), "float16", scope="wmma.accumulator")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0_0_ax2_0_0_fused in T.thread_binding(2, thread="blockIdx.x"):
|
||||
for ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
|
||||
for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
|
||||
with T.sblock("compute_o_init"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
|
||||
T.reads()
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("compute_init_o"):
|
||||
v1_i_init_o = T.axis.spatial(1, 0)
|
||||
v2_i_init_o = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
C = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
|
||||
for ax3_0_0 in range(4, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
|
||||
for ax0_ax1_fused_0 in range(4):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("X_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
|
||||
v2 = T.axis.spatial(256, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
|
||||
T.reads(X[v1, v2])
|
||||
T.writes(X_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
X_reindex_shared_dyn[v0, v1, v2] = X[v1, v2]
|
||||
for ax0_ax1_fused_0 in range(4):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("W_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
|
||||
v2 = T.axis.spatial(256, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
|
||||
T.reads(W[v1, v2])
|
||||
T.writes(W_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
W_reindex_shared_dyn[v0, v1, v2] = W[v1, v2]
|
||||
for ax3_0_1 in range(4, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("X_reindex_shared.dyn_wmma.matrix_a_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
|
||||
T.reads(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
|
||||
T.reads(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
|
||||
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
|
||||
with T.sblock("compute_o_update"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
|
||||
v3_o = T.axis.reduce(16, ax3_0_0 * 4 + ax3_0_1)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("compute_o"):
|
||||
v1_i_o = T.axis.spatial(1, 0)
|
||||
v2_i_o = T.axis.spatial(1, 0)
|
||||
v3_i_o = T.axis.reduce(1, 0)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
B = T.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
C = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
|
||||
for ax0_0, ax1_0 in T.grid(2, 2):
|
||||
with T.sblock("compute_reindex_shared.dyn_wmma.accumulator_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
C = T.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="shared.dyn", offset_factor=16)
|
||||
T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
|
||||
for ax0_ax1_fused_0 in range(8):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(4):
|
||||
with T.sblock("compute_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
|
||||
v2 = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
|
||||
T.reads(compute_reindex_shared_dyn[v0, v1, v2])
|
||||
T.writes(compute[v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
|
||||
compute[v1, v2] = compute_reindex_shared_dyn[v0, v1, v2]
|
||||
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-rtx-2080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_tensorize_too_small():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_X: T.handle, W: T.Buffer((15, 256), "float16"), var_compute: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
m = T.int32()
|
||||
X = T.match_buffer(var_X, (m, 256), "float16")
|
||||
compute = T.match_buffer(var_compute, (m, 15))
|
||||
# with T.sblock("root"):
|
||||
for i, j, k in T.grid(m, 15, 256):
|
||||
with T.sblock("compute"):
|
||||
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
|
||||
T.reads(X[v_i, v_k], W[v_j, v_k])
|
||||
T.writes(compute[v_i, v_j])
|
||||
with T.init():
|
||||
compute[v_i, v_j] = T.float32(0)
|
||||
compute[v_i, v_j] = compute[v_i, v_j] + T.Cast("float32", X[v_i, v_k]) * T.Cast("float32", W[v_j, v_k])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_X: T.handle, W: T.Buffer((15, 256), "float16"), var_compute: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
m = T.int32()
|
||||
X = T.match_buffer(var_X, (m, 256), "float16")
|
||||
compute = T.match_buffer(var_compute, (m, 15))
|
||||
# with T.sblock("root"):
|
||||
compute_reindex_pad_local = T.sblock_alloc_buffer((1, (m + 31) // 32 * 32, 64), scope="local")
|
||||
X_reindex_pad_shared = T.sblock_alloc_buffer((1, (m + 31) // 32 * 32, 256), "float16", scope="shared")
|
||||
W_reindex_pad_shared = T.sblock_alloc_buffer((1, 64, 256), "float16", scope="shared")
|
||||
for ax0_ax2_0_fused in T.thread_binding(1, thread="blockIdx.y"):
|
||||
for ax1_0 in T.thread_binding((m + 31) // 32, thread="blockIdx.x"):
|
||||
for ax2_1 in T.thread_binding(1, thread="vthread.y"):
|
||||
for ax1_1 in T.thread_binding(1, thread="vthread.x"):
|
||||
for ax2_2 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_2 in T.thread_binding(8, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_3_init, ax2_3_0_init in T.grid(4, 2):
|
||||
for ax2_3_1_init in T.vectorized(2):
|
||||
with T.sblock("compute_init"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3_init)
|
||||
v2 = T.axis.spatial(64, ax2_1 * 64 + ax2_2 * 4 + ax2_3_0_init * 2 + ax2_3_1_init)
|
||||
T.reads()
|
||||
T.writes(compute_reindex_pad_local[0, v1, v2])
|
||||
compute_reindex_pad_local[0, v1, v2] = T.float32(0)
|
||||
for ax3_0 in range(16):
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(2):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("X_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(256, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 32 + ax0_ax1_ax2_fused_1 * 4 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 16)
|
||||
T.reads(X[v1, v2])
|
||||
T.writes(X_reindex_pad_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
X_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < m, X[v1, v2], T.float16(0))
|
||||
for ax0_ax1_ax2_fused_0 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in range(4):
|
||||
for ax0_ax1_ax2_fused_3 in T.vectorized(2):
|
||||
with T.sblock("W_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(64, (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) // 16)
|
||||
v2 = T.axis.spatial(256, ax3_0 * 16 + (ax0_ax1_ax2_fused_0 * 64 + ax0_ax1_ax2_fused_1 * 8 + ax0_ax1_ax2_fused_2 * 2 + ax0_ax1_ax2_fused_3) % 16)
|
||||
T.reads(W[v1, v2])
|
||||
T.writes(W_reindex_pad_shared[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 8, 2]]})
|
||||
W_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < 15, W[v1, v2], T.float16(0))
|
||||
for ax3_1, ax1_3, ax2_3_0 in T.grid(16, 4, 2):
|
||||
for ax2_3_1 in T.vectorized(2):
|
||||
with T.sblock("compute_update"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_1 * 32 + ax1_2 * 4 + ax1_3)
|
||||
v2 = T.axis.spatial(64, ax2_1 * 64 + ax2_2 * 4 + ax2_3_0 * 2 + ax2_3_1)
|
||||
v3 = T.axis.reduce(256, ax3_0 * 16 + ax3_1)
|
||||
T.reads(compute_reindex_pad_local[0, v1, v2], X_reindex_pad_shared[0, v1, v3], W_reindex_pad_shared[0, v2, v3])
|
||||
T.writes(compute_reindex_pad_local[0, v1, v2])
|
||||
compute_reindex_pad_local[0, v1, v2] = compute_reindex_pad_local[0, v1, v2] + T.Cast("float32", X_reindex_pad_shared[0, v1, v3]) * T.Cast("float32", W_reindex_pad_shared[0, v2, v3])
|
||||
for ax0, ax1, ax2_0 in T.grid(1, 4, 2):
|
||||
for ax2_1_1 in T.vectorized(2):
|
||||
with T.sblock("compute_reindex_pad_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial((m + 31) // 32 * 32, ax1_0 * 32 + ax1_2 * 4 + ax1)
|
||||
v2 = T.axis.spatial(64, ax2_2 * 4 + ax2_0 * 2 + ax2_1_1)
|
||||
T.where(ax1_0 * 32 + ax1_2 * 4 + ax1 < m and ax2_2 * 4 + ax2_0 * 2 + ax2_1_1 < 15)
|
||||
T.reads(compute_reindex_pad_local[v0, v1, v2])
|
||||
T.writes(compute[v1, v2])
|
||||
compute[v1, v2] = compute_reindex_pad_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-rtx-2080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_tensorize_epilogue():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(lv686: T.Buffer((T.int32(4096), T.int32(256)), "uint32"), lv687: T.Buffer((T.int32(4096), T.int32(64)), "float16"), p_lv42: T.handle, p_lv3: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int32()
|
||||
lv42 = T.match_buffer(p_lv42, (T.int32(1), n, T.int32(2048)), "float16")
|
||||
lv3 = T.match_buffer(p_lv3, (T.int32(1), n, T.int32(4096)), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (T.int32(1), n, T.int32(4096)), "float16")
|
||||
# with T.sblock("root"):
|
||||
p_output0_intermediate_1 = T.sblock_alloc_buffer((T.int32(4096), T.int32(2048)), "float16")
|
||||
var_NT_matmul_intermediate = T.sblock_alloc_buffer((T.int32(1), n, T.int32(4096)), "float16")
|
||||
var_T_divide_intermediate = T.sblock_alloc_buffer((T.int32(1), n, T.int32(4096)), "float16")
|
||||
for i, j in T.grid(T.int32(4096), T.int32(2048)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(lv686[v_i, v_j // T.int32(8)], lv687[v_i, v_j // T.int32(32)])
|
||||
T.writes(p_output0_intermediate_1[v_i, v_j])
|
||||
p_output0_intermediate_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv686[v_i, v_j // T.int32(8)], T.Cast("uint32", v_j % T.int32(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv687[v_i, v_j // T.int32(32)]
|
||||
for i0, i1, i2, k in T.grid(T.int32(1), n, T.int32(4096), T.int32(2048)):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(lv42[v_i0, v_i1, v_k], p_output0_intermediate_1[v_i2, v_k])
|
||||
T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
var_NT_matmul_intermediate[v_i0, v_i1, v_i2] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2] + lv42[v_i0, v_i1, v_k] * p_output0_intermediate_1[v_i2, v_k]
|
||||
for ax0, ax1, ax2 in T.grid(T.int32(1), n, T.int32(4096)):
|
||||
with T.sblock("T_divide"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(lv3[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
var_T_divide_intermediate[v_ax0, v_ax1, v_ax2] = lv3[v_ax0, v_ax1, v_ax2] * T.float16(0.5)
|
||||
for ax0, ax1, ax2 in T.grid(T.int32(1), n, T.int32(4096)):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(var_T_divide_intermediate[v_ax0, v_ax1, v_ax2], var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
|
||||
p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_T_divide_intermediate[v_ax0, v_ax1, v_ax2] + var_NT_matmul_intermediate[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(lv686: T.Buffer((4096, 256), "uint32"), lv687: T.Buffer((4096, 64), "float16"), p_lv42: T.handle, p_lv3: T.handle, p_output0: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int32()
|
||||
lv42 = T.match_buffer(p_lv42, (1, n, 2048), "float16")
|
||||
lv3 = T.match_buffer(p_lv3, (1, n, 4096), "float16")
|
||||
p_output0_intermediate = T.match_buffer(p_output0, (1, n, 4096), "float16")
|
||||
# with T.sblock("root"):
|
||||
lv42_reindex_pad_shared_dyn = T.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 2048), "float16", scope="shared.dyn")
|
||||
p_output0_intermediate_1_reindex_shared_dyn = T.sblock_alloc_buffer((1, 4096, 2048), "float16", scope="shared.dyn")
|
||||
lv42_reindex_pad_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 2048), "float16", scope="wmma.matrix_a")
|
||||
p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer((1, 4096, 2048), "float16", scope="wmma.matrix_b")
|
||||
var_NT_matmul_intermediate_reindex_pad_shared_dyn = T.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 4096), "float16", scope="shared.dyn")
|
||||
var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator = T.sblock_alloc_buffer((1, (n + 127) // 128 * 128, 4096), "float16", scope="wmma.accumulator")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0_0_ax2_0_0_fused in T.thread_binding((n + 127) // 128, thread="blockIdx.x"):
|
||||
for ax1_0_1_ax2_0_1_fused in T.thread_binding(32, thread="blockIdx.y"):
|
||||
for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
|
||||
with T.sblock("NT_matmul_o_init"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial((n + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
|
||||
T.reads()
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("NT_matmul_init_o"):
|
||||
v1_i_init_o = T.axis.spatial(1, 0)
|
||||
v2_i_init_o = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
C = T.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
|
||||
for ax3_0_0 in range(32, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
|
||||
for ax0_ax1_fused_0 in range(4):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("lv42_reindex_pad_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((n + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
|
||||
v2 = T.axis.spatial(2048, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
|
||||
T.reads(lv42[v0, v1, v2])
|
||||
T.writes(lv42_reindex_pad_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
lv42_reindex_pad_shared_dyn[v0, v1, v2] = T.if_then_else(v1 < n, lv42[v0, v1, v2], T.float16(0))
|
||||
for ax0_ax1_fused_0 in range(4):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("p_output0_intermediate_1_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 64)
|
||||
v2 = T.axis.spatial(2048, ax3_0_0 * 64 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 64)
|
||||
T.reads(lv686[v1, v2 // 8], lv687[v1, v2 // 32])
|
||||
T.writes(p_output0_intermediate_1_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 8]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
p_output0_intermediate_1_reindex_shared_dyn[v0, v1, v2] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv686[v1, v2 // 8], T.Cast("uint32", v2 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv687[v1, v2 // 32]
|
||||
for ax3_0_1 in range(4, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("lv42_reindex_pad_shared.dyn_wmma.matrix_a_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(8 * ((n + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(128, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
|
||||
T.reads(lv42_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(lv42_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(lv42_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("p_output0_intermediate_1_reindex_shared.dyn_wmma.matrix_b_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(128, ax3_0_0 * 4 + ax3_0_1 + ax1_0)
|
||||
T.reads(p_output0_intermediate_1_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(p_output0_intermediate_1_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
|
||||
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
|
||||
with T.sblock("NT_matmul_o_update"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial((n + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
|
||||
v3_o = T.axis.reduce(128, ax3_0_0 * 4 + ax3_0_1)
|
||||
T.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("NT_matmul_o"):
|
||||
v1_i_o = T.axis.spatial(1, 0)
|
||||
v2_i_o = T.axis.spatial(1, 0)
|
||||
v3_i_o = T.axis.reduce(1, 0)
|
||||
T.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(lv42_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
B = T.match_buffer(p_output0_intermediate_1_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "float16", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
C = T.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
|
||||
for ax0_0, ax1_0 in T.grid(2, 2):
|
||||
with T.sblock("var_NT_matmul_intermediate_reindex_pad_shared.dyn_wmma.accumulator_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(8 * ((n + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
|
||||
T.reads(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
C = T.match_buffer(var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "float16", strides=("C_s0", "C_s1"), scope="shared.dyn", offset_factor=16)
|
||||
T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("float16"), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
|
||||
for ax0_ax1_fused_0 in range(8):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(4):
|
||||
with T.sblock("var_NT_matmul_intermediate_reindex_pad_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((n + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
|
||||
v2 = T.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
|
||||
T.where(ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + ((ax0_ax1_fused_0 * 32 + ax0_ax1_fused_1) * 4 + ax0_ax1_fused_2) // 32 < n)
|
||||
T.reads(lv3[0, v1, v2], var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0, v1, v2])
|
||||
T.writes(p_output0_intermediate[0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
|
||||
p_output0_intermediate[0, v1, v2] = lv3[0, v1, v2] * T.float16(0.5) + var_NT_matmul_intermediate_reindex_pad_shared_dyn[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-rtx-2080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_int8_tensorize():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(X: T.Buffer((256, 256), "int8"), W: T.Buffer((256, 256), "int8"), compute: T.Buffer((256, 256), "int32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
for i, j, r in T.grid(256, 256, 256):
|
||||
with T.sblock("compute"):
|
||||
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, r])
|
||||
T.reads(X[v_i, v_k], W[v_j, v_k])
|
||||
T.writes(compute[v_i, v_j])
|
||||
with T.init():
|
||||
compute[v_i, v_j] = 0
|
||||
compute[v_i, v_j] = compute[v_i, v_j] + T.Cast("int32", X[v_i, v_k]) * T.Cast("int32", W[v_j, v_k])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(X: T.Buffer((256, 256), "int8"), W: T.Buffer((256, 256), "int8"), compute: T.Buffer((256, 256), "int32")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
X_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "int8", scope="shared.dyn")
|
||||
W_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "int8", scope="shared.dyn")
|
||||
X_reindex_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer((1, 256, 256), "int8", scope="wmma.matrix_a")
|
||||
W_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer((1, 256, 256), "int8", scope="wmma.matrix_b")
|
||||
compute_reindex_shared_dyn = T.sblock_alloc_buffer((1, 256, 256), "int32", scope="shared.dyn")
|
||||
compute_reindex_shared_dyn_wmma_accumulator = T.sblock_alloc_buffer((1, 256, 256), "int32", scope="wmma.accumulator")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0_0_ax2_0_0_fused in T.thread_binding(2, thread="blockIdx.x"):
|
||||
for ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
|
||||
for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
|
||||
with T.sblock("compute_o_init"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
|
||||
T.reads()
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("compute_init_o"):
|
||||
v1_i_init_o = T.axis.spatial(1, 0)
|
||||
v2_i_init_o = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
C = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
|
||||
for ax3_0_0 in T.serial(16, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
|
||||
for ax0_ax1_fused_0 in range(1):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("X_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
|
||||
v2 = T.axis.spatial(256, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
|
||||
T.reads(X[v1, v2])
|
||||
T.writes(X_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
X_reindex_shared_dyn[v0, v1, v2] = X[v1, v2]
|
||||
for ax0_ax1_fused_0 in range(1):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("W_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
|
||||
v2 = T.axis.spatial(256, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
|
||||
T.reads(W[v1, v2])
|
||||
T.writes(W_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
W_reindex_shared_dyn[v0, v1, v2] = W[v1, v2]
|
||||
for ax3_0_1 in T.serial(1, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("X_reindex_shared.dyn_wmma.matrix_a_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax3_0_0 + ax1_0)
|
||||
T.reads(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(X_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int8"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "row_major")
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("W_reindex_shared.dyn_wmma.matrix_b_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax3_0_0 + ax1_0)
|
||||
T.reads(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(W_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int8"), A.data, A.elem_offset, A.strides[0] * 16, 1), A.strides[0], "col_major")
|
||||
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
|
||||
with T.sblock("compute_o_update"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
|
||||
v3_o = T.axis.reduce(16, ax3_0_0 + ax3_0_1)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("compute_o"):
|
||||
v1_i_o = T.axis.spatial(1, 0)
|
||||
v2_i_o = T.axis.spatial(1, 0)
|
||||
v3_i_o = T.axis.reduce(1, 0)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(X_reindex_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
B = T.match_buffer(W_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
C = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A.data, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, B.data, B.elem_offset // B.strides[0] // 16 * (B.strides[0] // 16) + B.elem_offset % B.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
|
||||
for ax0_0, ax1_0 in T.grid(2, 2):
|
||||
with T.sblock("compute_reindex_shared.dyn_wmma.accumulator_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(16, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(16, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
|
||||
T.reads(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A = T.match_buffer(compute_reindex_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
C = T.match_buffer(compute_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="shared.dyn", offset_factor=16)
|
||||
T.tvm_store_matrix_sync(A.data, 16, 16, 16, A.elem_offset // A.strides[0] // 16 * (A.strides[0] // 16) + A.elem_offset % A.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int32"), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
|
||||
for ax0_ax1_fused_0 in range(8):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(4):
|
||||
with T.sblock("compute_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(256, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
|
||||
v2 = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
|
||||
T.reads(compute_reindex_shared_dyn[v0, v1, v2])
|
||||
T.writes(compute[v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
|
||||
compute[v1, v2] = compute_reindex_shared_dyn[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-rtx-2080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_int8_tensorize_3d2d_dyn():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(var_A: T.handle, B: T.Buffer((4096, 22016), "int8"), var_matmul: T.handle):
|
||||
T.func_attr({"op_pattern": 4, "tirx.noalias": True})
|
||||
m = T.int32()
|
||||
A = T.match_buffer(var_A, (1, m, 22016), "int8")
|
||||
matmul_1 = T.match_buffer(var_matmul, (1, m, 4096), "int32")
|
||||
# with T.sblock("root"):
|
||||
for i0, i1, i2, k in T.grid(1, m, 4096, 22016):
|
||||
with T.sblock("matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.reads(A[v_i0, v_i1, v_k], B[v_i2, v_k])
|
||||
T.writes(matmul_1[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
matmul_1[v_i0, v_i1, v_i2] = 0
|
||||
matmul_1[v_i0, v_i1, v_i2] = matmul_1[v_i0, v_i1, v_i2] + T.Cast("int32", A[v_i0, v_i1, v_k]) * T.Cast("int32", B[v_i2, v_k])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_A: T.handle, B: T.Buffer((4096, 22016), "int8"), var_matmul: T.handle):
|
||||
T.func_attr({"op_pattern": 4, "tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
m = T.int32()
|
||||
A = T.match_buffer(var_A, (1, m, 22016), "int8")
|
||||
matmul_1 = T.match_buffer(var_matmul, (1, m, 4096), "int32")
|
||||
# with T.sblock("root"):
|
||||
A_reindex_pad_shared_dyn = T.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 22016), "int8", scope="shared.dyn")
|
||||
B_reindex_shared_dyn = T.sblock_alloc_buffer((1, 4096, 22016), "int8", scope="shared.dyn")
|
||||
A_reindex_pad_shared_dyn_wmma_matrix_a = T.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 22016), "int8", scope="wmma.matrix_a")
|
||||
B_reindex_shared_dyn_wmma_matrix_b = T.sblock_alloc_buffer((1, 4096, 22016), "int8", scope="wmma.matrix_b")
|
||||
matmul_1_reindex_pad_shared_dyn = T.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 4096), "int32", scope="shared.dyn")
|
||||
matmul_1_reindex_pad_shared_dyn_wmma_accumulator = T.sblock_alloc_buffer((1, (m + 127) // 128 * 128, 4096), "int32", scope="wmma.accumulator")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0_0_ax2_0_0_fused in T.thread_binding((m + 127) // 128, thread="blockIdx.x"):
|
||||
for ax1_0_1_ax2_0_1_fused in T.thread_binding(32, thread="blockIdx.y"):
|
||||
for ax2_0_2_ax1_0_2_fused in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax1_0_3_init, ax2_0_3_init in T.grid(2, 2):
|
||||
with T.sblock("matmul_o_init"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial((m + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3_init)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3_init)
|
||||
T.reads()
|
||||
T.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("matmul_init_o"):
|
||||
v1_i_init_o = T.axis.spatial(1, 0)
|
||||
v2_i_init_o = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
C = T.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_fill_fragment(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.float32(0))
|
||||
for ax3_0_0 in T.serial(1376, annotations={"software_pipeline_order": [0, 3, 1, 4, 5, 2, 6], "software_pipeline_stage": [0, 0, 0, 0, 0, 1, 1]}):
|
||||
for ax0_ax1_fused_0 in range(1):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("A_reindex_pad_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
|
||||
v2 = T.axis.spatial(22016, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
|
||||
T.reads(A[v0, v1, v2])
|
||||
T.writes(A_reindex_pad_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
A_reindex_pad_shared_dyn[v0, v1, v2] = T.if_then_else(v1 < m, A[v0, v1, v2], T.int8(0))
|
||||
for ax0_ax1_fused_0 in range(1):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("B_reindex_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) // 16)
|
||||
v2 = T.axis.spatial(22016, ax3_0_0 * 16 + (ax0_ax1_fused_0 * 2048 + ax0_ax1_fused_1 * 128 + ax0_ax1_fused_2 * 4 + ax0_ax1_fused_3) % 16)
|
||||
T.reads(B[v1, v2])
|
||||
T.writes(B_reindex_shared_dyn[v0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 32, 16]], "double_buffer_scope": 0, "tirx.manifest_shared_memory_local_stage": 1})
|
||||
B_reindex_shared_dyn[v0, v1, v2] = B[v1, v2]
|
||||
for ax3_0_1 in T.serial(1, annotations={"software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, 1]}):
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("A_reindex_pad_shared.dyn_wmma.matrix_a_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(8 * ((m + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(1376, ax3_0_0 + ax1_0)
|
||||
T.reads(A_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(A_reindex_pad_shared_dyn_wmma_matrix_a[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("C_s0", "C_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int8"), A_1.data, A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "row_major")
|
||||
for ax0_0 in T.unroll(2):
|
||||
for ax1_0 in T.unroll(1):
|
||||
with T.sblock("B_reindex_shared.dyn_wmma.matrix_b_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(1376, ax3_0_0 + ax1_0)
|
||||
T.reads(B_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A_1 = T.match_buffer(B_reindex_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="shared.dyn", offset_factor=16)
|
||||
C = T.match_buffer(B_reindex_shared_dyn_wmma_matrix_b[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int8", strides=("C_s0", "C_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
T.tvm_load_matrix_sync(C.data, 16, 16, 16, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int8"), A_1.data, A_1.elem_offset, A_1.strides[0] * 16, 1), A_1.strides[0], "col_major")
|
||||
for ax1_0_3, ax2_0_3 in T.grid(2, 2):
|
||||
with T.sblock("matmul_o_update"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial((m + 127) // 128 * 8, ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax1_0_3)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax2_0_3)
|
||||
v3_o = T.axis.reduce(1376, ax3_0_0 + ax3_0_1)
|
||||
T.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
with T.sblock("matmul_o"):
|
||||
v1_i_o = T.axis.spatial(1, 0)
|
||||
v2_i_o = T.axis.spatial(1, 0)
|
||||
v3_i_o = T.axis.reduce(1, 0)
|
||||
T.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16])
|
||||
T.writes(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared_dyn_wmma_matrix_a[0, v1_o * 16:v1_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=("A_s0", "A_s1"), scope="wmma.matrix_a", offset_factor=16)
|
||||
B_1 = T.match_buffer(B_reindex_shared_dyn_wmma_matrix_b[0, v2_o * 16:v2_o * 16 + 16, v3_o * 16:v3_o * 16 + 16], (16, 16), "int8", strides=("B_s0", "B_s1"), scope="wmma.matrix_b", offset_factor=16)
|
||||
C = T.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[0, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
T.tvm_mma_sync(C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16, A_1.data, A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) + A_1.elem_offset % A_1.strides[0] // 16, B_1.data, B_1.elem_offset // B_1.strides[0] // 16 * (B_1.strides[0] // 16) + B_1.elem_offset % B_1.strides[0] // 16, C.data, C.elem_offset // C.strides[0] // 16 * (C.strides[0] // 16) + C.elem_offset % C.strides[0] // 16)
|
||||
for ax0_0, ax1_0 in T.grid(2, 2):
|
||||
with T.sblock("matmul_1_reindex_pad_shared.dyn_wmma.accumulator_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(8 * ((m + 127) // 128), ax1_0_0_ax2_0_0_fused * 8 + ax2_0_2_ax1_0_2_fused % 4 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(256, ax1_0_1_ax2_0_1_fused * 8 + ax2_0_2_ax1_0_2_fused // 4 * 2 + ax1_0)
|
||||
T.reads(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
T.writes(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16])
|
||||
A_1 = T.match_buffer(matmul_1_reindex_pad_shared_dyn_wmma_accumulator[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("A_s0", "A_s1"), scope="wmma.accumulator", offset_factor=16)
|
||||
C = T.match_buffer(matmul_1_reindex_pad_shared_dyn[v0_o, v1_o * 16:v1_o * 16 + 16, v2_o * 16:v2_o * 16 + 16], (16, 16), "int32", strides=("C_s0", "C_s1"), scope="shared.dyn", offset_factor=16)
|
||||
T.tvm_store_matrix_sync(A_1.data, 16, 16, 16, A_1.elem_offset // A_1.strides[0] // 16 * (A_1.strides[0] // 16) + A_1.elem_offset % A_1.strides[0] // 16, T.tvm_access_ptr(T.type_annotation("int32"), C.data, C.elem_offset, C.strides[0] * 16, 2), C.strides[0], "row_major")
|
||||
for ax0_ax1_fused_0 in range(8):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(4):
|
||||
with T.sblock("matmul_1_reindex_pad_shared.dyn"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial((m + 127) // 128 * 128, ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 32)
|
||||
v2 = T.axis.spatial(4096, ax1_0_1_ax2_0_1_fused * 128 + ax2_0_2_ax1_0_2_fused // 4 * 32 + (ax0_ax1_fused_0 * 128 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 32)
|
||||
T.where(ax1_0_0_ax2_0_0_fused * 128 + ax2_0_2_ax1_0_2_fused % 4 * 32 + ((ax0_ax1_fused_0 * 32 + ax0_ax1_fused_1) * 4 + ax0_ax1_fused_2) // 32 < m)
|
||||
T.reads(matmul_1_reindex_pad_shared_dyn[v0, v1, v2])
|
||||
T.writes(matmul_1[0, v1, v2])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 1, 16, 4]]})
|
||||
matmul_1[0, v1, v2] = matmul_1_reindex_pad_shared_dyn[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("nvidia/geforce-rtx-2080-ti"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_metal():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
var_A: T.handle,
|
||||
B: T.Buffer((28672, 4096), "float16"),
|
||||
var_C: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 28672), "float16")
|
||||
for i0, i1, i2, k in T.grid(batch_size, 1, 28672, 4096):
|
||||
with T.sblock("C"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
T.writes(C[v_i0, v_i1, v_i2])
|
||||
with T.init():
|
||||
C[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
C[v_i0, v_i1, v_i2] += A[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(var_A: T.handle, B: T.Buffer((28672, 4096), "float16"), var_C: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 28672), "float16")
|
||||
# with T.sblock("root"):
|
||||
A_reindex_pad_shared = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 4096), "float16", scope="shared")
|
||||
B_reindex_shared = T.sblock_alloc_buffer((1, 28672, 4096), "float16", scope="shared")
|
||||
A_reindex_pad_shared_metal_simdgroup = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 4096), "float16", scope="metal.simdgroup")
|
||||
B_reindex_shared_metal_simdgroup = T.sblock_alloc_buffer((1, 4096, 28672), "float16", scope="metal.simdgroup")
|
||||
C_reindex_pad_metal_simdgroup = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 28672), "float16", scope="metal.simdgroup")
|
||||
C_reindex_pad_shared = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 28672), "float16", scope="shared")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0 in T.thread_binding((batch_size + 15) // 16, thread="blockIdx.x"):
|
||||
for ax2_0 in T.thread_binding(448, thread="blockIdx.y"):
|
||||
for ax1_1 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax2_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_2_init, ax2_2_init, ax1_3_init_0, ax2_3_init_0 in T.grid(2, 2, 1, 1):
|
||||
with T.sblock("C_init_o"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_1 * 2 + ax1_2_init + ax1_3_init_0)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_2_init + ax2_3_init_0)
|
||||
T.reads()
|
||||
T.writes(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.make_filled_simdgroup_matrix(A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, T.float32(0), 8, 8)
|
||||
for ax3_0 in range(128):
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 1):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("A_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial((batch_size + 15) // 16 * 16, ax1_0 * 16 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 32)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 32 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 32)
|
||||
T.reads(A[v1, 0, v2])
|
||||
T.writes(A_reindex_pad_shared[v0, v1, v2])
|
||||
A_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < batch_size, A[v1, 0, v2], T.float16(0))
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 4):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("B_reindex_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial(28672, ax2_0 * 64 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 32)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 32 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 32)
|
||||
T.reads(B[v1, v2])
|
||||
T.writes(B_reindex_shared[v0, v1, v2])
|
||||
B_reindex_shared[v0, v1, v2] = B[v1, v2]
|
||||
for ax3_1 in range(4):
|
||||
for ax0_0, ax1_0_1 in T.grid(2, 1):
|
||||
with T.sblock("A_reindex_pad_shared_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(512, ax3_0 * 4 + ax3_1 + ax1_0_1)
|
||||
T.reads(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=1)
|
||||
C_1 = T.match_buffer(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_load(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), A_1.data, A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
|
||||
for ax0_0, ax1_0_1 in T.grid(2, 1):
|
||||
with T.sblock("B_reindex_shared_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(512, ax3_0 * 4 + ax3_1 + ax1_0_1)
|
||||
T.reads(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o * 8:v1_o * 8 + 8])
|
||||
A_1 = T.match_buffer(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=1)
|
||||
C_1 = T.match_buffer(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o * 8:v1_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_load(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), A_1.data, A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
|
||||
for ax1_2, ax2_2 in T.grid(2, 2):
|
||||
with T.sblock("C_update_o"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_1 * 2 + ax1_2)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_2)
|
||||
v3_o = T.axis.reduce(512, ax3_0 * 4 + ax3_1)
|
||||
T.reads(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], A_reindex_pad_shared_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v3_o * 8:v3_o * 8 + 8], B_reindex_shared_metal_simdgroup[0, v3_o * 8:v3_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v3_o * 8:v3_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
B_1 = T.match_buffer(B_reindex_shared_metal_simdgroup[0, v3_o * 8:v3_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("B_s0", "B_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
C_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_multiply_accumulate(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, B_1.data, B_1.elem_offset // B_1.strides[0] // 8 * (B_1.strides[0] // 8) + B_1.elem_offset % B_1.strides[0] // 8, C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8)
|
||||
for ax0_1, ax1_0_1, ax2_0_1 in T.grid(1, 2, 2):
|
||||
with T.sblock("C_reindex_pad_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, ax0_1)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_0_1)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_0_1)
|
||||
T.reads(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
C_1 = T.match_buffer(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="shared", offset_factor=1)
|
||||
T.metal.simdgroup_store(A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), C_1.data, C_1.elem_offset, C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 2):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("C_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial((batch_size + 15) // 16 * 16, ax1_0 * 16 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 64)
|
||||
v2 = T.axis.spatial(28672, ax2_0 * 64 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 64)
|
||||
T.where(ax1_0 * 16 + (((ax1_ax2_fused_0 * 4 + ax1_ax2_fused_1 + ax1_ax2_fused_2) * 32 + ax1_ax2_fused_3) * 4 + ax1_ax2_fused_4) // 64 < batch_size)
|
||||
T.reads(C_reindex_pad_shared[v0, v1, v2])
|
||||
T.writes(C[v1, 0, v2])
|
||||
C[v1, 0, v2] = C_reindex_pad_shared[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_matmul_metal_int4_quant():
|
||||
# fmt: off
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
B0: T.Buffer((28672, 512), "uint32"),
|
||||
B1: T.Buffer((28672, 128), "float16"),
|
||||
var_A: T.handle,
|
||||
var_C: T.handle
|
||||
):
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 28672), "float16")
|
||||
compute = T.sblock_alloc_buffer((28672, 4096), "float16")
|
||||
B = T.sblock_alloc_buffer((28672, 4096), "float16")
|
||||
for i0, i1 in T.grid(28672, 4096):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
compute[v_i0, v_i1] = T.Cast("float16", T.bitwise_and(T.shift_right(B0[v_i0, v_i1 // 8], T.Cast("uint32", v_i1 % 8 * 4)), T.uint32(15)))
|
||||
for i0, i1 in T.grid(28672, 4096):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
B[v_i0, v_i1] = (compute[v_i0, v_i1] - T.float16(7)) * B1[v_i0, v_i1 // 32]
|
||||
for i0, i1, i2, k in T.grid(batch_size, 1, 28672, 4096):
|
||||
with T.sblock("NT_matmul"):
|
||||
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
||||
with T.init():
|
||||
C[v_i0, v_i1, v_i2] = T.float16(0)
|
||||
C[v_i0, v_i1, v_i2] = C[v_i0, v_i1, v_i2] + A[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(B0: T.Buffer((28672, 512), "uint32"), B1: T.Buffer((28672, 128), "float16"), var_A: T.handle, var_C: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
batch_size = T.int32()
|
||||
A = T.match_buffer(var_A, (batch_size, 1, 4096), "float16")
|
||||
C = T.match_buffer(var_C, (batch_size, 1, 28672), "float16")
|
||||
# with T.sblock("root"):
|
||||
A_reindex_pad_shared = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 4096), "float16", scope="shared")
|
||||
B_reindex_shared = T.sblock_alloc_buffer((1, 28672, 4096), "float16", scope="shared")
|
||||
A_reindex_pad_shared_metal_simdgroup = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 4096), "float16", scope="metal.simdgroup")
|
||||
B_reindex_shared_metal_simdgroup = T.sblock_alloc_buffer((1, 4096, 28672), "float16", scope="metal.simdgroup")
|
||||
C_reindex_pad_metal_simdgroup = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 28672), "float16", scope="metal.simdgroup")
|
||||
C_reindex_pad_shared = T.sblock_alloc_buffer((1, (batch_size + 15) // 16 * 16, 28672), "float16", scope="shared")
|
||||
for ax0 in T.thread_binding(1, thread="blockIdx.z"):
|
||||
for ax1_0 in T.thread_binding((batch_size + 15) // 16, thread="blockIdx.x"):
|
||||
for ax2_0 in T.thread_binding(448, thread="blockIdx.y"):
|
||||
for ax1_1 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax2_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_2_init, ax2_2_init, ax1_3_init_0, ax2_3_init_0 in T.grid(2, 2, 1, 1):
|
||||
with T.sblock("NT_matmul_init_o"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_1 * 2 + ax1_2_init + ax1_3_init_0)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_2_init + ax2_3_init_0)
|
||||
T.reads()
|
||||
T.writes(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.make_filled_simdgroup_matrix(A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, T.float32(0), 8, 8)
|
||||
for ax3_0 in range(128):
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 1):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("A_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial((batch_size + 15) // 16 * 16, ax1_0 * 16 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 32)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 32 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 32)
|
||||
T.reads(A[v1, 0, v2])
|
||||
T.writes(A_reindex_pad_shared[v0, v1, v2])
|
||||
A_reindex_pad_shared[v0, v1, v2] = T.if_then_else(v1 < batch_size, A[v1, 0, v2], T.float16(0))
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 4):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("B_reindex_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial(28672, ax2_0 * 64 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 32)
|
||||
v2 = T.axis.spatial(4096, ax3_0 * 32 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 32)
|
||||
T.reads(B0[v1, v2 // 8], B1[v1, v2 // 32])
|
||||
T.writes(B_reindex_shared[v0, v1, v2])
|
||||
B_reindex_shared[v0, v1, v2] = (T.Cast("float16", T.bitwise_and(T.shift_right(B0[v1, v2 // 8], T.Cast("uint32", v2 % 8 * 4)), T.uint32(15))) - T.float16(7)) * B1[v1, v2 // 32]
|
||||
for ax3_1 in range(4):
|
||||
for ax0_0, ax1_0_1 in T.grid(2, 1):
|
||||
with T.sblock("A_reindex_pad_shared_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(512, ax3_0 * 4 + ax3_1 + ax1_0_1)
|
||||
T.reads(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=1)
|
||||
C_1 = T.match_buffer(A_reindex_pad_shared_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_load(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), A_1.data, A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(False))
|
||||
for ax0_0, ax1_0_1 in T.grid(2, 1):
|
||||
with T.sblock("B_reindex_shared_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax0_0)
|
||||
v2_o = T.axis.spatial(512, ax3_0 * 4 + ax3_1 + ax1_0_1)
|
||||
T.reads(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o * 8:v1_o * 8 + 8])
|
||||
A_1 = T.match_buffer(B_reindex_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="shared", offset_factor=1)
|
||||
C_1 = T.match_buffer(B_reindex_shared_metal_simdgroup[v0_o, v2_o * 8:v2_o * 8 + 8, v1_o * 8:v1_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_load(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), A_1.data, A_1.elem_offset, A_1.strides[0] * 8, 1), A_1.strides[0], 8, 8, T.bool(True))
|
||||
for ax1_2, ax2_2 in T.grid(2, 2):
|
||||
with T.sblock("NT_matmul_update_o"):
|
||||
v0_o = T.axis.spatial(1, ax0)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_1 * 2 + ax1_2)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_2)
|
||||
v3_o = T.axis.reduce(512, ax3_0 * 4 + ax3_1)
|
||||
T.reads(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], A_reindex_pad_shared_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v3_o * 8:v3_o * 8 + 8], B_reindex_shared_metal_simdgroup[0, v3_o * 8:v3_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(A_reindex_pad_shared_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v3_o * 8:v3_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
B = T.match_buffer(B_reindex_shared_metal_simdgroup[0, v3_o * 8:v3_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("B_s0", "B_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
C_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[0, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
T.metal.simdgroup_multiply_accumulate(C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8, A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, B.data, B.elem_offset // B.strides[0] // 8 * (B.strides[0] // 8) + B.elem_offset % B.strides[0] // 8, C_1.data, C_1.elem_offset // C_1.strides[0] // 8 * (C_1.strides[0] // 8) + C_1.elem_offset % C_1.strides[0] // 8)
|
||||
for ax0_1, ax1_0_1, ax2_0_1 in T.grid(1, 2, 2):
|
||||
with T.sblock("C_reindex_pad_metal.simdgroup_o"):
|
||||
v0_o = T.axis.spatial(1, ax0_1)
|
||||
v1_o = T.axis.spatial(2 * ((batch_size + 15) // 16), ax1_0 * 2 + ax1_0_1)
|
||||
v2_o = T.axis.spatial(3584, ax2_0 * 8 + ax2_1 * 2 + ax2_0_1)
|
||||
T.reads(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
T.writes(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8])
|
||||
A_1 = T.match_buffer(C_reindex_pad_metal_simdgroup[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("A_s0", "A_s1"), scope="metal.simdgroup", offset_factor=1)
|
||||
C_1 = T.match_buffer(C_reindex_pad_shared[v0_o, v1_o * 8:v1_o * 8 + 8, v2_o * 8:v2_o * 8 + 8], (8, 8), "float16", strides=("C_s0", "C_s1"), scope="shared", offset_factor=1)
|
||||
T.metal.simdgroup_store(A_1.data, A_1.elem_offset // A_1.strides[0] // 8 * (A_1.strides[0] // 8) + A_1.elem_offset % A_1.strides[0] // 8, T.tvm_access_ptr(T.type_annotation("float16"), C_1.data, C_1.elem_offset, C_1.strides[0] * 8, 2), C_1.strides[0], 8, 8, T.bool(False))
|
||||
for ax0_1, ax1_ax2_fused_0 in T.grid(1, 2):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(4, thread="threadIdx.z"):
|
||||
for ax1_ax2_fused_2 in T.thread_binding(1, thread="threadIdx.y"):
|
||||
for ax1_ax2_fused_3 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax1_ax2_fused_4 in T.vectorized(4):
|
||||
with T.sblock("C_reindex_pad_shared"):
|
||||
v0 = T.axis.spatial(1, ax0_1)
|
||||
v1 = T.axis.spatial((batch_size + 15) // 16 * 16, ax1_0 * 16 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) // 64)
|
||||
v2 = T.axis.spatial(28672, ax2_0 * 64 + (ax1_ax2_fused_0 * 512 + ax1_ax2_fused_1 * 128 + ax1_ax2_fused_2 * 128 + ax1_ax2_fused_3 * 4 + ax1_ax2_fused_4) % 64)
|
||||
T.where(ax1_0 * 16 + (((ax1_ax2_fused_0 * 4 + ax1_ax2_fused_1 + ax1_ax2_fused_2) * 32 + ax1_ax2_fused_3) * 4 + ax1_ax2_fused_4) // 64 < batch_size)
|
||||
T.reads(C_reindex_pad_shared[v0, v1, v2])
|
||||
T.writes(C[v1, 0, v2])
|
||||
C[v1, 0, v2] = C_reindex_pad_shared[v0, v1, v2]
|
||||
# fmt: on
|
||||
|
||||
mod = tvm.IRModule({"main": before})
|
||||
with Target("metal"):
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,287 @@
|
||||
# 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
|
||||
import tvm.testing
|
||||
from tvm.ir import IRModule, assert_structural_equal
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _check(mod_before: IRModule, mod_after: IRModule):
|
||||
target = Target("nvidia/geforce-rtx-3090-ti")
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.RMSNorm(),
|
||||
)(mod_before)
|
||||
assert_structural_equal(mod, mod_after)
|
||||
|
||||
|
||||
def test_rms_norm_with_casting():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_data: T.handle, weight: T.Buffer((4096,), "float16"), var_T_cast: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int32()
|
||||
data = T.match_buffer(var_data, (1, n, 4096), "float16")
|
||||
T_cast = T.match_buffer(var_T_cast, (1, n, 4096), "float16")
|
||||
# with T.sblock("root"):
|
||||
T_cast_1 = T.sblock_alloc_buffer((1, n, 4096))
|
||||
T_multiply = T.sblock_alloc_buffer((1, n, 4096))
|
||||
T_multiply_red = T.sblock_alloc_buffer((1, n))
|
||||
rsqrt = T.sblock_alloc_buffer((1, n))
|
||||
T_cast_2 = T.sblock_alloc_buffer((4096,))
|
||||
T_rms_norm = T.sblock_alloc_buffer((1, n, 4096))
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_cast"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(data[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_cast_1[v_ax0, v_ax1, v_ax2])
|
||||
T_cast_1[v_ax0, v_ax1, v_ax2] = T.Cast("float32", data[v_ax0, v_ax1, v_ax2])
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_multiply"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(T_cast_1[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_multiply[v_ax0, v_ax1, v_ax2])
|
||||
T_multiply[v_ax0, v_ax1, v_ax2] = T_cast_1[v_ax0, v_ax1, v_ax2] * T_cast_1[v_ax0, v_ax1, v_ax2]
|
||||
for ax0, ax1, k2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_multiply_red"):
|
||||
v_ax0, v_ax1, v_k2 = T.axis.remap("SSR", [ax0, ax1, k2])
|
||||
T.reads(T_multiply[v_ax0, v_ax1, v_k2])
|
||||
T.writes(T_multiply_red[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
T_multiply_red[v_ax0, v_ax1] = T.float32(0)
|
||||
T_multiply_red[v_ax0, v_ax1] = T_multiply_red[v_ax0, v_ax1] + T_multiply[v_ax0, v_ax1, v_k2]
|
||||
for ax0, ax1 in T.grid(1, n):
|
||||
with T.sblock("rsqrt"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(T_multiply_red[v_ax0, v_ax1])
|
||||
T.writes(rsqrt[v_ax0, v_ax1])
|
||||
rsqrt[v_ax0, v_ax1] = T.rsqrt(T_multiply_red[v_ax0, v_ax1] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))
|
||||
for ax0 in range(4096):
|
||||
with T.sblock("T_cast_1"):
|
||||
v_ax0 = T.axis.spatial(4096, ax0)
|
||||
T.reads(weight[v_ax0])
|
||||
T.writes(T_cast_2[v_ax0])
|
||||
T_cast_2[v_ax0] = T.Cast("float32", weight[v_ax0])
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_rms_norm"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(rsqrt[v_ax0, v_ax1], T_cast_1[v_ax0, v_ax1, v_ax2], T_cast_2[v_ax2])
|
||||
T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2])
|
||||
T_rms_norm[v_ax0, v_ax1, v_ax2] = rsqrt[v_ax0, v_ax1] * T_cast_1[v_ax0, v_ax1, v_ax2] * T_cast_2[v_ax2]
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_cast_2"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(T_rms_norm[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_cast[v_ax0, v_ax1, v_ax2])
|
||||
T_cast[v_ax0, v_ax1, v_ax2] = T.Cast("float16", T_rms_norm[v_ax0, v_ax1, v_ax2])
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_data: T.handle, weight: T.Buffer((4096,), "float16"), var_T_cast: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int32()
|
||||
data = T.match_buffer(var_data, (1, n, 4096), "float16")
|
||||
T_cast = T.match_buffer(var_T_cast, (1, n, 4096), "float16")
|
||||
# with T.sblock("root"):
|
||||
T_multiply_local = T.sblock_alloc_buffer((1, n, 4096), scope="local")
|
||||
T_multiply_red_local = T.sblock_alloc_buffer((1, n), scope="local")
|
||||
rsqrt_shared = T.sblock_alloc_buffer((1, n), scope="shared")
|
||||
T_rms_norm_local = T.sblock_alloc_buffer((1, n, 4096), scope="local")
|
||||
data_local = T.sblock_alloc_buffer((1, n, 4096), "float16", scope="local")
|
||||
for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x"):
|
||||
for ax2_0 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
for ax2_1 in range(1):
|
||||
for ax2_2 in T.vectorized(8):
|
||||
with T.sblock("data_local"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v2 = T.axis.spatial(4096, ax2_0 * 8 + ax2_1 * 8 + ax2_2)
|
||||
T.reads(data[v0, v1, v2])
|
||||
T.writes(data_local[v0, v1, v2])
|
||||
data_local[v0, v1, v2] = data[v0, v1, v2]
|
||||
for ax0 in range(8):
|
||||
with T.sblock("T_multiply"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_ax2 = T.axis.spatial(4096, ax2_0 * 8 + ax0)
|
||||
T.reads(data_local[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_multiply_local[v_ax0, v_ax1, v_ax2])
|
||||
T_multiply_local[v_ax0, v_ax1, v_ax2] = T.Cast("float32", data_local[v_ax0, v_ax1, v_ax2]) * T.Cast("float32", data_local[v_ax0, v_ax1, v_ax2])
|
||||
for ax0 in range(8):
|
||||
with T.sblock("T_multiply_red"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_k2 = T.axis.reduce(4096, ax2_0 * 8 + ax0)
|
||||
T.reads(T_multiply_local[v_ax0, v_ax1, v_k2])
|
||||
T.writes(T_multiply_red_local[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
T_multiply_red_local[v_ax0, v_ax1] = T.float32(0)
|
||||
T_multiply_red_local[v_ax0, v_ax1] = T_multiply_red_local[v_ax0, v_ax1] + T_multiply_local[v_ax0, v_ax1, v_k2]
|
||||
with T.sblock("rsqrt"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
T.reads(T_multiply_red_local[v_ax0, v_ax1])
|
||||
T.writes(rsqrt_shared[v_ax0, v_ax1])
|
||||
rsqrt_shared[v_ax0, v_ax1] = T.rsqrt(T_multiply_red_local[v_ax0, v_ax1] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))
|
||||
for ax0_0 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
for ax0_1, ax0_2 in T.grid(1, 8):
|
||||
with T.sblock("T_rms_norm"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_ax2 = T.axis.spatial(4096, ax0_0 * 8 + ax0_1 * 8 + ax0_2)
|
||||
T.reads(rsqrt_shared[v_ax0, v_ax1], data_local[v_ax0, v_ax1, v_ax2], weight[v_ax2])
|
||||
T.writes(T_rms_norm_local[v_ax0, v_ax1, v_ax2])
|
||||
T_rms_norm_local[v_ax0, v_ax1, v_ax2] = rsqrt_shared[v_ax0, v_ax1] * T.Cast("float32", data_local[v_ax0, v_ax1, v_ax2]) * T.Cast("float32", weight[v_ax2])
|
||||
for ax0 in T.vectorized(8):
|
||||
with T.sblock("T_cast_local"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v2 = T.axis.spatial(4096, ax0_0 * 8 + ax0)
|
||||
T.reads(T_rms_norm_local[v0, v1, v2])
|
||||
T.writes(T_cast[v0, v1, v2])
|
||||
T_cast[v0, v1, v2] = T.Cast("float16", T_rms_norm_local[v0, v1, v2])
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_rms_norm_without_casting():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_data: T.handle, weight: T.Buffer((4096,), "float32"), var_T_cast: T.handle):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
n = T.int32()
|
||||
data = T.match_buffer(var_data, (1, n, 4096))
|
||||
T_cast = T.match_buffer(var_T_cast, (1, n, 4096))
|
||||
# with T.sblock("root"):
|
||||
T_multiply = T.sblock_alloc_buffer((1, n, 4096))
|
||||
T_multiply_red = T.sblock_alloc_buffer((1, n))
|
||||
rsqrt = T.sblock_alloc_buffer((1, n))
|
||||
T_rms_norm = T.sblock_alloc_buffer((1, n, 4096))
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_multiply"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(data[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_multiply[v_ax0, v_ax1, v_ax2])
|
||||
T_multiply[v_ax0, v_ax1, v_ax2] = data[v_ax0, v_ax1, v_ax2] * data[v_ax0, v_ax1, v_ax2]
|
||||
for ax0, ax1, k2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_multiply_red"):
|
||||
v_ax0, v_ax1, v_k2 = T.axis.remap("SSR", [ax0, ax1, k2])
|
||||
T.reads(T_multiply[v_ax0, v_ax1, v_k2])
|
||||
T.writes(T_multiply_red[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
T_multiply_red[v_ax0, v_ax1] = T.float32(0)
|
||||
T_multiply_red[v_ax0, v_ax1] = T_multiply_red[v_ax0, v_ax1] + T_multiply[v_ax0, v_ax1, v_k2]
|
||||
for ax0, ax1 in T.grid(1, n):
|
||||
with T.sblock("rsqrt"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(T_multiply_red[v_ax0, v_ax1])
|
||||
T.writes(rsqrt[v_ax0, v_ax1])
|
||||
rsqrt[v_ax0, v_ax1] = T.rsqrt(T_multiply_red[v_ax0, v_ax1] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_rms_norm"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(rsqrt[v_ax0, v_ax1], data[v_ax0, v_ax1, v_ax2], weight[v_ax2])
|
||||
T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2])
|
||||
T_rms_norm[v_ax0, v_ax1, v_ax2] = rsqrt[v_ax0, v_ax1] * data[v_ax0, v_ax1, v_ax2] * weight[v_ax2]
|
||||
for ax0, ax1, ax2 in T.grid(1, n, 4096):
|
||||
with T.sblock("T_cast_2"):
|
||||
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(T_rms_norm[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_cast[v_ax0, v_ax1, v_ax2])
|
||||
T_cast[v_ax0, v_ax1, v_ax2] = T_rms_norm[v_ax0, v_ax1, v_ax2]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_data: T.handle, weight: T.Buffer((4096,), "float32"), var_T_cast: T.handle):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
n = T.int32()
|
||||
data = T.match_buffer(var_data, (1, n, 4096))
|
||||
T_cast = T.match_buffer(var_T_cast, (1, n, 4096))
|
||||
# with T.sblock("root"):
|
||||
T_multiply_local = T.sblock_alloc_buffer((1, n, 4096), scope="local")
|
||||
T_multiply_red_local = T.sblock_alloc_buffer((1, n), scope="local")
|
||||
rsqrt_shared = T.sblock_alloc_buffer((1, n), scope="shared")
|
||||
T_rms_norm_local = T.sblock_alloc_buffer((1, n, 4096), scope="local")
|
||||
data_local = T.sblock_alloc_buffer((1, n, 4096), scope="local")
|
||||
for ax0_ax1_fused in T.thread_binding(n, thread="blockIdx.x"):
|
||||
for ax2_0 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
for ax2_1 in range(1):
|
||||
for ax2_2 in T.vectorized(8):
|
||||
with T.sblock("data_local"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v2 = T.axis.spatial(4096, ax2_0 * 8 + ax2_1 * 8 + ax2_2)
|
||||
T.reads(data[v0, v1, v2])
|
||||
T.writes(data_local[v0, v1, v2])
|
||||
data_local[v0, v1, v2] = data[v0, v1, v2]
|
||||
for ax0 in range(8):
|
||||
with T.sblock("T_multiply"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_ax2 = T.axis.spatial(4096, ax2_0 * 8 + ax0)
|
||||
T.reads(data_local[v_ax0, v_ax1, v_ax2])
|
||||
T.writes(T_multiply_local[v_ax0, v_ax1, v_ax2])
|
||||
T_multiply_local[v_ax0, v_ax1, v_ax2] = data_local[v_ax0, v_ax1, v_ax2] * data_local[v_ax0, v_ax1, v_ax2]
|
||||
for ax0 in range(8):
|
||||
with T.sblock("T_multiply_red"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_k2 = T.axis.reduce(4096, ax2_0 * 8 + ax0)
|
||||
T.reads(T_multiply_local[v_ax0, v_ax1, v_k2])
|
||||
T.writes(T_multiply_red_local[v_ax0, v_ax1])
|
||||
with T.init():
|
||||
T_multiply_red_local[v_ax0, v_ax1] = T.float32(0)
|
||||
T_multiply_red_local[v_ax0, v_ax1] = T_multiply_red_local[v_ax0, v_ax1] + T_multiply_local[v_ax0, v_ax1, v_k2]
|
||||
with T.sblock("rsqrt"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
T.reads(T_multiply_red_local[v_ax0, v_ax1])
|
||||
T.writes(rsqrt_shared[v_ax0, v_ax1])
|
||||
rsqrt_shared[v_ax0, v_ax1] = T.rsqrt(T_multiply_red_local[v_ax0, v_ax1] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))
|
||||
for ax0_0 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
for ax0_1, ax0_2 in T.grid(1, 8):
|
||||
with T.sblock("T_rms_norm"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v_ax2 = T.axis.spatial(4096, ax0_0 * 8 + ax0_1 * 8 + ax0_2)
|
||||
T.reads(rsqrt_shared[v_ax0, v_ax1], data_local[v_ax0, v_ax1, v_ax2], weight[v_ax2])
|
||||
T.writes(T_rms_norm_local[v_ax0, v_ax1, v_ax2])
|
||||
T_rms_norm_local[v_ax0, v_ax1, v_ax2] = rsqrt_shared[v_ax0, v_ax1] * data_local[v_ax0, v_ax1, v_ax2] * weight[v_ax2]
|
||||
for ax0 in T.vectorized(8):
|
||||
with T.sblock("T_cast_local"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(n, ax0_ax1_fused)
|
||||
v2 = T.axis.spatial(4096, ax0_0 * 8 + ax0)
|
||||
T.reads(T_rms_norm_local[v0, v1, v2])
|
||||
T.writes(T_cast[v0, v1, v2])
|
||||
T_cast[v0, v1, v2] = T_rms_norm_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,190 @@
|
||||
# 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, F401
|
||||
import tvm
|
||||
from tvm.ir import IRModule, assert_structural_equal
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _check(mod_before: IRModule, mod_after: IRModule):
|
||||
target = Target("nvidia/geforce-rtx-3090-ti")
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
|
||||
dl.gpu.Transpose(),
|
||||
)(mod_before)
|
||||
assert_structural_equal(mod, mod_after)
|
||||
|
||||
|
||||
def test_transpose():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(rxplaceholder: T.Buffer((T.int64(512), T.int64(4096)), "float32"), T_transpose: T.Buffer((T.int64(4096), T.int64(512)), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
for ax0, ax1 in T.grid(T.int64(4096), T.int64(512)):
|
||||
with T.sblock("T_transpose"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(rxplaceholder: T.Buffer((T.int64(512), T.int64(4096)), "float32"), T_transpose: T.Buffer((T.int64(4096), T.int64(512)), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
rxplaceholder_shared = T.sblock_alloc_buffer((T.int64(512), T.int64(4096)), scope="shared")
|
||||
for ax0_0_0 in T.thread_binding(T.int64(512), thread="blockIdx.y", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0 in T.thread_binding(T.int64(32), thread="blockIdx.x"):
|
||||
for ax0_ax1_fused_0 in range(T.int64(1)):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.unroll(T.int64(1)):
|
||||
with T.sblock("rxplaceholder_shared"):
|
||||
v0 = T.axis.spatial(T.int64(512), ax1_0 * T.int64(16) + (ax0_ax1_fused_0 * T.int64(128) + ax0_ax1_fused_1 * T.int64(16) + ax0_ax1_fused_2 + ax0_ax1_fused_3) // T.int64(8))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax0_0_0 * T.int64(8) + (ax0_ax1_fused_0 * T.int64(128) + ax0_ax1_fused_1 * T.int64(16) + ax0_ax1_fused_2 + ax0_ax1_fused_3) % T.int64(8))
|
||||
T.reads(rxplaceholder[v0, v1])
|
||||
T.writes(rxplaceholder_shared[v0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 32, 1]]})
|
||||
rxplaceholder_shared[v0, v1] = rxplaceholder[v0, v1]
|
||||
for ax0_0_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax1_1 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_1_0 in range(T.int64(1)):
|
||||
for ax0_1_1 in range(T.int64(1)):
|
||||
with T.sblock("T_transpose"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax0_0_0 * T.int64(8) + ax0_0_1 + ax0_1_0 + ax0_1_1)
|
||||
v1 = T.axis.spatial(T.int64(512), ax1_0 * T.int64(16) + ax1_1)
|
||||
T.reads(rxplaceholder_shared[v1, v0])
|
||||
T.writes(T_transpose[v0, v1])
|
||||
T_transpose[v0, v1] = rxplaceholder_shared[v1, v0]
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_decode_transpose():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(rxplaceholder: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), rxplaceholder_1: T.Buffer((T.int64(128), T.int64(4096)), "uint32"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
decode = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)))
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(rxplaceholder[v_i // T.int64(8), v_j], rxplaceholder_1[v_i // T.int64(32), v_j])
|
||||
T.writes(decode[v_i, v_j])
|
||||
decode[v_i, v_j] = T.Cast("float32", T.bitwise_and(T.shift_right(rxplaceholder[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8) * T.int64(4))), T.uint32(15))) * T.reinterpret("float32", T.shift_left(T.bitwise_and(rxplaceholder_1[v_i // T.int64(32), v_j], T.uint32(65535)), T.uint32(16))) + T.reinterpret("float32", T.shift_left(T.bitwise_and(T.shift_right(rxplaceholder_1[v_i // T.int64(32), v_j], T.uint32(16)), T.uint32(65535)), T.uint32(16)))
|
||||
for ax0, ax1 in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("T_transpose"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(decode[v_ax1, v_ax0])
|
||||
T.writes(T_transpose[v_ax0, v_ax1])
|
||||
T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(rxplaceholder: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), rxplaceholder_1: T.Buffer((T.int64(128), T.int64(4096)), "uint32"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float32")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
decode_shared = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), scope="shared")
|
||||
for ax0_0_0 in T.thread_binding(T.int64(64), thread="blockIdx.y", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
||||
for ax0_ax1_fused_0 in range(T.int64(1)):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.unroll(T.int64(8)):
|
||||
with T.sblock("decode_shared"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax1_0 * T.int64(16) + (ax0_ax1_fused_0 * T.int64(1024) + ax0_ax1_fused_1 * T.int64(128) + ax0_ax1_fused_2 * T.int64(8) + ax0_ax1_fused_3) // T.int64(64))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax0_0_0 * T.int64(64) + (ax0_ax1_fused_0 * T.int64(1024) + ax0_ax1_fused_1 * T.int64(128) + ax0_ax1_fused_2 * T.int64(8) + ax0_ax1_fused_3) % T.int64(64))
|
||||
T.reads(rxplaceholder[v0 // T.int64(8), v1], rxplaceholder_1[v0 // T.int64(32), v1])
|
||||
T.writes(decode_shared[v0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 32, 1]]})
|
||||
decode_shared[v0, v1] = T.Cast("float32", T.bitwise_and(T.shift_right(rxplaceholder[v0 // T.int64(8), v1], T.Cast("uint32", v0 % T.int64(8) * T.int64(4))), T.uint32(15))) * T.reinterpret("float32", T.shift_left(T.bitwise_and(rxplaceholder_1[v0 // T.int64(32), v1], T.uint32(65535)), T.uint32(16))) + T.reinterpret("float32", T.shift_left(T.bitwise_and(T.shift_right(rxplaceholder_1[v0 // T.int64(32), v1], T.uint32(16)), T.uint32(65535)), T.uint32(16)))
|
||||
for ax0_0_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax1_1 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_1_0 in range(T.int64(2)):
|
||||
for ax0_1_1 in T.vectorized(T.int64(4)):
|
||||
with T.sblock("T_transpose"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax0_0_0 * T.int64(64) + ax0_0_1 * T.int64(8) + ax0_1_0 * T.int64(4) + ax0_1_1)
|
||||
v1 = T.axis.spatial(T.int64(4096), ax1_0 * T.int64(16) + ax1_1)
|
||||
T.reads(decode_shared[v1, v0])
|
||||
T.writes(T_transpose[v0, v1])
|
||||
T_transpose[v0, v1] = decode_shared[v1, v0]
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
|
||||
|
||||
def test_decode_int3_transpose():
|
||||
# fmt: off
|
||||
@I.ir_module(s_tir=True)
|
||||
class Before:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((T.int64(412), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(103), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float16")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
decode_1 = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
|
||||
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("decode"):
|
||||
v_i, v_j = T.axis.remap("SS", [i, j])
|
||||
T.reads(A[v_i // T.int64(10), v_j], B[v_i // T.int64(40), v_j])
|
||||
T.writes(decode_1[v_i, v_j])
|
||||
decode_1[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(10), v_j], T.Cast("uint32", v_i % T.int64(10)) * T.uint32(3)), T.uint32(7))) - T.float16(3)) * B[v_i // T.int64(40), v_j]
|
||||
for ax0, ax1 in T.grid(T.int64(4096), T.int64(4096)):
|
||||
with T.sblock("T_transpose"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads(decode_1[v_ax1, v_ax0])
|
||||
T.writes(T_transpose[v_ax0, v_ax1])
|
||||
T_transpose[v_ax0, v_ax1] = decode_1[v_ax1, v_ax0]
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class After:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((T.int64(412), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(103), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float16")):
|
||||
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
decode_1_shared = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16", scope="shared")
|
||||
for ax0_0_0 in T.thread_binding(T.int64(52), thread="blockIdx.y", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
||||
for ax1_0 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
||||
for ax0_ax1_fused_0 in range(T.int64(2)):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.unroll(T.int64(10)):
|
||||
with T.sblock("decode_1_shared"):
|
||||
v0 = T.axis.spatial(T.int64(4096), ax1_0 * T.int64(16) + (ax0_ax1_fused_0 * T.int64(1280) + ax0_ax1_fused_1 * T.int64(160) + ax0_ax1_fused_2 * T.int64(10) + ax0_ax1_fused_3) // T.int64(82))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax0_0_0 * T.int64(80) + (ax0_ax1_fused_0 * T.int64(1280) + ax0_ax1_fused_1 * T.int64(160) + ax0_ax1_fused_2 * T.int64(10) + ax0_ax1_fused_3) % T.int64(82))
|
||||
T.where(ax0_0_0 * T.int64(80) + (((ax0_ax1_fused_0 * T.int64(8) + ax0_ax1_fused_1) * T.int64(16) + ax0_ax1_fused_2) * T.int64(10) + ax0_ax1_fused_3) % T.int64(82) < T.int64(4096) and ((ax0_ax1_fused_0 * T.int64(8) + ax0_ax1_fused_1) * T.int64(16) + ax0_ax1_fused_2) * T.int64(10) + ax0_ax1_fused_3 < T.int64(1312))
|
||||
T.reads(A[v0 // T.int64(10), v1], B[v0 // T.int64(40), v1])
|
||||
T.writes(decode_1_shared[v0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 32, 1]]})
|
||||
decode_1_shared[v0, v1] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v0 // T.int64(10), v1], T.Cast("uint32", v0 % T.int64(10)) * T.uint32(3)), T.uint32(7))) - T.float16(3)) * B[v0 // T.int64(40), v1]
|
||||
for ax0_0_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"):
|
||||
for ax1_1 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
||||
for ax0_1_0 in range(T.int64(3)):
|
||||
for ax0_1_1 in T.vectorized(T.int64(4)):
|
||||
with T.sblock("T_transpose"):
|
||||
v0 = T.axis.spatial(T.int64(4096), (ax0_0_0 * T.int64(8) + ax0_0_1) * T.int64(10) + (ax0_1_0 * T.int64(4) + ax0_1_1))
|
||||
v1 = T.axis.spatial(T.int64(4096), ax1_0 * T.int64(16) + ax1_1)
|
||||
T.where((ax0_0_0 * T.int64(8) + ax0_0_1) * T.int64(10) + (ax0_1_0 * T.int64(4) + ax0_1_1) < T.int64(4096) and ax0_0_0 * T.int64(8) + ax0_0_1 < T.int64(410) and ax0_1_0 * T.int64(4) + ax0_1_1 < T.int64(10))
|
||||
T.reads(decode_1_shared[v1, v0])
|
||||
T.writes(T_transpose[v0, v1])
|
||||
T_transpose[v0, v1] = decode_1_shared[v1, v0]
|
||||
# fmt: on
|
||||
_check(Before, After)
|
||||
@@ -0,0 +1,65 @@
|
||||
# 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: F841
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((), "int32"), T_stack: T.Buffer((T.int64(3),), "int32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
compile_engine_const = T.sblock_alloc_buffer((), "int32")
|
||||
compile_engine_const_1 = T.sblock_alloc_buffer((), "int32")
|
||||
with T.sblock("compile_engine_const"):
|
||||
vi = T.axis.spatial(1, T.int64(0))
|
||||
T.reads()
|
||||
T.writes(compile_engine_const[()])
|
||||
compile_engine_const[()] = 16
|
||||
with T.sblock("compile_engine_const_1"):
|
||||
vi = T.axis.spatial(1, T.int64(0))
|
||||
T.reads()
|
||||
T.writes(compile_engine_const_1[()])
|
||||
compile_engine_const_1[()] = 20
|
||||
for ax0 in range(T.int64(3)):
|
||||
with T.sblock("T_stack"):
|
||||
v_ax0 = T.axis.spatial(T.int64(3), ax0)
|
||||
T.reads(compile_engine_const[()], p0[()], compile_engine_const_1[()])
|
||||
T.writes(T_stack[v_ax0])
|
||||
T_stack[v_ax0] = T.if_then_else(
|
||||
v_ax0 == T.int64(2),
|
||||
compile_engine_const[()],
|
||||
T.if_then_else(v_ax0 == T.int64(1), p0[()], compile_engine_const_1[()]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_normalize_primfunc_with_scalar():
|
||||
sch = tvm.s_tir.Schedule(main)
|
||||
f_normalize_prim_func = tvm.get_global_func("s_tir.schedule.NormalizePrimFunc")
|
||||
assert f_normalize_prim_func(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user