wehub-resource-sync
191 lines
14 KiB
Python
191 lines
14 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, 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)
|