wehub-resource-sync
263 lines
12 KiB
Python
263 lines
12 KiB
Python
# 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()
|