wehub-resource-sync
1184 lines
80 KiB
Python
1184 lines
80 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,line-too-long,invalid-name,too-few-public-methods,too-many-locals
|
|
# ruff: noqa: E501, 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_decode_gemv_1():
|
|
# NK layout + K as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[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] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
|
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
W = T.match_buffer(W_handle, (4096, 512), "uint32")
|
|
S = T.match_buffer(S_handle, (4096, 128), "float16")
|
|
V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
|
|
C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
C_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), "float16", scope="local")
|
|
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
|
|
for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads()
|
|
T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
|
|
for ax1_0_fused_0 in range(1):
|
|
for ax1_1 in range(8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
|
|
vax1_1 = T.axis.reduce(8, ax1_1)
|
|
T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_0_fused_1, 0, 0, v0] = C_rf_local[vax1_0_fused_1, 0, 0, v0] + V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
for ax1_fused in range(1):
|
|
for ax0 in T.thread_binding(512, thread="threadIdx.x"):
|
|
with T.sblock("matmul"):
|
|
vax1_0_fused_1 = T.axis.reduce(512, ax0)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
T.writes(C[0, 0, v0])
|
|
with T.init():
|
|
C[0, 0, v0] = T.float16(0)
|
|
C[0, 0, v0] = C[0, 0, v0] + C_rf_local[vax1_0_fused_1, 0, 0, v0]
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_decode_gemv_2():
|
|
# KN layout + K as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i // 8, v_j], S[v_i // 32, v_j])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i // 32, v_j]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[v_i0, v_i1, v_k], B[v_k, v_i2])
|
|
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] + V[v_i0, v_i1, v_k] * B[v_k, v_i2]
|
|
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
C_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
|
|
for i2_i0_i1_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
|
|
for i2_i0_i1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for k_0_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul_rf_init"):
|
|
vk_0_fused_1 = T.axis.spatial(16, k_0_fused_1)
|
|
v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + i2_i0_i1_fused_1)
|
|
C_rf_local[vk_0_fused_1, 0, 0, v_i2] = T.float16(0)
|
|
for k_0_fused_0, k_1 in T.grid(32, 8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vk_0_fused_1 = T.axis.spatial(16, k_0_fused_1)
|
|
v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + i2_i0_i1_fused_1)
|
|
vk_0_fused_0, vk_1 = T.axis.remap("RR", [k_0_fused_0, k_1])
|
|
C_rf_local[vk_0_fused_1, 0, 0, v_i2] = C_rf_local[vk_0_fused_1, 0, 0, v_i2] + V[0, 0, vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[(vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) // 8, v_i2], T.Cast("uint32", (vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[(vk_0_fused_0 * 128 + vk_0_fused_1 * 8 + vk_1) // 32, v_i2])
|
|
for ax1_ax2_ax3_fused in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0_fused in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vk_0_fused_1 = T.axis.reduce(16, ax0_fused)
|
|
v_i2 = T.axis.spatial(4096, i2_i0_i1_fused_0 * 16 + ax1_ax2_ax3_fused)
|
|
with T.init():
|
|
C[0, 0, v_i2] = T.float16(0)
|
|
C[0, 0, v_i2] = C[0, 0, v_i2] + C_rf_local[vk_0_fused_1, 0, 0, v_i2]
|
|
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_decode_gemv_3():
|
|
# NK layout + N as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((512, 4096), "uint32"), S: T.Buffer((128, 4096), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i // 8, v_j], S[v_i // 32, v_j])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i // 8, v_j], T.Cast("uint32", v_i % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i // 32, v_j]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[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] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
W = T.match_buffer(W_handle, (512, 4096), "uint32")
|
|
S = T.match_buffer(S_handle, (128, 4096), "float16")
|
|
V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
|
|
C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
C_rf_local = T.sblock_alloc_buffer((1024, 1, 1, 4096), "float16", scope="local")
|
|
for ax0_0_fused in T.thread_binding(512, thread="blockIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(1024, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
for ax0_1_init in range(8):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax0_1_init)
|
|
T.reads()
|
|
T.writes(C_rf_local[vax1_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_fused_1, 0, 0, v0] = T.float16(0)
|
|
for ax1_fused_0 in range(4):
|
|
for ax0_1 in range(8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax0_1)
|
|
vax1_fused_0 = T.axis.reduce(4, ax1_fused_0)
|
|
T.reads(C_rf_local[vax1_fused_1, 0, 0, v0], V[0, 0, vax1_fused_0 * 1024 + vax1_fused_1], W[v0 // 8, vax1_fused_0 * 1024 + vax1_fused_1], S[v0 // 32, vax1_fused_0 * 1024 + vax1_fused_1])
|
|
T.writes(C_rf_local[vax1_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_fused_1, 0, 0, v0] = C_rf_local[vax1_fused_1, 0, 0, v0] + V[0, 0, vax1_fused_0 * 1024 + vax1_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0 // 8, vax1_fused_0 * 1024 + vax1_fused_1], T.Cast("uint32", v0 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0 // 32, vax1_fused_0 * 1024 + vax1_fused_1])
|
|
for ax1_fused_0 in range(1):
|
|
for ax0 in T.thread_binding(1024, thread="threadIdx.x"):
|
|
for ax1_fused_1 in range(8):
|
|
with T.sblock("matmul"):
|
|
vax1_fused_1 = T.axis.reduce(1024, ax0)
|
|
v0 = T.axis.spatial(4096, ax0_0_fused * 8 + ax1_fused_1)
|
|
T.reads(C_rf_local[vax1_fused_1, 0, 0, v0])
|
|
T.writes(C[0, 0, v0])
|
|
with T.init():
|
|
C[0, 0, v0] = T.float16(0)
|
|
C[0, 0, v0] = C[0, 0, v0] + C_rf_local[vax1_fused_1, 0, 0, v0]
|
|
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_decode_gemv_4():
|
|
# KN layout + N as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[v_i0, v_i1, v_k], B[v_k, v_i2])
|
|
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] + V[v_i0, v_i1, v_k] * B[v_k, v_i2]
|
|
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
C_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
|
|
for i2_0_i0_i1_fused_0 in T.thread_binding(32, thread="blockIdx.x"):
|
|
for i2_0_i0_i1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for k_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
|
for i2_1_init in range(8):
|
|
with T.sblock("matmul_rf_init"):
|
|
vk_fused_1 = T.axis.spatial(16, k_fused_1)
|
|
v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + i2_0_i0_i1_fused_1 * 8 + i2_1_init)
|
|
C_rf_local[vk_fused_1, 0, 0, v_i2] = T.float16(0)
|
|
for k_fused_0, i2_1 in T.grid(256, 8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vk_fused_1 = T.axis.spatial(16, k_fused_1)
|
|
v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + i2_0_i0_i1_fused_1 * 8 + i2_1)
|
|
vk_fused_0 = T.axis.reduce(256, k_fused_0)
|
|
C_rf_local[vk_fused_1, 0, 0, v_i2] = C_rf_local[vk_fused_1, 0, 0, v_i2] + V[0, 0, vk_fused_0 * 16 + vk_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[vk_fused_0 * 16 + vk_fused_1, v_i2 // 8], T.Cast("uint32", v_i2 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[vk_fused_0 * 16 + vk_fused_1, v_i2 // 32])
|
|
for ax1_ax2_ax3_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax1_ax2_ax3_fused_1 in range(8):
|
|
for ax0_fused in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vk_fused_1 = T.axis.reduce(16, ax0_fused)
|
|
v_i2 = T.axis.spatial(4096, i2_0_i0_i1_fused_0 * 128 + ax1_ax2_ax3_fused_0 * 8 + ax1_ax2_ax3_fused_1)
|
|
with T.init():
|
|
C[0, 0, v_i2] = T.float16(0)
|
|
C[0, 0, v_i2] = C[0, 0, v_i2] + C_rf_local[vk_fused_1, 0, 0, v_i2]
|
|
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_decode_gemv_sigmoid():
|
|
# NK layout + K as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), D: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
C = T.sblock_alloc_buffer((1, 1, 4096), "float16")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[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] + V[v_i0, v_i1, v_k] * B[v_i2, v_k]
|
|
for i0, i1, i2 in T.grid(1, 1, 4096):
|
|
with T.sblock("sigmoid"):
|
|
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(C[v_i0, v_i1, v_i2])
|
|
T.writes(D[v_i0, v_i1, v_i2])
|
|
D[v_i0, v_i1, v_i2] = T.sigmoid(C[v_i0, v_i1, v_i2])
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, D_handle: T.handle):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
W = T.match_buffer(W_handle, (4096, 512), "uint32")
|
|
S = T.match_buffer(S_handle, (4096, 128), "float16")
|
|
V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
|
|
D = T.match_buffer(D_handle, (1, 1, 4096), "float16")
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
C_local = T.sblock_alloc_buffer((1, 1, 4096), "float16", scope="local")
|
|
C_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), "float16", scope="local")
|
|
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
|
|
for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads()
|
|
T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
|
|
for ax1_0_fused_0 in range(1):
|
|
for ax1_1 in range(8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
|
|
vax1_1 = T.axis.reduce(8, ax1_1)
|
|
T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
T.writes(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_rf_local[vax1_0_fused_1, 0, 0, v0] = C_rf_local[vax1_0_fused_1, 0, 0, v0] + V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
for ax1_fused in range(1):
|
|
for ax0 in T.thread_binding(512, thread="threadIdx.x"):
|
|
with T.sblock("matmul"):
|
|
vax1_0_fused_1 = T.axis.reduce(512, ax0)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads(C_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
T.writes(C_local[0, 0, v0])
|
|
with T.init():
|
|
C_local[0, 0, v0] = T.float16(0)
|
|
C_local[0, 0, v0] = C_local[0, 0, v0] + C_rf_local[vax1_0_fused_1, 0, 0, v0]
|
|
for ax0 in range(1):
|
|
with T.sblock("sigmoid"):
|
|
v0 = T.axis.spatial(4096, ax0_fused + ax0)
|
|
T.reads(C_local[0, 0, v0])
|
|
T.writes(D[0, 0, v0])
|
|
D[0, 0, v0] = T.sigmoid(C_local[0, 0, v0])
|
|
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_decode_gemv_1_fp32():
|
|
# NK layout + K as decode dim
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W: T.Buffer((4096, 512), "uint32"), S: T.Buffer((4096, 128), "float16"), V: T.Buffer((1, 1, 4096), "float16"), C: T.Buffer((1, 1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
B = T.sblock_alloc_buffer((4096, 4096), "float16")
|
|
C_fp32 = T.sblock_alloc_buffer((1, 1, 4096), "float32")
|
|
for i, j in T.grid(4096, 4096):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(W[v_i, v_j // 8], S[v_i, v_j // 32])
|
|
T.writes(B[v_i, v_j])
|
|
B[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(W[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v_i, v_j // 32]
|
|
for i0, i1, i2, k in T.grid(1, 1, 4096, 4096):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(V[v_i0, v_i1, v_k], B[v_i2, v_k])
|
|
T.writes(C_fp32[v_i0, v_i1, v_i2])
|
|
with T.init():
|
|
C_fp32[v_i0, v_i1, v_i2] = T.float16(0)
|
|
C_fp32[v_i0, v_i1, v_i2] = C_fp32[v_i0, v_i1, v_i2] + T.Cast("float32", V[v_i0, v_i1, v_k]) * T.Cast("float32", B[v_i2, v_k])
|
|
for i0, i1, i2 in T.grid(1, 1, 4096):
|
|
with T.sblock("cast"):
|
|
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(C_fp32[v_i0, v_i1, v_i2])
|
|
T.writes(C[v_i0, v_i1, v_i2])
|
|
C[v_i0, v_i1, v_i2] = T.Cast("float16", C_fp32[v_i0, v_i1, v_i2])
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(s_tir=True)
|
|
def func(W_handle: T.handle, S_handle: T.handle, V_handle: T.handle, C_handle: T.handle):
|
|
T.func_attr({"global_symbol": "main", "tirx.is_scheduled": True, "tirx.noalias": True})
|
|
W = T.match_buffer(W_handle, (4096, 512), "uint32")
|
|
S = T.match_buffer(S_handle, (4096, 128), "float16")
|
|
V = T.match_buffer(V_handle, (1, 1, 4096), "float16")
|
|
C = T.match_buffer(C_handle, (1, 1, 4096), "float16")
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
C_fp32_local = T.sblock_alloc_buffer((1, 1, 4096), scope="local")
|
|
C_fp32_rf_local = T.sblock_alloc_buffer((512, 1, 1, 4096), scope="local")
|
|
for ax0_fused in T.thread_binding(4096, thread="blockIdx.x"):
|
|
for ax1_0_fused_1 in T.thread_binding(512, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads()
|
|
T.writes(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float32(0)
|
|
for ax1_0_fused_0 in range(1):
|
|
for ax1_1 in range(8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_0_fused_1 = T.axis.spatial(512, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
vax1_0_fused_0 = T.axis.reduce(1, ax1_0_fused_0)
|
|
vax1_1 = T.axis.reduce(8, ax1_1)
|
|
T.reads(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0], V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1], W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
T.writes(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] = C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0] + T.Cast("float32", V[0, 0, vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1]) * T.Cast("float32", (T.Cast("float16", T.bitwise_and(T.shift_right(W[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 8], T.Cast("uint32", (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * S[v0, (vax1_0_fused_0 * 4096 + vax1_0_fused_1 * 8 + vax1_1) // 32])
|
|
for ax1_fused in range(1):
|
|
for ax0 in T.thread_binding(512, thread="threadIdx.x"):
|
|
with T.sblock("matmul"):
|
|
vax1_0_fused_1 = T.axis.reduce(512, ax0)
|
|
v0 = T.axis.spatial(4096, ax0_fused)
|
|
T.reads(C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
T.writes(C_fp32_local[0, 0, v0])
|
|
with T.init():
|
|
C_fp32_local[0, 0, v0] = T.float32(0)
|
|
C_fp32_local[0, 0, v0] = C_fp32_local[0, 0, v0] + C_fp32_rf_local[vax1_0_fused_1, 0, 0, v0]
|
|
for ax0 in range(1):
|
|
with T.sblock("cast"):
|
|
v0 = T.axis.spatial(4096, ax0_fused + ax0)
|
|
T.reads(C_fp32_local[0, 0, v0])
|
|
T.writes(C[0, 0, v0])
|
|
C[0, 0, v0] = T.Cast("float16", C_fp32_local[0, 0, v0])
|
|
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_reduction_no_spatial():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((1, 1, 4096), "float16"), B: T.Buffer((4096,), "float16"), rms_norm: T.Buffer((1, 4096), "float16")):
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
Ared_temp = T.sblock_alloc_buffer((1, 1))
|
|
for ax0 in range(4096):
|
|
with T.sblock("Ared_temp"):
|
|
v0 = T.axis.reduce(4096, ax0)
|
|
with T.init():
|
|
Ared_temp[0, 0] = T.float32(0)
|
|
Ared_temp[0, 0] = Ared_temp[0, 0] + T.Cast("float32", A[0, 0, v0]) * T.Cast("float32", A[0, 0, v0])
|
|
for ax0 in range(4096):
|
|
with T.sblock("rms_norm"):
|
|
v0 = T.axis.spatial(4096, ax0)
|
|
rms_norm[0, v0] = T.Cast("float16", T.Cast("float32", B[v0]) * (T.Cast("float32", A[0, 0, v0]) / T.sqrt(Ared_temp[0, 0] * 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(A_handle: T.handle, B_handle: T.handle, rms_norm_handle: T.handle):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
A = T.match_buffer(A_handle, (1, 1, 4096), "float16")
|
|
B = T.match_buffer(B_handle, (4096,), "float16")
|
|
rms_norm = T.match_buffer(rms_norm_handle, (1, 4096), "float16")
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
Ared_temp_shared = T.sblock_alloc_buffer((1, 1), scope="shared")
|
|
Ared_temp_rf_local = T.sblock_alloc_buffer((1024, 1, 1), scope="local")
|
|
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(1024, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("Ared_temp_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
T.reads()
|
|
T.writes(Ared_temp_rf_local[vax1_fused_1, 0, 0])
|
|
Ared_temp_rf_local[vax1_fused_1, 0, 0] = T.float32(0)
|
|
for ax1_fused_0 in range(4):
|
|
for u in range(1):
|
|
with T.sblock("Ared_temp_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(1024, ax1_fused_1)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
vax1_fused_0 = T.axis.reduce(4, ax1_fused_0)
|
|
T.reads(Ared_temp_rf_local[vax1_fused_1, 0, 0], A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1])
|
|
T.writes(Ared_temp_rf_local[vax1_fused_1, 0, 0])
|
|
Ared_temp_rf_local[vax1_fused_1, 0, 0] = Ared_temp_rf_local[vax1_fused_1, 0, 0] + T.Cast("float32", A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1]) * T.Cast("float32", A[0, 0, vax1_fused_0 * 1024 + vax1_fused_1])
|
|
for ax1_fused in range(T.int64(1)):
|
|
for ax0 in T.thread_binding(1024, thread="threadIdx.x"):
|
|
with T.sblock("Ared_temp"):
|
|
vax1_fused_1 = T.axis.reduce(1024, ax0)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
T.reads(Ared_temp_rf_local[vax1_fused_1, 0, 0])
|
|
T.writes(Ared_temp_shared[0, 0])
|
|
with T.init():
|
|
Ared_temp_shared[0, 0] = T.float32(0)
|
|
Ared_temp_shared[0, 0] = Ared_temp_shared[0, 0] + Ared_temp_rf_local[vax1_fused_1, 0, 0]
|
|
for ax0_fused_0 in range(4):
|
|
for ax0_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
|
with T.sblock("rms_norm"):
|
|
v0 = T.axis.spatial(4096, ax0_fused_0 * 1024 + ax0_fused_1)
|
|
T.reads(B[v0], A[0, 0, v0], Ared_temp_shared[0, 0])
|
|
T.writes(rms_norm[0, v0])
|
|
rms_norm[0, v0] = T.Cast("float16", T.Cast("float32", B[v0]) * (T.Cast("float32", A[0, 0, v0]) / T.sqrt(Ared_temp_shared[0, 0] * T.float32(0.000244140625) + T.float32(9.9999999999999995e-07))))
|
|
# fmt: on
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, After)
|
|
|
|
|
|
def test_spatial_inner_no_broadcasting():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(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})
|
|
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)) * 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)
|
|
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]
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main(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.is_scheduled": True, "tirx.noalias": True})
|
|
var_matmul_intermediate_local = T.sblock_alloc_buffer((1, 1, 4096), "float16", scope="local")
|
|
var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((16, 1, 1, 4096), "float16", scope="local")
|
|
for ax0_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
|
|
for ax0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax1_0_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_0_fused_1 = T.axis.spatial(16, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_1)
|
|
T.reads()
|
|
T.writes(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] = T.float16(0)
|
|
for ax1_0_fused_0, ax1_1 in T.grid(86, 8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_0_fused_1 = T.axis.spatial(16, ax1_0_fused_1)
|
|
v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_1)
|
|
vax1_0_fused_0, vax1_1 = T.axis.remap("RR", [ax1_0_fused_0, ax1_1])
|
|
T.reads(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0], lv574[0, 0, vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1], lv575[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 8, v0], lv576[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 32, v0])
|
|
T.writes(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] = var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0] + lv574[0, 0, vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv575[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 8, v0], T.Cast("uint32", (vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv576[(vax1_0_fused_0 * 128 + vax1_0_fused_1 * 8 + vax1_1) // 32, v0])
|
|
for ax1_fused in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vax1_0_fused_1 = T.axis.reduce(16, ax0)
|
|
v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax1_fused)
|
|
T.reads(var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0])
|
|
T.writes(var_matmul_intermediate_local[0, 0, v0])
|
|
with T.init():
|
|
var_matmul_intermediate_local[0, 0, v0] = T.float16(0)
|
|
var_matmul_intermediate_local[0, 0, v0] = var_matmul_intermediate_local[0, 0, v0] + var_matmul_intermediate_rf_local[vax1_0_fused_1, 0, 0, v0]
|
|
for ax0_fused_0_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0_fused_1 in range(1):
|
|
with T.sblock("T_add"):
|
|
v0 = T.axis.spatial(4096, ax0_fused_0 * 16 + ax0_fused_0_1 + ax0_fused_1)
|
|
T.reads(lv570[0, 0, v0], var_matmul_intermediate_local[0, 0, v0])
|
|
T.writes(p_output0_intermediate[0, 0, v0])
|
|
p_output0_intermediate[0, 0, v0] = lv570[0, 0, v0] + var_matmul_intermediate_local[0, 0, v0]
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_spatial_inner_broadcasting():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
temp_local = T.sblock_alloc_buffer((256,))
|
|
for j in T.serial(256):
|
|
for k in T.serial(256):
|
|
with T.sblock("sum"):
|
|
vj, vk = T.axis.remap("SR", [j, k])
|
|
T.reads(A[vk, vj])
|
|
T.writes(temp_local[vj])
|
|
with T.init():
|
|
temp_local[vj] = T.float32(0)
|
|
temp_local[vj] = temp_local[vj] + A[vk, vj]
|
|
for i, j in T.grid(256, 256):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
T.reads(temp_local[vj])
|
|
T.writes(B[vi, vj])
|
|
B[vi, vj] = A[vi, vj] + temp_local[vj]
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256, 256), "float32")):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
temp_local_shared = T.sblock_alloc_buffer((256,), scope="shared")
|
|
temp_local_rf_local = T.sblock_alloc_buffer((16, 256), scope="local")
|
|
for ax0_fused_0 in T.thread_binding(16, thread="blockIdx.x"):
|
|
for ax0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("sum_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
|
|
v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax0_fused_1)
|
|
T.reads()
|
|
T.writes(temp_local_rf_local[vax1_fused_1, v0])
|
|
temp_local_rf_local[vax1_fused_1, v0] = T.float32(0)
|
|
for ax1_fused_0, u in T.grid(16, 1):
|
|
with T.sblock("sum_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
|
|
v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax0_fused_1)
|
|
vax1_fused_0 = T.axis.reduce(16, ax1_fused_0)
|
|
T.reads(temp_local_rf_local[vax1_fused_1, v0], A[vax1_fused_0 * 16 + vax1_fused_1, v0])
|
|
T.writes(temp_local_rf_local[vax1_fused_1, v0])
|
|
temp_local_rf_local[vax1_fused_1, v0] = temp_local_rf_local[vax1_fused_1, v0] + A[vax1_fused_0 * 16 + vax1_fused_1, v0]
|
|
for ax1_fused in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("sum"):
|
|
vax1_fused_1 = T.axis.reduce(16, ax0)
|
|
v0 = T.axis.spatial(256, ax0_fused_0 * 16 + ax1_fused)
|
|
T.reads(temp_local_rf_local[vax1_fused_1, v0])
|
|
T.writes(temp_local_shared[v0])
|
|
with T.init():
|
|
temp_local_shared[v0] = T.float32(0)
|
|
temp_local_shared[v0] = temp_local_shared[v0] + temp_local_rf_local[vax1_fused_1, v0]
|
|
for ax0_ax1_fused_0 in range(16):
|
|
for ax0_ax1_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0_ax1_fused_2 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("add"):
|
|
v0 = T.axis.spatial(256, (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) // 16)
|
|
v1 = T.axis.spatial(256, ax0_fused_0 * 16 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 16 + ax0_ax1_fused_2) % 16)
|
|
T.reads(temp_local_shared[v1])
|
|
T.writes(B[v0, v1])
|
|
B[v0, v1] = A[v0, v1] + temp_local_shared[v1]
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reduction_inner_no_broadcasting():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
temp_local = T.sblock_alloc_buffer((256,))
|
|
for i in T.serial(256):
|
|
for k in T.serial(256):
|
|
with T.sblock("sum"):
|
|
vi, vk = T.axis.remap("SR", [i, k])
|
|
T.reads(A[vi, vk])
|
|
T.writes(temp_local[vi])
|
|
with T.init():
|
|
temp_local[vi] = T.float32(0)
|
|
temp_local[vi] = temp_local[vi] + A[vi, vk]
|
|
for i in T.grid(256):
|
|
with T.sblock("add"):
|
|
vi = T.axis.remap("S", [i])
|
|
T.reads(temp_local[vi])
|
|
T.writes(B[vi,])
|
|
B[vi] = temp_local[vi] + T.float32(1)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main(A: T.Buffer((256, 256), "float32"), B: T.Buffer((256,), "float32")):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
temp_local_local = T.sblock_alloc_buffer((256,), scope="local")
|
|
temp_local_rf_local = T.sblock_alloc_buffer((256, 256), scope="local")
|
|
for ax0_fused in T.thread_binding(256, thread="blockIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(256, thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("sum_rf_init"):
|
|
vax1_fused_1, v0 = T.axis.remap("SS", [ax1_fused_1, ax0_fused])
|
|
T.reads()
|
|
T.writes(temp_local_rf_local[vax1_fused_1, v0])
|
|
temp_local_rf_local[vax1_fused_1, v0] = T.float32(0)
|
|
for ax1_fused_0, u in T.grid(1, 1):
|
|
with T.sblock("sum_rf_update"):
|
|
vax1_fused_1, v0, vax1_fused_0 = T.axis.remap("SSR", [ax1_fused_1, ax0_fused, ax1_fused_0])
|
|
T.reads(temp_local_rf_local[vax1_fused_1, v0], A[v0, vax1_fused_0 * 256 + vax1_fused_1])
|
|
T.writes(temp_local_rf_local[vax1_fused_1, v0])
|
|
temp_local_rf_local[vax1_fused_1, v0] = temp_local_rf_local[vax1_fused_1, v0] + A[v0, vax1_fused_0 * 256 + vax1_fused_1]
|
|
for ax1_fused in range(1):
|
|
for ax0 in T.thread_binding(256, thread="threadIdx.x"):
|
|
with T.sblock("sum"):
|
|
vax1_fused_1, v0 = T.axis.remap("RS", [ax0, ax0_fused])
|
|
T.reads(temp_local_rf_local[vax1_fused_1, v0])
|
|
T.writes(temp_local_local[v0])
|
|
with T.init():
|
|
temp_local_local[v0] = T.float32(0)
|
|
temp_local_local[v0] = temp_local_local[v0] + temp_local_rf_local[vax1_fused_1, v0]
|
|
for ax0 in range(1):
|
|
with T.sblock("add"):
|
|
v0 = T.axis.spatial(256, ax0_fused + ax0)
|
|
T.reads(temp_local_local[v0])
|
|
T.writes(B[v0])
|
|
B[v0] = temp_local_local[v0] + T.float32(1)
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reduction_inner_no_broadcasting2():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(lv9: T.Buffer((2560, 320), "uint32"), lv10: T.Buffer((2560, 80), "float16"), lv1: T.Buffer((1, 2560), "float16"), p_output0_intermediate: T.Buffer((1, 2560), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
p_output0_intermediate_1 = T.sblock_alloc_buffer((2560, 2560), "float16")
|
|
var_matmul_intermediate = T.sblock_alloc_buffer((1, 2560), "float16")
|
|
for i, j in T.grid(2560, 2560):
|
|
with T.sblock("decode"):
|
|
v_i, v_j = T.axis.remap("SS", [i, j])
|
|
T.reads(lv9[v_i, v_j // 8], lv10[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(lv9[v_i, v_j // 8], T.Cast("uint32", v_j % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv10[v_i, v_j // 32]
|
|
for i0, i1, k in T.grid(1, 2560, 2560):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
|
T.reads(lv1[v_i0, v_k], p_output0_intermediate_1[v_k, v_i1])
|
|
T.writes(var_matmul_intermediate[v_i0, v_i1])
|
|
with T.init():
|
|
var_matmul_intermediate[v_i0, v_i1] = T.float16(0)
|
|
var_matmul_intermediate[v_i0, v_i1] = var_matmul_intermediate[v_i0, v_i1] + lv1[v_i0, v_k] * p_output0_intermediate_1[v_k, v_i1]
|
|
for i0, i1 in T.grid(1, 2560):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(var_matmul_intermediate[v_i0, v_i1])
|
|
T.writes(p_output0_intermediate[v_i0, v_i1])
|
|
p_output0_intermediate[v_i0, v_i1] = T.Cast("float32", var_matmul_intermediate[v_i0, v_i1])
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main(lv9: T.Buffer((2560, 320), "uint32"), lv10: T.Buffer((2560, 80), "float16"), lv1: T.Buffer((1, 2560), "float16"), p_output0_intermediate: T.Buffer((1, 2560), "float32")):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
var_matmul_intermediate_local = T.sblock_alloc_buffer((1, 2560), "float16", scope="local")
|
|
var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((16, 1, 2560), "float16", scope="local")
|
|
for ax0_0_fused_0 in T.thread_binding(20, thread="blockIdx.x"):
|
|
for ax0_0_fused_1 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(16, thread="threadIdx.y"):
|
|
for ax0_1_init in range(8):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
|
|
v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_0_fused_1 * 8 + ax0_1_init)
|
|
T.reads()
|
|
T.writes(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
|
|
var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] = T.float16(0)
|
|
for ax1_fused_0, ax0_1 in T.grid(160, 8):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(16, ax1_fused_1)
|
|
v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_0_fused_1 * 8 + ax0_1)
|
|
vax1_fused_0 = T.axis.reduce(160, ax1_fused_0)
|
|
T.reads(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0], lv1[0, vax1_fused_0 * 16 + vax1_fused_1], lv9[vax1_fused_0 * 16 + vax1_fused_1, v0 // 8], lv10[vax1_fused_0 * 16 + vax1_fused_1, v0 // 32])
|
|
T.writes(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
|
|
var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] = var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0] + lv1[0, vax1_fused_0 * 16 + vax1_fused_1] * ((T.Cast("float16", T.bitwise_and(T.shift_right(lv9[vax1_fused_0 * 16 + vax1_fused_1, v0 // 8], T.Cast("uint32", v0 % 8) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv10[vax1_fused_0 * 16 + vax1_fused_1, v0 // 32])
|
|
for ax1_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax1_fused_1 in range(8):
|
|
for ax0 in T.thread_binding(16, thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vax1_fused_1 = T.axis.reduce(16, ax0)
|
|
v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax1_fused_0 * 8 + ax1_fused_1)
|
|
T.reads(var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0])
|
|
T.writes(var_matmul_intermediate_local[0, v0])
|
|
with T.init():
|
|
var_matmul_intermediate_local[0, v0] = T.float16(0)
|
|
var_matmul_intermediate_local[0, v0] = var_matmul_intermediate_local[0, v0] + var_matmul_intermediate_rf_local[vax1_fused_1, 0, v0]
|
|
for ax0_fused_0 in T.thread_binding(16, thread="threadIdx.x"):
|
|
for ax0_fused_1 in range(8):
|
|
with T.sblock("compute"):
|
|
v0 = T.axis.spatial(2560, ax0_0_fused_0 * 128 + ax0_fused_0 * 8 + ax0_fused_1)
|
|
T.reads(var_matmul_intermediate_local[0, v0])
|
|
T.writes(p_output0_intermediate[0, v0])
|
|
p_output0_intermediate[0, v0] = T.Cast("float32", var_matmul_intermediate_local[0, v0])
|
|
# fmt: on
|
|
|
|
with Target("nvidia/geforce-rtx-3090-ti"):
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reduction_inner_spatial_choose_perfect_factor():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16")):
|
|
T.func_attr({"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(100)), "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(100), 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]
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def main(var_A: T.handle, var_B: T.handle, matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16")):
|
|
T.func_attr({"tirx.is_scheduled": True, "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(100)), "float16")
|
|
# with T.sblock("root"):
|
|
matmul_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(32), T.int64(1), T.int64(100)), "float16", scope="local")
|
|
for ax0_ax1_fused_0 in T.thread_binding(T.int64(320), thread="blockIdx.x"):
|
|
for ax0_ax1_fused_1 in T.thread_binding(T.int64(10), thread="threadIdx.x"):
|
|
for ax2_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
|
|
v0 = T.axis.spatial(T.int64(32), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) // T.int64(100))
|
|
v1 = T.axis.spatial(T.int64(100), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) % T.int64(100))
|
|
T.reads()
|
|
T.writes(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = T.float16(0)
|
|
for ax2_fused_0, u in T.grid((n + T.int64(15)) // T.int64(16), 1):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
|
|
v0 = T.axis.spatial(T.int64(32), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) // T.int64(100))
|
|
v1 = T.axis.spatial(T.int64(100), (ax0_ax1_fused_0 * T.int64(10) + ax0_ax1_fused_1) % T.int64(100))
|
|
vax2_fused_0 = T.axis.reduce((n + T.int64(15)) // T.int64(16), ax2_fused_0)
|
|
T.where(ax2_fused_0 * T.int64(16) + ax2_fused_1 < n)
|
|
T.reads(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1], A[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1], B[T.int64(0), v0, vax2_fused_0 * T.int64(16) + vax2_fused_1, v1])
|
|
T.writes(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] + A[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1] * B[T.int64(0), v0, vax2_fused_0 * T.int64(16) + vax2_fused_1, v1]
|
|
for ax1_ax2_fused in T.thread_binding(T.int64(10), thread="threadIdx.x"):
|
|
for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vax2_fused_1 = T.axis.reduce(T.int64(16), ax0)
|
|
v0 = T.axis.spatial(T.int64(32), ax0_ax1_fused_0 // T.int64(10))
|
|
v1 = T.axis.spatial(T.int64(100), ax0_ax1_fused_0 % T.int64(10) * T.int64(10) + ax1_ax2_fused)
|
|
T.reads(matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
T.writes(matmul[T.int64(0), v0, T.int64(0), v1])
|
|
with T.init():
|
|
matmul[T.int64(0), v0, T.int64(0), v1] = T.float16(0)
|
|
matmul[T.int64(0), v0, T.int64(0), v1] = matmul[T.int64(0), v0, T.int64(0), v1] + matmul_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1]
|
|
# fmt: on
|
|
|
|
with Target("nvidia/geforce-rtx-3090-ti"):
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_repeat_transpose_gemv():
|
|
# fmt: off
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_relax_repeat_relax_permute_dims_relax_matmul1(p_lv716: T.handle, p_astype66: T.handle, var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
kv_seq_len = T.int64()
|
|
lv716 = T.match_buffer(p_lv716, (T.int64(1), kv_seq_len, T.int64(8), T.int64(128)), "float16")
|
|
astype66 = T.match_buffer(p_astype66, (T.int64(1), T.int64(32), T.int64(1), kv_seq_len), "float16")
|
|
# with T.sblock("root"):
|
|
var_T_repeat_intermediate = T.sblock_alloc_buffer((T.int64(1), kv_seq_len, T.int64(32), T.int64(128)), "float16")
|
|
var_T_transpose_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(32), kv_seq_len, T.int64(128)), "float16")
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), kv_seq_len, T.int64(32), T.int64(128)):
|
|
with T.sblock("T_repeat"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(lv716[v_ax0, v_ax1, v_ax2 // T.int64(4), v_ax3])
|
|
T.writes(var_T_repeat_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
var_T_repeat_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv716[v_ax0, v_ax1, v_ax2 // T.int64(4), v_ax3]
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), kv_seq_len, T.int64(128)):
|
|
with T.sblock("T_transpose"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(var_T_repeat_intermediate[v_ax0, v_ax2, v_ax1, v_ax3])
|
|
T.writes(var_T_transpose_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
var_T_transpose_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = var_T_repeat_intermediate[v_ax0, v_ax2, v_ax1, v_ax3]
|
|
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(astype66[v_i0, v_i1, v_i2, v_k], var_T_transpose_intermediate[v_i0, v_i1, v_k, v_i3])
|
|
T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3])
|
|
with T.init():
|
|
var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = T.float16(0)
|
|
var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + astype66[v_i0, v_i1, v_i2, v_k] * var_T_transpose_intermediate[v_i0, v_i1, v_k, v_i3]
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_relax_repeat_relax_permute_dims_relax_matmul1(p_lv716: T.handle, p_astype66: T.handle, var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16")):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
kv_seq_len = T.int64()
|
|
lv716 = T.match_buffer(p_lv716, (T.int64(1), kv_seq_len, T.int64(8), T.int64(128)), "float16")
|
|
astype66 = T.match_buffer(p_astype66, (T.int64(1), T.int64(32), T.int64(1), kv_seq_len), "float16")
|
|
# with T.sblock("root"):
|
|
var_matmul_intermediate_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(32), T.int64(1), T.int64(128)), "float16", scope="local")
|
|
for ax0_0_ax1_fused_0 in T.thread_binding(T.int64(64), thread="blockIdx.x"):
|
|
for ax0_0_ax1_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
|
for ax2_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
for ax0_1_init in range(T.int64(4)):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
|
|
v0 = T.axis.spatial(T.int64(32), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) // T.int64(128) * T.int64(4) + ax0_1_init)
|
|
v1 = T.axis.spatial(T.int64(128), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) % T.int64(128))
|
|
T.reads()
|
|
T.writes(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = T.float16(0)
|
|
for ax2_fused_0, ax0_1 in T.grid((kv_seq_len + T.int64(15)) // T.int64(16), T.int64(4)):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax2_fused_1 = T.axis.spatial(T.int64(16), ax2_fused_1)
|
|
v0 = T.axis.spatial(T.int64(32), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) // T.int64(128) * T.int64(4) + ax0_1)
|
|
v1 = T.axis.spatial(T.int64(128), (ax0_0_ax1_fused_0 * T.int64(16) + ax0_0_ax1_fused_1) % T.int64(128))
|
|
vax2_fused_0 = T.axis.reduce((kv_seq_len + T.int64(15)) // T.int64(16), ax2_fused_0)
|
|
T.where(ax2_fused_0 * T.int64(16) + ax2_fused_1 < kv_seq_len)
|
|
T.reads(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1], astype66[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1], lv716[T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1, v0 // T.int64(4), v1])
|
|
T.writes(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] = var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1] + astype66[T.int64(0), v0, T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1] * lv716[T.int64(0), vax2_fused_0 * T.int64(16) + vax2_fused_1, v0 // T.int64(4), v1]
|
|
for ax1_0_ax2_fused in T.thread_binding(T.int64(16), thread="threadIdx.x"):
|
|
for ax1_1 in range(T.int64(4)):
|
|
for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vax2_fused_1 = T.axis.reduce(T.int64(16), ax0)
|
|
v0 = T.axis.spatial(T.int64(32), ax0_0_ax1_fused_0 // T.int64(8) * T.int64(4) + ax1_1)
|
|
v1 = T.axis.spatial(T.int64(128), ax0_0_ax1_fused_0 % T.int64(8) * T.int64(16) + ax1_0_ax2_fused)
|
|
T.reads(var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1])
|
|
T.writes(var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1])
|
|
with T.init():
|
|
var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] = T.float16(0)
|
|
var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] = var_matmul_intermediate[T.int64(0), v0, T.int64(0), v1] + var_matmul_intermediate_rf_local[vax2_fused_1, T.int64(0), v0, T.int64(0), v1]
|
|
# fmt: on
|
|
|
|
with Target("nvidia/geforce-rtx-3090-ti"):
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_gemv_dyn_shape_epilogue():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(
|
|
var_A: T.handle,
|
|
B: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"),
|
|
var_C: T.handle,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
vocab_size = T.int64()
|
|
A = T.match_buffer(var_A, (T.int64(4096), vocab_size), "float16")
|
|
C = T.match_buffer(var_C, (T.int64(1), T.int64(1), vocab_size))
|
|
C_temp = T.sblock_alloc_buffer((T.int64(1), T.int64(1), vocab_size), "float16")
|
|
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), vocab_size, 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(B[v_i0, v_i1, v_k], A[v_k, v_i2])
|
|
T.writes(C_temp[v_i0, v_i1, v_i2])
|
|
with T.init():
|
|
C_temp[v_i0, v_i1, v_i2] = T.float16(0)
|
|
C_temp[v_i0, v_i1, v_i2] = (
|
|
C_temp[v_i0, v_i1, v_i2] + B[v_i0, v_i1, v_k] * A[v_k, v_i2]
|
|
)
|
|
for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), vocab_size):
|
|
with T.sblock("epilogue"):
|
|
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(C_temp[v_i0, v_i1, v_i2])
|
|
T.writes(C[v_i0, v_i1, v_i2])
|
|
C[v_i0, v_i1, v_i2] = T.Cast("float32", C_temp[v_i0, v_i1, v_i2])
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(var_A: T.handle, B: T.Buffer((T.int64(1), T.int64(1), T.int64(4096)), "float16"), var_C: T.handle):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
vocab_size = T.int64()
|
|
A = T.match_buffer(var_A, (T.int64(4096), vocab_size), "float16")
|
|
C = T.match_buffer(var_C, (T.int64(1), T.int64(1), vocab_size))
|
|
# with T.sblock("root"):
|
|
C_temp_local = T.sblock_alloc_buffer((T.int64(1), T.int64(1), vocab_size), "float16", scope="local")
|
|
C_temp_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(1), T.int64(1), vocab_size), "float16", scope="local")
|
|
for ax0_fused_0 in T.thread_binding(vocab_size, thread="blockIdx.x"):
|
|
for ax0_fused_1 in T.thread_binding(T.int64(1), thread="threadIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
with T.sblock("matmul_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(T.int64(16), ax1_fused_1)
|
|
v0 = T.axis.spatial(vocab_size, ax0_fused_0 + ax0_fused_1)
|
|
T.reads()
|
|
T.writes(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
|
|
C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] = T.float16(0)
|
|
for ax1_fused_0, u in T.grid(T.int64(256), 1):
|
|
with T.sblock("matmul_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(T.int64(16), ax1_fused_1)
|
|
v0 = T.axis.spatial(vocab_size, ax0_fused_0 + ax0_fused_1)
|
|
vax1_fused_0 = T.axis.reduce(T.int64(256), ax1_fused_0)
|
|
T.reads(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0], B[T.int64(0), T.int64(0), vax1_fused_0 * T.int64(16) + vax1_fused_1], A[vax1_fused_0 * T.int64(16) + vax1_fused_1, v0])
|
|
T.writes(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
|
|
C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] = C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0] + B[T.int64(0), T.int64(0), vax1_fused_0 * T.int64(16) + vax1_fused_1] * A[vax1_fused_0 * T.int64(16) + vax1_fused_1, v0]
|
|
for ax1_fused in T.thread_binding(T.int64(1), thread="threadIdx.x"):
|
|
for ax0 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
|
|
with T.sblock("matmul"):
|
|
vax1_fused_1, v0 = T.axis.remap("RS", [ax0, ax0_fused_0])
|
|
T.reads(C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0])
|
|
T.writes(C_temp_local[T.int64(0), T.int64(0), v0])
|
|
with T.init():
|
|
C_temp_local[T.int64(0), T.int64(0), v0] = T.float16(0)
|
|
C_temp_local[T.int64(0), T.int64(0), v0] = C_temp_local[T.int64(0), T.int64(0), v0] + C_temp_rf_local[vax1_fused_1, T.int64(0), T.int64(0), v0]
|
|
for ax0_fused_0_1 in T.thread_binding(T.int64(1), thread="threadIdx.x"):
|
|
for ax0_fused_1 in range(T.int64(1)):
|
|
with T.sblock("epilogue"):
|
|
v0 = T.axis.spatial(vocab_size, ax0_fused_0)
|
|
T.reads(C_temp_local[T.int64(0), T.int64(0), v0])
|
|
T.writes(C[T.int64(0), T.int64(0), v0])
|
|
C[T.int64(0), T.int64(0), v0] = T.Cast("float32", C_temp_local[T.int64(0), T.int64(0), v0])
|
|
# fmt: on
|
|
|
|
with Target("nvidia/geforce-rtx-3090-ti"):
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Module) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_gemv_output_one_element():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(A: T.Buffer((T.int64(1), T.int64(2048)), "float16"), weight: T.Buffer((T.int64(1), T.int64(2048)), "float16"), out: T.Buffer((T.int64(1), T.int64(1)), "float16")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
NT_matmul_intermediate = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), "float16")
|
|
for i0, i1, k in T.grid(T.int64(1), T.int64(1), T.int64(2048)):
|
|
with T.sblock("NT_matmul"):
|
|
v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
|
|
with T.init():
|
|
NT_matmul_intermediate[v_i0, v_i1] = T.float16(0)
|
|
NT_matmul_intermediate[v_i0, v_i1] = NT_matmul_intermediate[v_i0, v_i1] + A[v_i0, v_k] * weight[v_i1, v_k]
|
|
for i0, i1 in T.grid(T.int64(1), T.int64(1)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
out[v_i0, v_i1] = T.sigmoid(NT_matmul_intermediate[v_i0, v_i1])
|
|
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def main(A: T.Buffer((T.int64(1), T.int64(2048)), "float16"), weight: T.Buffer((T.int64(1), T.int64(2048)), "float16"), out: T.Buffer((T.int64(1), T.int64(1)), "float16")):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
NT_matmul_intermediate_shared = T.sblock_alloc_buffer((T.int64(1), T.int64(1)), "float16", scope="shared")
|
|
NT_matmul_intermediate_rf_local = T.sblock_alloc_buffer((T.int64(1024), T.int64(1), T.int64(1)), "float16", scope="local")
|
|
for ax0_fused in T.thread_binding(T.int64(1), thread="blockIdx.x"):
|
|
for ax1_fused_1 in T.thread_binding(T.int64(1024), thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 256, "pragma_unroll_explicit": 1}):
|
|
with T.sblock("NT_matmul_rf_init"):
|
|
vax1_fused_1 = T.axis.spatial(T.int64(1024), ax1_fused_1)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] = T.float16(0)
|
|
for ax1_fused_0, u in T.grid(T.int64(2), 1):
|
|
with T.sblock("NT_matmul_rf_update"):
|
|
vax1_fused_1 = T.axis.spatial(T.int64(1024), ax1_fused_1)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
vax1_fused_0 = T.axis.reduce(T.int64(2), ax1_fused_0)
|
|
NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] = NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)] + A[T.int64(0), vax1_fused_0 * T.int64(1024) + vax1_fused_1] * weight[T.int64(0), vax1_fused_0 * T.int64(1024) + vax1_fused_1]
|
|
for ax1_fused in range(T.int64(1)):
|
|
for ax0 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
|
|
with T.sblock("NT_matmul"):
|
|
vax1_fused_1 = T.axis.reduce(T.int64(1024), ax0)
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
with T.init():
|
|
NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] = T.float16(0)
|
|
NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] = NT_matmul_intermediate_shared[T.int64(0), T.int64(0)] + NT_matmul_intermediate_rf_local[vax1_fused_1, T.int64(0), T.int64(0)]
|
|
for ax0_fused_0 in range(T.int64(1)):
|
|
for ax0_fused_1 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
|
|
with T.sblock("compute"):
|
|
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
|
T.where(ax0_fused_0 * T.int64(1024) + ax0_fused_1 < T.int64(1))
|
|
out[T.int64(0), T.int64(0)] = T.sigmoid(NT_matmul_intermediate_shared[T.int64(0), T.int64(0)])
|
|
# fmt: on
|
|
|
|
with Target("nvidia/geforce-rtx-3090-ti"):
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_no_reduction_loop_check():
|
|
# The normalized prime func will not contain a reduction loop since its extent is one.
|
|
# This checks that the Reduction schedule is correctly not applied in this case
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def matmul(lv43: T.Buffer((T.int64(1), T.int64(32), T.int64(1)), "float16"), lv44: T.Buffer((T.int64(1), T.int64(1), T.int64(1)), "float16"), matmul: T.Buffer((T.int64(1), T.int64(32), T.int64(1)), "float16")):
|
|
T.func_attr({"op_pattern": 4, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for i0, i1, i2, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(1)):
|
|
with T.sblock("matmul"):
|
|
v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
|
|
T.reads(lv43[v_i0, v_i1, v_k], lv44[v_i0, v_k, v_i2])
|
|
T.writes(matmul[v_i0, v_i1, v_i2])
|
|
with T.init():
|
|
matmul[v_i0, v_i1, v_i2] = T.float16(0.0)
|
|
matmul[v_i0, v_i1, v_i2] = matmul[v_i0, v_i1, v_i2] + lv43[v_i0, v_i1, v_k] * lv44[v_i0, v_k, v_i2]
|
|
# fmt: on
|
|
|
|
target = Target("nvidia/geforce-rtx-3090-ti")
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(dl.gpu.Reduction())(Before) # pylint: disable=not-callable
|
|
assert_structural_equal(mod, Before)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|