wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
2d78e8e00decbd011b711126aa2f757165e941cd
apache--tvm/tests/python/tvmscript/test_tvmscript_roundtrip.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

3495 lines
133 KiB
Python
Raw Blame History

# 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.
# ruff: noqa: E501, F401, F841
import sys
import numpy as np
import pytest
import tvm
import tvm.testing
from tvm import tirx
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
def opt_gemm_lower():
@tvm.script.ir_module
class Module:
@T.prim_func(s_tir=True)
def mmult(A: T.handle, B: T.handle, C: T.handle) -> None:
# function attr dict
T.func_attr({"tirx.noalias": True})
A_1 = T.match_buffer(A, [16384], elem_offset=0, align=64, offset_factor=1)
B_1 = T.match_buffer(B, [1024, 1024], elem_offset=0, align=64, offset_factor=1)
C_1 = T.match_buffer(C, [16384], elem_offset=0, align=64, offset_factor=1)
# body
packedB = T.alloc_buffer((32768,))
for x in T.parallel(0, 32):
for y in T.serial(0, 1024):
packedB[T.ramp(((x * 32768) + (y * 32)), 1, 32)] = B_1[y, T.ramp(x * 32, 1, 32)]
for x_outer in T.parallel(0, 32):
C_global = T.alloc_buffer((1024,))
for y_outer in T.serial(0, 32):
for x_c_init in T.serial(0, 32):
C_global[T.ramp((x_c_init * 32), 1, 32)] = T.broadcast(T.float32(0), 32)
for k_outer in T.serial(0, 256):
for x_c in T.serial(0, 32):
C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
T.ramp((x_c * 32), 1, 32)
] + (
T.broadcast(
A_1[(((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4))],
32,
)
* packedB[T.ramp(((y_outer * 32768) + (k_outer * 128)), 1, 32)]
)
C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
T.ramp((x_c * 32), 1, 32)
] + (
T.broadcast(
A_1[
((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 1),
],
32,
)
* packedB[
T.ramp((((y_outer * 32768) + (k_outer * 128)) + 32), 1, 32)
]
)
C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
T.ramp((x_c * 32), 1, 32)
] + (
T.broadcast(
A_1[
((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 2),
],
32,
)
* packedB[
T.ramp((((y_outer * 32768) + (k_outer * 128)) + 64), 1, 32)
]
)
C_global[T.ramp((x_c * 32), 1, 32)] = C_global[
T.ramp((x_c * 32), 1, 32)
] + (
T.broadcast(
A_1[
((((x_outer * 32768) + (x_c * 1024)) + (k_outer * 4)) + 3),
],
32,
)
* packedB[
T.ramp((((y_outer * 32768) + (k_outer * 128)) + 96), 1, 32)
]
)
for x_inner in T.serial(0, 32):
for y_inner in T.serial(0, 32):
C_1[
(
(((x_outer * 32768) + (x_inner * 1024)) + (y_outer * 32))
+ y_inner
)
] = C_global[((x_inner * 32) + y_inner)]
return Module
def launch_env_thread():
@T.prim_func(s_tir=True)
def main(inputs: T.Buffer((64, 2, 4), "float32")) -> None:
bx = T.launch_thread("blockIdx.x", 64)
for i, j in T.grid(2, 4):
T.evaluate(inputs[bx, i, j])
return main
def opt_conv_tensorcore_lower():
@T.prim_func(s_tir=True)
def func(
A: T.Buffer((16, 14, 14, 16, 16, 16), "float16"),
W: T.Buffer((3, 3, 16, 32, 16, 16), "float16"),
Conv: T.Buffer((16, 14, 14, 32, 16, 16), "float32"),
) -> None:
# function attr dict
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
# body
A_1 = T.decl_buffer([12845056], dtype="float16", data=A.data)
W_1 = T.decl_buffer([1179648], dtype="float16", data=W.data)
Conv_1 = T.decl_buffer([25690112], data=Conv.data)
bx = T.env_thread("blockIdx.x")
by = T.env_thread("blockIdx.y")
bz = T.env_thread("blockIdx.z")
tx = T.env_thread("threadIdx.x")
ty = T.env_thread("threadIdx.y")
tz = T.env_thread("threadIdx.z")
T.launch_thread(bz, 196)
Conv_wmma_accumulator = T.alloc_buffer((2048,), scope="wmma.accumulator")
Apad_shared = T.alloc_buffer((12288,), "float16", scope="shared")
W_shared = T.alloc_buffer((12288,), "float16", scope="shared")
Apad_shared_wmma_matrix_a = T.alloc_buffer((512,), "float16", scope="wmma.matrix_a")
W_shared_wmma_matrix_b = T.alloc_buffer((1024,), "float16", scope="wmma.matrix_b")
T.launch_thread(bx, 2)
T.launch_thread(by, 4)
T.launch_thread(ty, 4)
T.launch_thread(tz, 2)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 0, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 1, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 2, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 3, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 4, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 5, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 6, T.float32(0), dtype="handle"
)
)
T.evaluate(
T.tvm_fill_fragment(
Conv_wmma_accumulator.data, 16, 16, 16, 7, T.float32(0), dtype="handle"
)
)
for ic_outer in T.serial(0, 8):
for kh in T.serial(0, 3):
for ax2 in T.serial(0, 3):
with T.launch_thread(tx, 32):
Apad_shared[((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61440
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 32)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61408
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 64)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61376
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 96)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61344
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 128)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61312
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 160)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61280
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 192)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61248
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 224)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61216
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 256)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61184
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 288)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61152
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 320)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61120
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 352)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61088
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 384)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61056
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 416)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 61024
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 448)] = (
T.if_then_else(
(
(
(
1 <= (T.floordiv(bz, 14) + kh)
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 60992
),
],
T.float16(0),
dtype="float16",
)
)
T.launch_thread(tx, 32)
Apad_shared[(((((ty * 3072) + (tz * 1536)) + (ax2 * 512)) + tx) + 480)] = (
T.if_then_else(
(
(
(
(1 <= (T.floordiv(bz, 14) + kh))
and ((T.floordiv(bz, 14) + kh) < 15)
)
and (1 <= (ax2 + T.floormod(bz, 14)))
)
and ((ax2 + T.floormod(bz, 14)) < 15)
),
A_1[
(
(
(
(
(
(
(
((bx * 6422528) + (ty * 1605632))
+ (tz * 802816)
)
+ (kh * 57344)
)
+ (bz * 4096)
)
+ (ax2 * 4096)
)
+ (ic_outer * 512)
)
+ tx
)
- 60960
),
],
T.float16(0),
dtype="float16",
)
)
with T.launch_thread(tx, 32):
W_shared[T.ramp((((ty * 512) + (tz * 256)) + (tx * 8)), 1, 8)] = W_1[
T.ramp(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
),
1,
8,
)
]
with T.launch_thread(tx, 32):
W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 2048), 1, 8)] = W_1[
T.ramp(
(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
)
+ 8192
),
1,
8,
)
]
with T.launch_thread(tx, 32):
W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 4096), 1, 8)] = W_1[
T.ramp(
(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
)
+ 131072
),
1,
8,
)
]
with T.launch_thread(tx, 32):
W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 6144), 1, 8)] = W_1[
T.ramp(
(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
)
+ 139264
),
1,
8,
)
]
with T.launch_thread(tx, 32):
W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 8192), 1, 8)] = W_1[
T.ramp(
(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
)
+ 262144
),
1,
8,
)
]
with T.launch_thread(tx, 32):
W_shared[T.ramp(((((ty * 512) + (tz * 256)) + (tx * 8)) + 10240), 1, 8)] = W_1[
T.ramp(
(
(
(
(
(((kh * 393216) + (ic_outer * 16384)) + (by * 2048))
+ (ty * 512)
)
+ (tz * 256)
)
+ (tx * 8)
)
+ 270336
),
1,
8,
)
]
for ic_inner in T.serial(0, 2):
for kw in T.serial(0, 3):
T.evaluate(
T.tvm_load_matrix_sync(
Apad_shared_wmma_matrix_a.data,
16,
16,
16,
0,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
Apad_shared.data,
(((ty * 3072) + (kw * 512)) + (ic_inner * 256)),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_load_matrix_sync(
Apad_shared_wmma_matrix_a.data,
16,
16,
16,
1,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
Apad_shared.data,
((((ty * 3072) + (kw * 512)) + (ic_inner * 256)) + 1536),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_load_matrix_sync(
W_shared_wmma_matrix_b.data,
16,
16,
16,
0,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
W_shared.data,
(((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_load_matrix_sync(
W_shared_wmma_matrix_b.data,
16,
16,
16,
1,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
W_shared.data,
((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 256),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_load_matrix_sync(
W_shared_wmma_matrix_b.data,
16,
16,
16,
2,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
W_shared.data,
((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 512),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_load_matrix_sync(
W_shared_wmma_matrix_b.data,
16,
16,
16,
3,
T.tvm_access_ptr(
T.type_annotation(dtype="float16"),
W_shared.data,
((((kw * 4096) + (ic_inner * 2048)) + (tz * 1024)) + 768),
256,
1,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
0,
Apad_shared_wmma_matrix_a.data,
0,
W_shared_wmma_matrix_b.data,
0,
Conv_wmma_accumulator.data,
0,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
1,
Apad_shared_wmma_matrix_a.data,
0,
W_shared_wmma_matrix_b.data,
1,
Conv_wmma_accumulator.data,
1,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
2,
Apad_shared_wmma_matrix_a.data,
0,
W_shared_wmma_matrix_b.data,
2,
Conv_wmma_accumulator.data,
2,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
3,
Apad_shared_wmma_matrix_a.data,
0,
W_shared_wmma_matrix_b.data,
3,
Conv_wmma_accumulator.data,
3,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
4,
Apad_shared_wmma_matrix_a.data,
1,
W_shared_wmma_matrix_b.data,
0,
Conv_wmma_accumulator.data,
4,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
5,
Apad_shared_wmma_matrix_a.data,
1,
W_shared_wmma_matrix_b.data,
1,
Conv_wmma_accumulator.data,
5,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
6,
Apad_shared_wmma_matrix_a.data,
1,
W_shared_wmma_matrix_b.data,
2,
Conv_wmma_accumulator.data,
6,
dtype="handle",
)
)
T.evaluate(
T.tvm_mma_sync(
Conv_wmma_accumulator.data,
7,
Apad_shared_wmma_matrix_a.data,
1,
W_shared_wmma_matrix_b.data,
3,
Conv_wmma_accumulator.data,
7,
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
0,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
1,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 256
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
2,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 512
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
3,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 768
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
4,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 1605632
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
5,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 1605888
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
6,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 1606144
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
T.evaluate(
T.tvm_store_matrix_sync(
Conv_wmma_accumulator.data,
16,
16,
16,
7,
T.tvm_access_ptr(
T.type_annotation(dtype="float32"),
Conv_1.data,
(
(
((((bx * 12845056) + (ty * 3211264)) + (bz * 8192)) + (by * 2048))
+ (tz * 1024)
)
+ 1606400
),
256,
2,
dtype="handle",
),
16,
"row_major",
dtype="handle",
)
)
return func
def opt_conv_tensorcore_mod_host():
@T.prim_func(s_tir=True)
def opt_conv_tensorcore_mod_host(
args: T.handle,
arg_type_ids: T.Buffer((3,), "int32"),
num_args: T.int32,
out_ret_value: T.handle,
out_ret_tcode: T.handle,
resource_handle: T.handle,
) -> T.int32:
# function attr dict
T.func_attr(
{
"tirx.noalias": True,
"global_symbol": "default_function",
"tirx.is_entry_func": True,
"calling_conv": 1,
}
)
# body
stack_tcode_data: T.let[T.handle("int32")] = T.tvm_stack_alloca(
"arg_tcode", 10, dtype="handle"
)
stack_tcode = T.decl_buffer([9], "int32", data=stack_tcode_data)
stack_value: T.let[T.handle] = T.tvm_stack_alloca("arg_value", 10, dtype="handle")
assert num_args == 3, "default_function: num_args should be 3"
arg0: T.let[T.handle] = T.tvm_struct_get(args, 0, 12, dtype="handle")
arg0_code: T.let[T.int32] = arg_type_ids[0]
arg1: T.let[T.handle] = T.tvm_struct_get(args, 1, 12, dtype="handle")
arg1_code: T.let[T.int32] = arg_type_ids[1]
arg2: T.let[T.handle] = T.tvm_struct_get(args, 2, 12, dtype="handle")
arg2_code: T.let[T.int32] = arg_type_ids[2]
A: T.let[T.handle] = T.tvm_struct_get(arg0, 0, 1, dtype="handle")
T.attr(A, "storage_alignment", 128)
arg0_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg0, 0, 2, dtype=T.handle("int64").ty
)
arg0_shape = T.decl_buffer([6], "int64", data=arg0_shape_data)
arg0_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg0, 0, 3, dtype=T.handle("int64").ty
)
arg0_strides = T.decl_buffer([6], "int64", data=arg0_strides_data)
dev_id: T.let[T.int32] = T.tvm_struct_get(arg0, 0, 9, dtype="int32")
W: T.let[T.handle] = T.tvm_struct_get(arg1, 0, 1, dtype="handle")
T.attr(W, "storage_alignment", 128)
arg1_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg1, 0, 2, dtype=T.handle("int64").ty
)
arg1_shape = T.decl_buffer([6], "int64", data=arg1_shape_data)
arg1_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg1, 0, 3, dtype=T.handle("int64").ty
)
arg1_strides = T.decl_buffer([6], "int64", data=arg1_strides_data)
Conv: T.let[T.handle] = T.tvm_struct_get(arg2, 0, 1, dtype="handle")
T.attr(Conv, "storage_alignment", 128)
arg2_shape_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg2, 0, 2, dtype=T.handle("int64").ty
)
arg2_shape = T.decl_buffer([6], "int64", data=arg2_shape_data)
arg2_strides_data: T.let[T.handle("int64")] = T.tvm_struct_get(
arg2, 0, 3, dtype=T.handle("int64").ty
)
arg2_strides = T.decl_buffer([6], "int64", data=arg2_strides_data)
assert (((arg0_code == 3) or (arg0_code == 13)) or (arg0_code == 7)) or (arg0_code == 4), (
"default_function: Expect arg[0] to be pointer"
)
assert (((arg1_code == 3) or (arg1_code == 13)) or (arg1_code == 7)) or (arg1_code == 4), (
"default_function: Expect arg[1] to be pointer"
)
assert (((arg2_code == 3) or (arg2_code == 13)) or (arg2_code == 7)) or (arg2_code == 4), (
"default_function: Expect arg[2] to be pointer"
)
assert 6 == T.tvm_struct_get(arg0, 0, 4, dtype="int32"), "arg0.ndim is expected to equal 6"
assert 6 == T.tvm_struct_get(arg0, 0, 4, dtype="int32"), "arg0.ndim is expected to equal 6"
assert (
(T.tvm_struct_get(arg0, 0, 5, dtype="uint8") == T.uint8(2))
and (T.tvm_struct_get(arg0, 0, 6, dtype="uint8") == T.uint8(16))
) and (T.tvm_struct_get(arg0, 0, 7, dtype="uint16") == T.uint16(1)), (
"arg0.dtype is expected to be float16"
)
assert 16 == T.cast(arg0_shape[0], "int32"), (
"Argument arg0.shape[0] has an unsatisfied constraint"
)
assert 14 == T.cast(arg0_shape[1], "int32"), (
"Argument arg0.shape[1] has an unsatisfied constraint"
)
assert 14 == T.cast(arg0_shape[2], "int32"), (
"Argument arg0.shape[2] has an unsatisfied constraint"
)
assert 16 == T.cast(arg0_shape[3], "int32"), (
"Argument arg0.shape[3] has an unsatisfied constraint"
)
assert 16 == T.cast(arg0_shape[4], "int32"), (
"Argument arg0.shape[4] has an unsatisfied constraint"
)
assert 16 == T.cast(arg0_shape[5], "int32"), (
"Argument arg0.shape[5] has an unsatisfied constraint"
)
if not (T.isnullptr(arg0_strides.data, dtype="bool")):
assert (
(
(
(
(1 == T.cast(arg0_strides[5], "int32"))
and (16 == T.cast(arg0_strides[4], "int32"))
)
and (256 == T.cast(arg0_strides[3], "int32"))
)
and (4096 == T.cast(arg0_strides[2], "int32"))
)
and (57344 == T.cast(arg0_strides[1], "int32"))
) and (802816 == T.cast(arg0_strides[0], "int32")), (
"arg0.strides: expected to be compact array"
)
T.evaluate(0)
assert T.uint64(0) == T.tvm_struct_get(arg0, 0, 8, dtype="uint64"), (
"Argument arg0.byte_offset has an unsatisfied constraint"
)
assert 2 == T.tvm_struct_get(arg0, 0, 10, dtype="int32"), (
"Argument arg0.device_type has an unsatisfied constraint"
)
assert 6 == T.tvm_struct_get(arg1, 0, 4, dtype="int32"), "arg1.ndim is expected to equal 6"
assert 6 == T.tvm_struct_get(arg1, 0, 4, dtype="int32"), "arg1.ndim is expected to equal 6"
assert (
(T.tvm_struct_get(arg1, 0, 5, dtype="uint8") == T.uint8(2))
and (T.tvm_struct_get(arg1, 0, 6, dtype="uint8") == T.uint8(16))
) and (T.tvm_struct_get(arg1, 0, 7, dtype="uint16") == T.uint16(1)), (
"arg1.dtype is expected to be float16"
)
assert 3 == T.cast(arg1_shape[0], "int32"), (
"Argument arg1.shape[0] has an unsatisfied constraint"
)
assert 3 == T.cast(arg1_shape[1], "int32"), (
"Argument arg1.shape[1] has an unsatisfied constraint"
)
assert 16 == T.cast(arg1_shape[2], "int32"), (
"Argument arg1.shape[2] has an unsatisfied constraint"
)
assert 32 == T.cast(arg1_shape[3], "int32"), (
"Argument arg1.shape[3] has an unsatisfied constraint"
)
assert 16 == T.cast(arg1_shape[4], "int32"), (
"Argument arg1.shape[4] has an unsatisfied constraint"
)
assert 16 == T.cast(arg1_shape[5], "int32"), (
"Argument arg1.shape[5] has an unsatisfied constraint"
)
if not (T.isnullptr(arg1_strides.data, dtype="bool")):
assert (
(
(
(
(1 == T.cast(arg1_strides[5], "int32"))
and (16 == T.cast(arg1_strides[4], "int32"))
)
and (256 == T.cast(arg1_strides[3], "int32"))
)
and (8192 == T.cast(arg1_strides[2], "int32"))
)
and (131072 == T.cast(arg1_strides[1], "int32"))
) and (393216 == T.cast(arg1_strides[0], "int32")), (
"arg1.strides: expected to be compact array"
)
T.evaluate(0)
assert T.uint64(0) == T.tvm_struct_get(arg1, 0, 8, dtype="uint64"), (
"Argument arg1.byte_offset has an unsatisfied constraint"
)
assert 2 == T.tvm_struct_get(arg1, 0, 10, dtype="int32"), (
"Argument arg1.device_type has an unsatisfied constraint"
)
assert dev_id == T.tvm_struct_get(arg1, 0, 9, dtype="int32"), (
"Argument arg1.device_id has an unsatisfied constraint"
)
assert 6 == T.tvm_struct_get(arg2, 0, 4, dtype="int32"), "arg2.ndim is expected to equal 6"
assert 6 == T.tvm_struct_get(arg2, 0, 4, dtype="int32"), "arg2.ndim is expected to equal 6"
assert (
(T.tvm_struct_get(arg2, 0, 5, dtype="uint8") == T.uint8(2))
and (T.tvm_struct_get(arg2, 0, 6, dtype="uint8") == T.uint8(32))
) and (T.tvm_struct_get(arg2, 0, 7, dtype="uint16") == T.uint16(1)), (
"arg2.dtype is expected to be float32"
)
assert 16 == T.cast(arg2_shape[0], "int32"), (
"Argument arg2.shape[0] has an unsatisfied constraint"
)
assert 14 == T.cast(arg2_shape[1], "int32"), (
"Argument arg2.shape[1] has an unsatisfied constraint"
)
assert 14 == T.cast(arg2_shape[2], "int32"), (
"Argument arg2.shape[2] has an unsatisfied constraint"
)
assert 32 == T.cast(arg2_shape[3], "int32"), (
"Argument arg2.shape[3] has an unsatisfied constraint"
)
assert 16 == T.cast(arg2_shape[4], "int32"), (
"Argument arg2.shape[4] has an unsatisfied constraint"
)
assert 16 == T.cast(arg2_shape[5], "int32"), (
"Argument arg2.shape[5] has an unsatisfied constraint"
)
if not (T.isnullptr(arg2_strides.data, dtype="bool")):
assert (
(
(
(
(1 == T.cast(arg2_strides[5], "int32"))
and (16 == T.cast(arg2_strides[4], "int32"))
)
and (256 == T.cast(arg2_strides[3], "int32"))
)
and (8192 == T.cast(arg2_strides[2], "int32"))
)
and (114688 == T.cast(arg2_strides[1], "int32"))
) and (1605632 == T.cast(arg2_strides[0], "int32")), (
"arg2.strides: expected to be compact array"
)
T.evaluate(0)
assert T.uint64(0) == T.tvm_struct_get(arg2, 0, 8, dtype="uint64"), (
"Argument arg2.byte_offset has an unsatisfied constraint"
)
assert 2 == T.tvm_struct_get(arg2, 0, 10, dtype="int32"), (
"Argument arg2.device_type has an unsatisfied constraint"
)
assert dev_id == T.tvm_struct_get(arg2, 0, 9, dtype="int32"), (
"Argument arg2.device_id has an unsatisfied constraint"
)
T.evaluate(T.tvm_struct_set(stack_value, 0, 12, T.cast(2, "int64"), dtype="int32"))
stack_tcode[0] = 0
T.evaluate(T.tvm_struct_set(stack_value, 1, 12, T.cast(dev_id, "int64"), dtype="int32"))
stack_tcode[1] = 0
T.evaluate(
T.tvm_call_packed_lowered(
"__tvm_set_device", stack_value, stack_tcode.data, 0, 2, dtype="int32"
)
)
T.attr(0, "compute_scope", "default_function_compute_")
T.evaluate(T.tvm_struct_set(stack_value, 0, 12, A, dtype="int32"))
stack_tcode[0] = 3
T.evaluate(T.tvm_struct_set(stack_value, 1, 12, W, dtype="int32"))
stack_tcode[1] = 3
T.evaluate(T.tvm_struct_set(stack_value, 2, 12, Conv, dtype="int32"))
stack_tcode[2] = 3
T.evaluate(T.tvm_struct_set(stack_value, 3, 12, T.cast(196, "int64"), dtype="int32"))
stack_tcode[3] = 0
T.evaluate(T.tvm_struct_set(stack_value, 4, 12, T.cast(2, "int64"), dtype="int32"))
stack_tcode[4] = 0
T.evaluate(T.tvm_struct_set(stack_value, 5, 12, T.cast(4, "int64"), dtype="int32"))
stack_tcode[5] = 0
T.evaluate(T.tvm_struct_set(stack_value, 6, 12, T.cast(4, "int64"), dtype="int32"))
stack_tcode[6] = 0
T.evaluate(T.tvm_struct_set(stack_value, 7, 12, T.cast(2, "int64"), dtype="int32"))
stack_tcode[7] = 0
T.evaluate(T.tvm_struct_set(stack_value, 8, 12, T.cast(32, "int64"), dtype="int32"))
stack_tcode[8] = 0
T.evaluate(
T.tvm_call_packed_lowered(
"default_function_kernel0", stack_value, stack_tcode.data, 0, 9, dtype="int32"
)
)
return opt_conv_tensorcore_mod_host
def vthread_func():
@T.prim_func(s_tir=True)
def vthread_func(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, [256], "float32")
C = T.match_buffer(c, [256], "float32")
i0 = T.env_thread("blockIdx.x")
i1 = T.env_thread("threadIdx.x")
i2 = T.env_thread("vthread")
T.launch_thread(i0, 4)
T.launch_thread(i1, 2)
T.launch_thread(i2, 2)
B = T.alloc_buffer((16,), scope="local")
for j in range(16):
B[j] = A[i0 * 64 + i1 * 32 + i2 * 16 + j] + T.float32(1)
for j in range(16):
C[i0 * 64 + i1 * 32 + i2 * 16 + j] = B[j] * T.float32(2)
return vthread_func
def matmul():
@T.prim_func(s_tir=True)
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, [128, 128])
B = T.match_buffer(b, [128, 128])
C = T.match_buffer(c, [128, 128])
for i, j, k in T.grid(128, 128, 128):
with T.sblock("update"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
return matmul
def matmul_original():
@T.prim_func(s_tir=True)
def matmul_original(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, [128, 128])
B = T.match_buffer(b, [128, 128])
C = T.match_buffer(c, [128, 128])
for i, j in T.grid(128, 128):
with T.sblock("init"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.float32(0)
for k in range(128):
with T.sblock("update"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
return matmul_original
def element_wise():
@T.prim_func(s_tir=True)
def element_wise(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
C = T.match_buffer(c, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * T.float32(2)
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + T.float32(1)
return element_wise
def predicate():
@T.prim_func(s_tir=True)
def predicate(b: T.handle, c: T.handle) -> None:
B = T.match_buffer(b, (16, 16), "float32")
C = T.match_buffer(c, (16, 16), "float32")
for i, jo, ji in T.grid(16, 4, 5):
with T.sblock("update"):
vi = T.axis.S(16, i)
vj = T.axis.S(16, jo * 4 + ji)
T.where(jo * 4 + ji < 16)
C[vi, vj] = B[vi, vj] + T.float32(1)
return predicate
def test_module_define():
func1 = tvm.ir.IRModule({"matmul": matmul()})["matmul"]
func2 = tvm.ir.IRModule({"element_wise": element_wise()})["element_wise"]
func3 = tvm.ir.IRModule({"predicate": predicate()})["predicate"]
mod1 = tvm.ir.IRModule({"func1": func1, "func2": func2, "func3": func3})
mod2 = tvm.ir.IRModule({"func1": matmul(), "func2": element_wise(), "func3": predicate()})
tvm.ir.assert_structural_equal(mod1, mod2)
def test_matmul_original():
func = matmul_original()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.SeqStmt)
assert isinstance(rt_func.body.block.body.body.body[0].block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body.body.body[1], tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body.body[1].body.block, tirx.stmt.SBlock)
def test_element_wise():
func = element_wise()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body, tirx.stmt.SeqStmt)
assert isinstance(rt_func.body.block.body[0], tirx.stmt.For)
assert isinstance(rt_func.body.block.body[0].body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body[0].body.body.block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body[1], tirx.stmt.For)
assert isinstance(rt_func.body.block.body[1].body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body[1].body.body.block, tirx.stmt.SBlock)
def test_predicate():
func = predicate()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.For)
assert isinstance(rt_func.body.block.body.body.body.body.block, tirx.stmt.SBlock)
def for_thread_binding():
@T.prim_func(s_tir=True)
def for_thread_binding(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16, 16), "float32")
B = T.match_buffer(b, (16, 16), "float32")
for i in T.thread_binding(0, 16, thread="threadIdx.x"):
for j in T.thread_binding(
0, 16, thread="threadIdx.y", annotations={"attr_key": "attr_value"}
):
A[i, j] = B[i, j] + T.float32(1)
return for_thread_binding
def test_for_thread_binding():
func = for_thread_binding()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body, tirx.stmt.For)
assert rt_func.body.kind == 4
assert rt_func.body.thread_binding.thread_tag == "threadIdx.x"
assert isinstance(rt_func.body.body, tirx.stmt.For)
assert rt_func.body.body.kind == 4
assert rt_func.body.body.thread_binding.thread_tag == "threadIdx.y"
assert rt_func.body.body.annotations["attr_key"] == "attr_value"
def match_buffer_region():
@T.prim_func(s_tir=True)
def match_buffer_region(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16, 16, 16), "float32")
B = T.match_buffer(b, (1), "float32")
for i, j in T.grid(16, 4):
with T.sblock():
vi, vj = T.axis.remap("SS", [i, j])
C = T.match_buffer(A[0:16, vi, vj * 4 : vj * 4 + 4], (16, 1, 4))
for ii in range(4):
with T.sblock():
vii = T.axis.S(4, ii)
D = T.match_buffer(C[vii * 4 : vii * 4 + 4, 0, 0:4], (4, 1, 4))
for i, j in T.grid(4, 4):
B[0] += D[i, 0, j]
return match_buffer_region
def test_match_buffer_region():
func = match_buffer_region()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body, tirx.stmt.SBlockRealize)
root = rt_func.body.block
assert isinstance(root.body, tirx.stmt.For)
assert isinstance(root.body.body, tirx.stmt.For)
assert isinstance(root.body.body.body, tirx.stmt.SBlockRealize)
outer_block = root.body.body.body.block
assert len(outer_block.match_buffers) == 1
buffer_C = outer_block.match_buffers[0].buffer
tvm.ir.assert_structural_equal(buffer_C.shape, [T.int32(16), T.int32(1), T.int32(4)])
assert isinstance(outer_block.body, tirx.stmt.For)
assert isinstance(outer_block.body.body, tirx.stmt.SBlockRealize)
inner_block = outer_block.body.body.block
assert len(inner_block.match_buffers) == 1
buffer_D = inner_block.match_buffers[0].buffer
tvm.ir.assert_structural_equal(buffer_D.shape, [T.int32(4), T.int32(1), T.int32(4)])
def block_elements():
@T.prim_func(s_tir=True)
def block_elements(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16, 16), "float32")
B = T.match_buffer(b, (1, 1), "float32")
with T.sblock("update"):
vi = T.axis.S(1, 0)
T.where(True)
T.reads(A[0:16, 0:16])
T.writes(B[0, 0])
T.sblock_attr({"attr_key": "attr_value"})
C = T.sblock_alloc_buffer((4, 4), dtype="float32")
D = T.match_buffer(A[0:4, 0], (4, 1))
with T.init():
B[0, 0] = T.float32(0)
B[0, 0] = A[0, 0] + B[0, 0] + C[1, 1] + D[2, 0]
return block_elements
def test_block_elements():
func = block_elements()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
assert isinstance(rt_func.body.block, tirx.stmt.SBlock)
assert isinstance(rt_func.body.block.body, tirx.stmt.SBlockRealize)
assert isinstance(rt_func.body.block.body.block, tirx.stmt.SBlock)
block = rt_func.body.block.body.block
assert isinstance(block.body, tirx.stmt.BufferStore)
assert isinstance(block.init, tirx.stmt.BufferStore)
assert len(block.annotations) == 1
assert block.annotations["attr_key"] == "attr_value"
def opaque_block():
@T.prim_func(s_tir=True)
def opaque_block(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16, 16), "float32")
B = T.match_buffer(b, (16, 16), "float32")
for i in range(16):
for j in range(16):
with T.sblock():
T.reads([])
T.writes(A[i, j])
A[i, j] = T.float32(0)
with T.sblock():
T.reads([A[i, 0:16]])
T.writes([B[i, 0:16]])
for j in range(16):
B[i, j] = A[i, j]
return opaque_block
def test_opaque_block():
func = opaque_block()
rt_func = tvm.script.from_source(func.script())
tvm.ir.assert_structural_equal(func, rt_func)
root_block = rt_func.body.block
assert isinstance(root_block, tirx.stmt.SBlock)
assert isinstance(root_block.body, tirx.stmt.For)
assert isinstance(root_block.body.body[0], tirx.stmt.For)
assert isinstance(root_block.body.body[0].body, tirx.stmt.SBlockRealize)
assert isinstance(root_block.body.body[0].body.block, tirx.stmt.SBlock)
assert len(root_block.body.body[0].body.block.iter_vars) == 0
assert isinstance(root_block.body.body[1], tirx.stmt.SBlockRealize)
assert isinstance(root_block.body.body[1].block, tirx.stmt.SBlock)
assert len(root_block.body.body[1].block.iter_vars) == 0
def rank0():
@T.prim_func(s_tir=True)
def rank0(a: T.handle) -> None:
A = T.match_buffer(a, (), "float32")
B = T.sblock_alloc_buffer((), "float32")
A[()] = 2
B[()] = A[()]
return rank0
def rank0_block():
@T.prim_func(s_tir=True)
def rank0_block(a: T.handle) -> None:
A = T.match_buffer(a, (), "float32")
B = T.sblock_alloc_buffer((), "float32")
B[()] = A[()]
with T.sblock("update"):
T.reads([A[()]])
T.writes([B[()]])
for i in range(1):
B[()] = A[()]
return rank0_block
def select():
@T.prim_func(s_tir=True)
def select(a: T.handle) -> None:
A = T.match_buffer(a, (), "float32")
A[()] = T.Select(True, 1, 2)
return select
def minmax():
@T.prim_func(s_tir=True)
def minmax(a: T.handle) -> None:
A = T.match_buffer(a, (), "float32")
A[()] = T.min(1, 2)
A[()] = T.max(1, 2)
return minmax
def abs():
@T.prim_func(s_tir=True)
def abs(a: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
for i, j in T.grid(128, 128):
with T.sblock("A"):
vi, vj = T.axis.remap("SS", [i, j])
A[vi, vj] = T.abs(A[vi, vj])
return abs
def constant_folding():
@T.prim_func(s_tir=True)
def constant_folding(a: T.handle) -> None:
A = T.match_buffer(a, (), "float32")
A[()] = T.min(2.2, 5.2)
A[()] = T.max(T.float32(2.2), T.float32(T.float32(5.2)))
A[()] = T.min(2.2, 5.0)
return constant_folding
def simplify_bracket():
# uninitialized variables
@T.prim_func(check_well_formed=False, s_tir=True)
def simplify_bracket() -> None:
a = T.int32()
b = T.int32()
c = T.int32()
d = T.int32()
T.evaluate(a + b * (c + d))
return simplify_bracket
def var_with_same_name():
@T.prim_func(s_tir=True)
def var_with_same_name(a: T.handle) -> None:
A = T.match_buffer(a, (16, 16), "float32")
for i, j in T.grid(16, 16):
with T.sblock():
vi, vj = T.axis.remap("SS", [i, j])
A[vi, vj] = 0
for i, j in T.grid(16, 16):
with T.sblock():
vi, vj = T.axis.remap("SS", [i, j])
A[vi, vj] = 0
return var_with_same_name
def test_same_name_var():
func = var_with_same_name()
out_str = func.script()
rt_func = tvm.script.from_source(out_str)
tvm.ir.assert_structural_equal(func, rt_func)
assert out_str.count("for i, j in T.grid(16, 16)") == 2
assert out_str.find("i_") == -1
assert out_str.find("i_") == -1
def while_loop():
@T.prim_func(s_tir=True)
def while_loop(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16,), "float32")
B = T.match_buffer(b, (16,), "float32")
i = T.sblock_alloc_buffer((), "int32", scope="local")
for ii in range(16):
with T.sblock():
vi = T.axis.S(16, ii)
B[vi] = 0
while i[()] < 10:
for j in range(16):
B[j] += A[j]
return while_loop
# fmt: off
def primfunc_with_allocate_annotations():
@T.prim_func(s_tir=True)
def primfunc_with_allocate_annotations(placeholder_28: T.handle, T_cast_6: T.handle) -> None:
# function attr dict
T.func_attr({"global_symbol": "tvmgen_default_fused_nn_max_pool2d_cast", "tirx.noalias": True})
placeholder_29 = T.match_buffer(placeholder_28, [802816], dtype="uint8", elem_offset=0, align=64, offset_factor=1)
T_cast_7 = T.match_buffer(T_cast_6, [200704], dtype="int16", elem_offset=0, align=64, offset_factor=1)
# body
tensor_2 = T.alloc_buffer((200704,), "uint8", annotations={"attr1_key": "attr1_value"})
for ax0_ax1_fused_4 in T.serial(0, 56):
for ax2_4 in T.serial(0, 56):
for ax3_init in T.serial(0, 64):
tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_init)] = T.uint8(0)
for rv0_rv1_fused_1, ax3_2 in T.grid(9, 64):
tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_2)] = T.max(tensor_2[(((ax0_ax1_fused_4*3584) + (ax2_4*64)) + ax3_2)], T.if_then_else(((((ax0_ax1_fused_4*2) + T.floordiv(rv0_rv1_fused_1, 3)) < 112) and (((ax2_4*2) + T.floormod(rv0_rv1_fused_1, 3)) < 112)), placeholder_29[(((((ax0_ax1_fused_4*14336) + (T.floordiv(rv0_rv1_fused_1, 3)*7168)) + (ax2_4*128)) + (T.floormod(rv0_rv1_fused_1, 3)*64)) + ax3_2)], T.uint8(0), dtype="uint8"))
for ax0_ax1_fused_5 in T.serial(0, 56):
for ax2_5, ax3_3 in T.grid(56, 64):
T_cast_7[(((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3)] = T.cast(tensor_2[(((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3)], "int16")
return primfunc_with_allocate_annotations
# fmt: on
# fmt: off
def comm_reducer_single_reduce_group():
@T.prim_func(s_tir=True)
def comm_reducer_single_reduce_group(a: T.handle, b: T.handle) -> None:
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
threadIdx_x = T.env_thread("threadIdx.x")
A = T.match_buffer(a, [16384], dtype="float32")
for i in T.serial(0, 128):
T.launch_thread(threadIdx_x, 128)
reduce_temp0 = T.alloc_buffer((1,), scope="local")
with T.attr(T.comm_reducer(lambda x, y: x + y, [T.float32(0)]), "reduce_scope", T.int32(0)):
T.evaluate(T.tvm_thread_allreduce(T.uint32(1), A[i * 128 + threadIdx_x], True, reduce_temp0.data, threadIdx_x, dtype="handle"))
return comm_reducer_single_reduce_group
def comm_reducer_multiple_reduce_groups():
@T.prim_func(s_tir=True)
def comm_reducer_multiple_reduce_groups(a: T.handle, b: T.handle) -> None:
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
threadIdx_x = T.env_thread("threadIdx.x")
A = T.match_buffer(a, [16384], dtype="float32")
for i in T.serial(0, 128):
T.launch_thread(threadIdx_x, 128)
reduce_temp0 = T.alloc_buffer((1,), scope="local")
with T.attr(T.comm_reducer(lambda x0, x1, y0, y1: (T.Select((x1 >= y1), x0, y0), T.Select((x1 >= y1), x1, y1)), [T.int32(-1), T.min_value("float32")]), "reduce_scope", T.int32(0)):
T.evaluate(T.tvm_thread_allreduce(T.uint32(1), A[i * 128 + threadIdx_x], True, reduce_temp0.data, threadIdx_x, dtype="handle"))
return comm_reducer_multiple_reduce_groups
def multiple_commreducer():
# normal_reduce_temp0 is treated as uninitialized value
@T.prim_func(check_well_formed=False, s_tir=True)
def multiple_commreducer() -> None:
normal_reduce_temp0 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
normal_reduce_temp1 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
reduce_temp0 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
reduce_temp1 = T.Buffer([1], dtype="float32", strides=[1], scope="local")
for ax0_1 in T.thread_binding(0, 32, thread="threadIdx.x"):
with T.sblock("T_softmax_maxelem_cross_thread_reduction"):
T.attr(T.comm_reducer(lambda x, y: T.max(x, y), [T.min_value("float32")]), "reduce_scope", T.int32(0))
T.evaluate(T.tvm_thread_allreduce(T.uint32(1), normal_reduce_temp0[0], True, reduce_temp0.data, ax0_1, dtype="handle"))
for ax0_1 in T.thread_binding(0, 32, thread="threadIdx.x"):
with T.sblock("T_softmax_expsum_cross_thread_reduction"):
T.attr(T.comm_reducer(lambda x, y: x + y, [T.float32(0)]), "reduce_scope", T.int32(0))
T.evaluate(T.tvm_thread_allreduce(T.uint32(1), normal_reduce_temp1[0], True, reduce_temp1.data, ax0_1, dtype="handle"))
return multiple_commreducer
# fmt: on
def func_div_mod():
# not well-formed: free variables
@T.prim_func(check_well_formed=False, s_tir=True)
def func_div_mod():
a = T.int32()
b = T.int32()
T.evaluate(a // b)
T.evaluate(a % b)
T.evaluate(T.truncmod(a, b))
return func_div_mod
def test_div_mod():
func = func_div_mod()
rt_func = tvm.script.from_source(func.script(), check_well_formed=False)
tvm.ir.assert_structural_equal(func, rt_func, True)
assert isinstance(func.body[0].value, tvm.tirx.FloorDiv)
assert isinstance(func.body[1].value, tvm.tirx.FloorMod)
assert isinstance(func.body[2].value, tvm.tirx.Mod)
def loop_extent_dependent():
@T.prim_func(s_tir=True)
def loop_extent_dependent(a: T.handle) -> None:
A = T.match_buffer(a, [], dtype="int32")
for i in T.serial(0, 128):
for j in T.serial(0, i):
A[()] = A[()] + j
return loop_extent_dependent
def nontrivial_range_axis():
@T.prim_func(s_tir=True)
def nontrivial_range_axis(a: T.handle) -> None:
A = T.match_buffer(a, (10), "float32")
for i in range(10):
with T.sblock("block"):
vi = T.axis.spatial((1, 11), i + 1)
A[vi - 1] = A[vi - 1] + 1.0
return nontrivial_range_axis
def func_with_target_spec_by_config():
@T.prim_func(s_tir=True)
def func_with_target_spec_by_config() -> None:
T.func_attr(
{
"kTarget": T.target(
{
"max_num_threads": 1024,
"arch": "sm_70",
"thread_warp_size": 32,
"kind": "cuda",
"tag": "",
"keys": ["cuda", "gpu"],
"host": T.target({"kind": "llvm", "tag": "", "keys": ["cpu"]}),
}
)
}
)
T.evaluate(0)
return func_with_target_spec_by_config
def func_with_target_spec_by_str():
@T.prim_func(s_tir=True)
def func_with_target_spec_by_str() -> None:
T.func_attr({"kTarget": T.target("nvidia/nvidia-a100")})
T.evaluate(0)
return func_with_target_spec_by_str
def func_with_target_and_host_spec_by_str():
@T.prim_func(s_tir=True)
def func():
T.func_attr({"target": T.target("nvidia/nvidia-a100", host="llvm")})
T.evaluate(0)
return func
def func_root_attr():
@T.prim_func(s_tir=True)
def func_root_attr():
with T.sblock("root"):
T.sblock_attr({"a": "0"})
T.evaluate(0)
return func_root_attr
def func_trivial_root_block():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(1, "int32")):
with T.sblock("root"):
A[0] = 0
return func
def func_nested_root_block():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(1, "int32")):
with T.sblock("root"):
with T.sblock("block"):
A[0] = 0
return func
def func_T_ptr_let_statement():
@T.prim_func(s_tir=True)
def func_T_ptr_let_statement(
args: T.handle, arg_type_ids_handle: T.handle("int32"), num_args: T.int32
) -> None:
# The T.Ptr declaration in the parameter list should parse
# correctly, and should be usable as the data pointer in a buffer.
arg_type_ids = T.decl_buffer([2], dtype="int32", data=arg_type_ids_handle)
arg0: T.let[T.handle] = T.tvm_struct_get(args, 0, 12, dtype="handle")
arg1: T.let[T.handle] = T.tvm_struct_get(args, 1, 12, dtype="handle")
# The ABI field is an opaque pointer. Retag it explicitly before
# binding it to the buffer's exact element pointer type.
A_data: T.let[T.handle("float32")] = T.reinterpret(
T.handle("float32").ty,
T.tvm_struct_get(arg0, 0, 1, dtype="handle"),
)
# The buffer declaration has a data pointer defined earlier in
# this function. It should only be defined after the data pointer
# has been defined, and should not be hoisted into the header of
# the function as other buffer_decl statements can be.
A = T.decl_buffer([1024], dtype="float32", data=A_data)
B_data: T.let[T.handle("float32")] = T.reinterpret(
T.handle("float32").ty,
T.tvm_struct_get(arg1, 0, 1, dtype="handle"),
)
B = T.decl_buffer([1024], dtype="float32", data=B_data)
B[0] = A[0]
return func_T_ptr_let_statement
def func_T_ptr_allocate():
@T.prim_func(s_tir=True)
def func_T_ptr_allocate() -> None:
A = T.alloc_buffer((1024,))
A[0] = 0.0
return func_T_ptr_allocate
def llvm_intrin_call():
@T.prim_func(s_tir=True)
def ctpop(A: T.Buffer((16,), "uint8"), B: T.Buffer((16,), "uint8")) -> None:
for i in range(0, 16):
with T.sblock("A"):
vi = T.axis.remap(
"S",
[
i,
],
)
B[vi] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.ctpop.i8"),
A[vi],
dtype="uint8",
)
return ctpop
def parse_bufferslice_as_range_bound():
# apparently the use of i in the "outer" block when it is defined outside of a block is wrong
@T.prim_func(check_well_formed=False, s_tir=True)
def segment_sum(
A_ptr: T.handle, B_ptr: T.handle, indptr_ptr: T.handle, n: T.int32, m: T.int32
) -> None:
A = T.match_buffer(A_ptr, [m], dtype="float32")
B = T.match_buffer(B_ptr, [n], dtype="float32")
indptr = T.match_buffer(indptr_ptr, [n + 1], dtype="int32")
for i in T.serial(n):
with T.sblock("outer"):
vi = T.axis.spatial(n, i)
T.reads(indptr[i : i + 2], B[vi], A[indptr[i] : indptr[i + 1]])
T.writes(B[vi])
for j in T.serial(indptr[i], indptr[i + 1]):
with T.sblock("inner"):
vj = T.axis.reduce(m, j)
T.reads(B[vi], A[vj])
T.writes(B[vi])
with T.init():
B[vi] = T.float32(0)
B[vi] = B[vi] + A[vj]
return segment_sum
def int64_support():
@T.prim_func(s_tir=True)
def elementwise_shape_int64(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (T.int64(128), T.int64(128)), dtype="float32")
B = T.sblock_alloc_buffer((T.int64(128), T.int64(128)), dtype="float32")
C = T.match_buffer(c, (T.int64(128), T.int64(128)), dtype="float32")
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(T.int64(128), T.int64(128)):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + 1.0
return elementwise_shape_int64
def string_annotation_escaping():
@T.prim_func(s_tir=True)
def string_annotation_of_special_chars():
T.func_attr(
{
"key1": '"\'hello\t\r"',
"key2": """
%1 = add i32 %0, %0
%2 = add i32 %0, %1
%3 = add i32 %1, %2
""",
}
)
T.evaluate(0)
return string_annotation_of_special_chars
def pointer_type():
@T.prim_func(s_tir=True)
def func_with_ptr_type_annotations(x: T.handle("int32"), y: T.handle("int32", "shared")):
xx = T.alloc_buffer((16,), "int32")
yy = T.alloc_buffer((16,), "int32", scope="shared")
a: T.let[T.handle("int32")] = T.address_of(xx[0], dtype="handle")
b: T.let[T.handle("int32", "shared")] = T.address_of(yy[0], dtype="handle")
T.evaluate(T.call_extern("copy", a, b, dtype=""))
return func_with_ptr_type_annotations
def buffer_axis_separator():
@T.prim_func(s_tir=True)
def element_wise(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32", axis_separators=[1])
C = T.match_buffer(c, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32", axis_separators=[1])
for i, j in T.grid(128, 128):
with T.sblock("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * T.float32(2)
for i, j in T.grid(128, 128):
with T.sblock("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = B[vi, vj] + T.float32(1)
return element_wise
def buffer_ramp_access_as_slice_index():
@T.prim_func(s_tir=True)
def buffer_ramp_access(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128,), "float32")
B = T.match_buffer(b, (128,), "float32")
C = T.match_buffer(c, (128,), "float32")
for i in range(128):
A[i : i + 1 : 1] = i
for i in range(4):
B[i * 32 : i * 32 + 32] = A[i * 32 : i * 32 + 32 : 1] + T.broadcast(1.0, 32)
for i in range(4):
C[i : i + 128 : 4] = B[i : i + 128 : 4] + T.broadcast(1.0, 32)
return buffer_ramp_access
def ramp_int64():
@T.prim_func(s_tir=True)
def func() -> None:
T.evaluate(T.Ramp(T.int64(0), 1, 3))
return func
def scalable_vectors():
@T.prim_func(s_tir=True)
def func(a: T.handle):
A = T.match_buffer(a, (200,), "float32")
A[T.Ramp(11, 2, 4 * tirx.vscale())] = T.Broadcast(125, 4 * tirx.vscale())
return func
def predicated_buffer_load_store():
@T.prim_func(s_tir=True)
def func(a: T.handle, b: T.handle):
A = T.match_buffer(a, (4,), "float32")
B = T.match_buffer(b, (8,), "float32")
for i_0 in range(4):
load_a = T.meta_var(
A.vload([T.Ramp(i_0, 1, 4)], predicate=T.Broadcast(T.bool(True), 4))
)
B.vstore([T.Ramp(0, 2, 4)], load_a, predicate=T.Broadcast(T.bool(True), 4))
return func
def let_expression():
@T.prim_func(s_tir=True)
def func():
x = T.int32()
T.evaluate(T.Let(x + 1, where={x: 1}))
return func
def test_void_ptr_vs_handle():
"""An untyped handle is the canonical void-pointer type."""
# Generates PointerType(PrimType::Void())
@T.prim_func(s_tir=True)
def void_ptr(out_ret_value: T.handle("void")):
T.evaluate(out_ret_value)
# Generates PointerType::VoidPointerTy()
@T.prim_func(s_tir=True)
def handle(out_ret_value: T.handle):
T.evaluate(out_ret_value)
tvm.ir.assert_structural_equal(void_ptr.params[0].ty, handle.params[0].ty)
script = void_ptr.script()
assert "out_ret_value: T.handle" in script
assert 'T.handle("void")' not in script
tvm.ir.assert_structural_equal(void_ptr, tvm.script.from_source(script))
@T.prim_func(s_tir=True)
def scoped_void_ptr(out_ret_value: T.handle("void", "shared")):
T.evaluate(out_ret_value)
scoped_script = scoped_void_ptr.script()
assert 'out_ret_value: T.handle(storage_scope="shared")' in scoped_script
assert 'T.handle("void"' not in scoped_script
tvm.ir.assert_structural_equal(scoped_void_ptr, tvm.script.from_source(scoped_script))
def void_ptr():
@T.prim_func(s_tir=True)
def func(out_ret_value: T.handle("void")):
T.evaluate(out_ret_value)
return func
def decl_buffer():
@T.prim_func(s_tir=True)
def func(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")) -> None:
A_flattened = T.decl_buffer(data=A.data, shape=(256,), dtype="float32")
B_flattened = T.decl_buffer(data=B.data, shape=(256,), dtype="float32")
C_alias = T.decl_buffer(data=A_flattened.data, shape=(256,), dtype="float32")
for i in range(256):
B_flattened[i] = A_flattened[i] + C_alias[i] + T.float32(1.0)
return func
def allocate_and_decl_buffer():
@T.prim_func(s_tir=True)
def func(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")) -> None:
D = T.alloc_buffer((16,))
for i in range(4):
C = T.alloc_buffer((4,))
for j in range(4):
C[j] = A[i * 4 + j] + T.float32(1.0)
for j in range(4):
D[j] = C[j]
for j in range(4):
B[i * 4 + j] = D[j]
return func
def alloc_buffer_example():
@T.prim_func(s_tir=True)
def func(a: T.handle, c: T.handle):
A = T.match_buffer(a, (128,), "float32")
C = T.match_buffer(c, (128,), "float32")
B = T.alloc_buffer((128,), "float32")
for i in range(128):
B[i] = A[i] * T.float32(2)
for i in range(128):
C[i] = B[i] + T.float32(1)
return func
def float_infinity():
@T.prim_func(s_tir=True)
def func(
placeholder: T.Buffer((1, 512, 768), "float32"), T_isinf: T.Buffer((1, 512, 768), "bool")
) -> None:
# function attr dict
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
# body
# with T.sblock("root")
for i0, i1, i2 in T.grid(1, 512, 768):
with T.sblock("T_isinf"):
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
T.reads(placeholder[ax0, ax1, ax2])
T.writes(T_isinf[ax0, ax1, ax2])
T_isinf[ax0, ax1, ax2] = T.fabs(
placeholder[ax0, ax1, ax2], dtype="float32"
) == T.float32("inf") and not (T.isnan(placeholder[ax0, ax1, ax2], dtype="bool"))
return func
def minimal_i32_literal():
@T.prim_func(s_tir=True)
def func() -> None:
T.evaluate(T.int32(-2147483648))
T.evaluate(-T.int64(2147483648))
return func
def boolean_argument():
@T.prim_func(s_tir=True)
def func(a: T.boolean) -> None:
T.evaluate(a)
return func
def bool_argument():
@T.prim_func(s_tir=True)
def func(a: T.bool) -> None:
T.evaluate(a)
return func
def bool_variable_annotation():
@T.prim_func(s_tir=True)
def func() -> None:
a: T.let[T.bool] = T.call_extern("dummy", dtype="bool")
T.evaluate(0)
return func
def return_none():
@T.prim_func(s_tir=True)
def func():
T.evaluate(0)
return func
def bool_primitive():
@T.prim_func(s_tir=True)
def func() -> None:
T.evaluate(T.bool(True))
return func
def bool_cast():
# uninitialized var
@T.prim_func(check_well_formed=False, s_tir=True)
def func() -> None:
a = T.bool()
T.evaluate(T.bool(T.int32(0)))
T.evaluate(a == T.bool(False))
return func
def implicit_evaluate():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(1, "int32")):
T.evaluate(T.assume(A[0] == 5))
A[0] = 10
return func
def if_true_else():
@T.prim_func(s_tir=True)
def func() -> None:
if True:
T.evaluate(0)
else:
T.evaluate(1)
return func
def elif_chain_without_else():
@T.prim_func(s_tir=True)
def func(i: T.int32) -> None:
if i == 0:
T.evaluate(0)
elif i == 1:
T.evaluate(1)
elif i == 2:
T.evaluate(2)
return func
def elif_chain_with_else():
@T.prim_func(s_tir=True)
def func(i: T.int32) -> None:
if i == 0:
T.evaluate(0)
elif i == 1:
T.evaluate(1)
elif i == 2:
T.evaluate(2)
else:
T.evaluate(3)
return func
def nested_boolean_expressions():
expressions = {
"and_lhs_and": lambda i, j, k: tirx.all(tirx.all(i, j), k),
"and_rhs_and": lambda i, j, k: tirx.all(i, tirx.all(j, k)),
"and_lhs_or": lambda i, j, k: tirx.all(tirx.any(i, j), k),
"and_rhs_or": lambda i, j, k: tirx.all(i, tirx.any(j, k)),
"or_lhs_and": lambda i, j, k: tirx.any(tirx.all(i, j), k),
"or_rhs_and": lambda i, j, k: tirx.any(i, tirx.all(j, k)),
"or_lhs_or": lambda i, j, k: tirx.any(tirx.any(i, j), k),
"or_rhs_or": lambda i, j, k: tirx.any(i, tirx.any(j, k)),
"and_of_ors": lambda i, j, k: tirx.all(
tirx.any(i, j), tirx.any(j, k), tirx.any(i, k), i, j, k
),
"or_of_ands": lambda i, j, k: tirx.any(
tirx.all(i, j), tirx.all(j, k), tirx.all(i, k), i, j, k
),
}
def make_ir_generator(name, expression):
def inner():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(1, "bool"), i: T.bool, j: T.bool, k: T.bool):
A[0] = expression(i, j, k)
return func
inner.__name__ = f"nested_boolean_expr_{name}"
return inner
for name, expression in expressions.items():
generator = make_ir_generator(name, expression)
yield generator
def multi_env_threads():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(128, "float32"), C: T.Buffer(128, "float32")):
B = T.sblock_alloc_buffer([128], dtype="float32")
for i in T.thread_binding(128, thread="threadIdx.x"):
B[i] = A[i] + 1.0
for i in T.thread_binding(128, thread="threadIdx.x"):
C[i] = B[i] + 2.0
mod = tvm.s_tir.transform.LowerOpaqueBlock()(
tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
)
return mod["main"]
def intrinsic_pow():
@T.prim_func(s_tir=True)
def func():
T.pow(T.float32(1), T.float32(1))
return func
def bind_var():
@T.prim_func(s_tir=True)
def func():
x = T.bind(0)
y = T.bind(0)
T.evaluate(0)
T.evaluate(0)
return func
def string_stride():
@T.prim_func(s_tir=True)
def main(a: T.handle, b: T.handle):
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
n = T.int32()
A = T.match_buffer(a, (n,), strides=("A_s0",), buffer_type="auto")
B = T.match_buffer(b, (n,), strides=("B_s0",), buffer_type="auto")
blockIdx_x = T.launch_thread("blockIdx.x", (n + 63) // 64)
threadIdx_x = T.launch_thread("threadIdx.x", 64)
if T.likely(blockIdx_x * 64 + threadIdx_x < n):
B2 = T.decl_buffer((B.strides[0] * n,), data=B.data)
A2 = T.decl_buffer((A.strides[0] * n,), data=A.data)
B2[(blockIdx_x * 64 + threadIdx_x) * B.strides[0]] = A2[
(blockIdx_x * 64 + threadIdx_x) * A.strides[0]
] * T.float32(2)
return main
def string_stride_int64():
@T.prim_func(s_tir=True)
def main(a: T.handle, b: T.handle):
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
n = T.int64()
A_s0 = T.int64()
B_s0 = T.int64()
A = T.match_buffer(a, (n,), strides=(A_s0,), buffer_type="auto")
B = T.match_buffer(b, (n,), strides=(B_s0,), buffer_type="auto")
for i in range(n):
B[i] = A[i]
return main
def merge_shape_var_def():
# uninitialized vars
@T.prim_func(check_well_formed=False, s_tir=True)
def main(A: T.handle, B: T.handle):
# fmt: off
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
m, n = T.int32(), T.int32()
A_1 = T.match_buffer(A, (m, n), strides=("A_1_s0", "A_1_s1"), buffer_type="auto")
B_1 = T.match_buffer(B, (m, n), strides=("B_1_s0", "B_1_s1"), buffer_type="auto")
for i_outer, j_outer, i_inner in T.grid((m + 9) // 10, (n + 4) // 5, 10):
if T.likely(i_outer * 10 + i_inner < m):
for j_inner in range(5):
if T.likely(j_outer * 5 + j_inner < n):
cse_v2: T.let[T.int32] = j_outer * 5 + j_inner
cse_v1: T.let[T.int32] = i_outer * 10 + i_inner
B_2 = T.decl_buffer(
(B_1.strides[0] * m,),
data=B_1.data,
strides=("B_2_s0",),
buffer_type="auto",
)
A_2 = T.decl_buffer(
(A_1.strides[0] * m,),
data=A_1.data,
strides=("A_2_s0",),
buffer_type="auto",
)
B_2[cse_v1 * B_1.strides[0] + cse_v2 * B_1.strides[1]] = A_2[
cse_v1 * A_1.strides[0] + cse_v2 * A_1.strides[1]
]
# fmt: on
return main
def if_then_else_var():
@T.prim_func(s_tir=True)
def main(n: T.int32):
if n == 0:
x = 5
T.evaluate(x)
else:
x = 10
T.evaluate(x)
return main
def tvm_shfl_builtins():
@T.prim_func(s_tir=True)
def func(
A: T.handle("float32"),
B: T.handle("float32"),
C: T.handle("float32"),
):
blockIdx_x = T.launch_thread("blockIdx.x", 1)
threadIdx_x = T.launch_thread("threadIdx.x", 32)
A_warp = T.alloc_buffer((1,), scope="local")
B_warp = T.alloc_buffer((1,), scope="local")
red_buf0 = T.alloc_buffer((1,), scope="local")
A_warp_1 = T.decl_buffer((32,), data=A_warp.data, scope="local")
A_1 = T.decl_buffer((32,), data=A) # A is a handle param
A_warp_1[0] = A_1[threadIdx_x]
B_warp_1 = T.decl_buffer((32,), data=B_warp.data, scope="local")
T.tvm_storage_sync("warp")
B_warp_1[0] = T.tvm_warp_shuffle(
T.tvm_warp_activemask(), A_warp_1[0], threadIdx_x % 4 * 8 + threadIdx_x // 4, 32, 32
) + T.float32(1)
red_buf0_1 = T.decl_buffer((1,), data=red_buf0.data, scope="local")
with T.attr(
T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
"reduce_scope",
T.int32(0),
):
mask = T.alloc_buffer((1,), "uint32", scope="local")
t0 = T.alloc_buffer((1,), scope="local")
red_buf0_1[0] = A_warp_1[0]
mask_1 = T.decl_buffer((1,), "uint32", data=mask.data, scope="local")
mask_1[0] = T.tvm_warp_activemask()
t0_1 = T.decl_buffer((1,), data=t0.data, scope="local")
t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 16, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 8, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 4, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 2, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
t0_1[0] = T.tvm_warp_shuffle_down(mask_1[0], red_buf0_1[0], 1, 32, 32)
red_buf0_1[0] = red_buf0_1[0] + t0_1[0]
red_buf0_1[0] = T.tvm_warp_shuffle(mask_1[0], red_buf0_1[0], 0, 32, 32)
# NOTE(Zihao): test tvm_warp_shuffle_up
red_buf0_1[0] = T.tvm_warp_shuffle_up(mask_1[0], red_buf0_1[0], 0, 32, 32)
if threadIdx_x == 0:
C_1 = T.decl_buffer((1,), data=C)
C_1[0] = red_buf0_1[0]
B_1 = T.decl_buffer((32,), data=B)
B_1[threadIdx_x] = B_warp_1[0]
return func
def make_packed_api_result():
@T.prim_func(s_tir=True)
def func(A: T.Buffer(64, "float32")):
T.func_attr({"global_symbol": "main", "target": T.target("cuda")})
bx = T.launch_thread("blockIdx.x", 64)
T.evaluate(A[bx])
mod = tvm.IRModule.from_expr(func)
return tvm.tirx.transform.MakePackedAPI()(mod)
def tvm_struct_set_generated_in_cpp():
"""Ensure same dtype for tvm_struct_set in Python/C++
The TVMStructSet method in C++, used internally by
LowerTVMBuiltin, and the Python method `T.tvm_struct_set`, used
when parsing TVMScript should use the same dtype "int32".
"""
@I.ir_module(s_tir=True)
class Module:
@T.prim_func(s_tir=True)
def tir_packed_call(A: T.Buffer(16)):
T.attr(0, "device_id", 0)
T.attr(0, "device_type", 0)
T.evaluate(
T.tvm_call_cpacked(
"tvm_test_cpacked",
T.tvm_stack_make_array(
A.data,
T.tvm_stack_make_shape(16, dtype="handle"),
T.reinterpret(T.uint64(0), dtype="handle"),
T.uint32(1),
T.Cast("float32", 0),
0,
dtype="handle",
),
dtype="int32",
)
)
return tvm.tirx.transform.LowerTVMBuiltin()(Module)
def ir_module_with_attrs():
@I.ir_module(s_tir=True)
class Module:
I.module_attrs({"attr": 10})
@T.prim_func(s_tir=True)
def tir_func(A: T.Buffer(16, "int32"), B: T.Buffer(16, "int32")):
for i in range(16):
B[i] = A[i]
return Module
def nested_seqstmt():
"""Nested SeqStmt should be normalized to flat SeqStmt
Nested SeqStmt are representable in the TIR structures, but are
flattened when converted to TVMScript. Previously, this could
cause failures to round-trip through TVMScript, including
erroneous use of TVMScript's concise-scoping rules. This was
resolved by normalizing nested SeqStmt in TIR, such that the use
of `tirx.SeqStmt` below results in a single flat `tirx.SeqStmt`
containing the three `tirx.Evaluate` calls.
"""
func = tvm.tirx.PrimFunc(
params=[],
body=tvm.tirx.SeqStmt(
[
tvm.tirx.SeqStmt([tvm.tirx.Evaluate(0), tvm.tirx.Evaluate(1)]),
tvm.tirx.Evaluate(2),
]
),
)
return func
def subroutine_call():
"""A GlobalVar may reference other functions in the module"""
@I.ir_module(s_tir=True)
class mod:
@T.prim_func(s_tir=True)
def main(A: T.Buffer(16, "float32")):
mod.subroutine(A.data, T.int32(16))
@T.prim_func(s_tir=True)
def subroutine(A_data: T.handle("float32"), n: T.int32):
T.evaluate(0)
return mod
def subroutine_call_returning_int():
"""An internal function call may return non-void"""
@I.ir_module(s_tir=True)
class mod:
@T.prim_func(s_tir=True)
def main(A: T.Buffer(2, "float32")):
mod.subroutine(A[0]) + mod.subroutine(A[1])
@T.prim_func(s_tir=True)
def subroutine(x: T.float32) -> T.float32:
T.ret(x * x)
return mod
def undefined_data_ptr_in_decl_buffer():
"""The T.decl_buffer syntax should not introduce an Allocate
While T.decl_buffer can be used to represent an
Allocate/DeclBuffer pair, performing a round-trip through
TVMScript should not introduce an Allocate node.
"""
# uninitialized var
@T.prim_func(check_well_formed=False, s_tir=True)
def func():
data_ptr = T.handle("float32")
buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr)
T.evaluate(buf[0])
return func
def undefined_shape_in_decl_buffer():
# uninitialized var
@T.prim_func(check_well_formed=False, s_tir=True)
def func():
size = T.int32()
buf = T.decl_buffer(shape=[size], dtype="float32")
T.evaluate(buf[0])
return func
def undefined_stride_in_decl_buffer():
# uninitialized var
@T.prim_func(check_well_formed=False, s_tir=True)
def func():
stride = T.int32()
data_ptr = T.handle("float32")
buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr, strides=[stride])
T.evaluate(buf[0])
return func
def undefined_elem_offset_in_decl_buffer():
# uninitialized var
@T.prim_func(check_well_formed=False, s_tir=True)
def func():
elem_offset = T.int32()
data_ptr = T.handle("float32")
buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr, elem_offset=elem_offset)
T.evaluate(buf[0])
return func
def subroutine_call_without_arguments():
@I.ir_module(s_tir=True)
class mod:
@T.prim_func(s_tir=True)
def main():
# Should be equivalent to the bare "mod.subroutine()", but
# that relies on `GlobalVar.__call__` returning the
# correct IR type.
tirx.call_tir(mod.subroutine)
@T.prim_func(s_tir=True)
def subroutine():
T.evaluate(0)
return mod
def return_zero():
@T.prim_func(s_tir=True)
def func() -> T.int32:
T.ret(0)
return func
def return_zero_private():
@T.prim_func(private=True, s_tir=True)
def func() -> T.int32:
T.ret(0)
return func
def return_zero_private_with_attr():
@T.prim_func(private=True, s_tir=True)
def func() -> T.int32:
T.func_attr({"greeting": "hello"})
T.ret(0)
return func
def func_attr_with_list():
@T.prim_func(s_tir=True)
def func(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
D: T.Buffer((128, 128), "float32"),
) -> None:
T.func_attr({"global_symbol": "main", "tirx.noalias": True, "layout_free_buffers": [1]})
C = T.sblock_alloc_buffer([128, 128], dtype="float32")
for i0, i1, i2 in T.grid(128, 128, 128):
with T.sblock("C"):
x, y, k = T.axis.remap("SSR", [i0, i1, i2])
with T.init():
C[x, y] = T.float32(0)
C[x, y] = C[x, y] + A[x, k] * B[y, k]
for i0, i1 in T.grid(128, 128):
with T.sblock("D"):
T.sblock_attr({"layout_free_placeholders": [C]})
x, y = T.axis.remap("SS", [i0, i1])
D[x, y] = C[x, y] + T.float32(1)
return func
def func_with_loop_jumps():
@T.prim_func(s_tir=True)
def func(In: T.Buffer((1,), "int32"), Out: T.Buffer((2,), "int32")):
Out[0] = 0
Out[1] = 0
for i in range(1000):
if i % 13 == 0:
Out[1] = Out[1] + 1
continue
Out[0] = Out[0] + 1
if Out[0] >= In[0]:
break
return func
def func_with_loop_steps():
@T.prim_func(s_tir=True)
def func(
A: T.Buffer((1024,)), B: T.Buffer((1024,)), C: T.Buffer((1024,)), tid: T.int32, v: T.int32
):
for i in T.serial(tid, 1024, step=2):
C[i] = A[i] + B[i]
for i in T.unroll(tid, 1024, step=3):
C[i] = A[i] + B[i]
for i in T.vectorized(tid, 1024, step=4):
C[i] = A[i] + B[i]
for i in T.parallel(tid, 1024, step=5):
C[i] = A[i] + B[i]
for i in range(tid, 1024, 6):
C[i] = A[i] + B[i]
return func
def op_of_literal():
op_list = [
(T.exp, 0),
(T.exp2, 0),
(T.exp10, 0),
(T.erf, 0.0),
(T.tanh, 0.0),
(T.sigmoid, 0.0),
(T.log, 0.0),
(T.log2, 0.0),
(T.log1p, 0.0),
(T.tan, 0.0),
(T.cos, 0.0),
(T.acos, 0.0),
(T.acosh, 0.0),
(T.sin, 0.0),
(T.sinh, 0.0),
(T.asin, 0.0),
(T.asinh, 0.0),
(T.atan, 0.0),
(T.atanh, 0.0),
(T.atan2, (1.0, 0.0)),
(T.sqrt, 0.0),
(T.rsqrt, 1.0),
(T.nextafter, (0.0, 1.0)),
(T.hypot, (1.0, 1.0)),
(T.copysign, (1.0, 1.0)),
(T.popcount, 0),
(T.fmod, (1.0, 1.0)),
]
def make_ir_generator(op, arg):
def inner():
call_expr = op(*arg) if isinstance(arg, tuple) else op(arg)
@T.prim_func(s_tir=True)
def func():
T.evaluate(call_expr)
return func
inner.__name__ = f"{op.__name__}_of_literal"
return inner
for op, arg in op_list:
yield make_ir_generator(op, arg)
def relax_extern_func():
@R.function
def func(A: R.Tensor([10, 20], "float32")):
func = R.ExternFunc("dummy_func")
B: R.Tensor([10, 20], "float32") = R.call_dps_packed(
func, [A], out_ty=R.Tensor([10, 20], "float32")
)
C: R.Tensor(ndim=2, dtype="float32") = R.call_dps_packed(
func, [B], out_ty=R.Tensor([10, 20], "float32")
)
return C
return func
def relax_match_cast_ty_proxy():
"""TypeProxy subclasses may be used as expressions
This is a regression test. The TVMScript parser allows Type
to be specified using a default-constructible class
(e.g. `R.Tensor` or `R.Shape`) rather than an instance of that
class (e.g. `R.Tensor()` or `R.Shape()`). In previous
implementations, this was only handled when the `Type` was
used in an annotation context. However, a `Type` may also
appear as an argument, which is passed to `R.match_cast`. Use of
a default-constructible class must be handled in this context as
well.
"""
def make_ir_generator(proxy_subclass):
def inner():
@R.function
def func(A: R.Any):
B = R.match_cast(A, proxy_subclass)
return B
return func
inner.__name__ = subclass.__name__
return inner
# Not all subclasses of TypeProxy are default-constructible.
# This list is a subset of `TypeProxy.__subclasses__()`,
# excluding `PrimProxy` and `DTensorProxy`.
subclasses = [
tvm.script.parser.relax.entry.AnyProxy,
tvm.script.parser.relax.entry.TensorProxy,
tvm.script.parser.relax.entry.CallableProxy,
tvm.script.parser.relax.entry.TupleProxy,
tvm.script.parser.relax.entry.ShapeProxy,
]
for subclass in subclasses:
yield make_ir_generator(subclass)
def relax_symbolic_var():
"""Relax tensors may use symbolic variables."""
N = tvm.tirx.Var("N", "int64")
@R.function
def func(A: R.Tensor([N], "float16")):
B: R.Tensor([N], "float16") = A
return B
return func
def relax_float_symbolic_var():
"""Relax scalar variables may use any dtype."""
@R.function
def func(value: R.Prim("float16")):
return value
return func
ir_generator = tvm.testing.parameter(
launch_env_thread,
opt_gemm_lower,
opt_conv_tensorcore_lower,
opt_conv_tensorcore_mod_host,
vthread_func,
matmul,
rank0,
rank0_block,
select,
minmax,
abs,
constant_folding,
simplify_bracket,
while_loop,
primfunc_with_allocate_annotations,
comm_reducer_single_reduce_group,
comm_reducer_multiple_reduce_groups,
multiple_commreducer,
loop_extent_dependent,
nontrivial_range_axis,
func_with_target_spec_by_config,
func_with_target_spec_by_str,
func_with_target_and_host_spec_by_str,
func_root_attr,
func_trivial_root_block,
func_nested_root_block,
func_T_ptr_let_statement,
func_T_ptr_allocate,
llvm_intrin_call,
parse_bufferslice_as_range_bound,
int64_support,
string_annotation_escaping,
pointer_type,
buffer_axis_separator,
buffer_ramp_access_as_slice_index,
ramp_int64,
scalable_vectors,
predicated_buffer_load_store,
let_expression,
void_ptr,
decl_buffer,
allocate_and_decl_buffer,
alloc_buffer_example,
float_infinity,
minimal_i32_literal,
boolean_argument,
bool_argument,
bool_variable_annotation,
bool_primitive,
bool_cast,
return_none,
implicit_evaluate,
if_true_else,
elif_chain_without_else,
elif_chain_with_else,
*nested_boolean_expressions(),
multi_env_threads,
intrinsic_pow,
bind_var,
string_stride,
string_stride_int64,
merge_shape_var_def,
if_then_else_var,
tvm_shfl_builtins,
make_packed_api_result,
tvm_struct_set_generated_in_cpp,
ir_module_with_attrs,
nested_seqstmt,
subroutine_call,
subroutine_call_returning_int,
undefined_data_ptr_in_decl_buffer,
undefined_shape_in_decl_buffer,
undefined_stride_in_decl_buffer,
undefined_elem_offset_in_decl_buffer,
subroutine_call_without_arguments,
return_zero,
return_zero_private,
return_zero_private_with_attr,
func_attr_with_list,
func_with_loop_jumps,
func_with_loop_steps,
*op_of_literal(),
*relax_match_cast_ty_proxy(),
relax_symbolic_var,
relax_float_symbolic_var,
)
relax_ir_generator = tvm.testing.parameter(
relax_extern_func,
)
show_all_relax_ty = tvm.testing.parameter(
by_dict={
"show_all_ty": True,
"hide_inferable_ty": False,
}
)
_NOT_ROUNDTRIP_STABLE: set[str] = set()
def test_roundtrip(ir_generator):
if getattr(ir_generator, "__name__", "") in _NOT_ROUNDTRIP_STABLE:
import pytest
pytest.skip(f"{ir_generator.__name__}: not round-trip stable here")
original = ir_generator()
after_roundtrip = tvm.script.from_source(
original.script(show_meta=True), check_well_formed=False
)
tvm.ir.assert_structural_equal(original, after_roundtrip, True)
def test_relax_roundtrip(relax_ir_generator, show_all_relax_ty):
original = relax_ir_generator()
after_roundtrip = tvm.script.from_source(
original.script(
show_meta=True,
show_all_ty=show_all_relax_ty,
)
)
tvm.ir.assert_structural_equal(original, after_roundtrip, True)
def test_return_none_no_trailing_type():
func = return_none()
script = func.script()
assert "-> None" not in script
def test_address_of_buffer():
@T.prim_func(s_tir=True)
def func(a: T.handle):
A = T.match_buffer(a, (128, 128), "float32")
T.evaluate(T.address_of(A))
assert "T.address_of(A[0, 0])" in func.script()
def test_assert_stmt_roundtrip_runtime_error():
"""RuntimeError assert roundtrips through print->parse."""
@T.prim_func(s_tir=True)
def func(x: T.int32):
assert x > 0, ("RuntimeError", ["x must be positive"])
script = func.script(show_meta=True)
roundtrip = tvm.script.from_source(script, check_well_formed=False)
tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
def test_assert_stmt_roundtrip_value_error():
"""ValueError assert roundtrips through print->parse."""
@T.prim_func(s_tir=True)
def func(x: T.int32):
assert x > 0, ("ValueError", ["Shape mismatch"])
script = func.script(show_meta=True)
roundtrip = tvm.script.from_source(script, check_well_formed=False)
tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
def test_assert_stmt_roundtrip_type_error():
"""TypeError assert roundtrips through print->parse."""
@T.prim_func(s_tir=True)
def func(x: T.int32):
assert x > 0, ("TypeError", ["Expected Tensor but got int"])
script = func.script(show_meta=True)
roundtrip = tvm.script.from_source(script, check_well_formed=False)
tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
def test_assert_stmt_roundtrip_multi_parts():
"""Multi-part message assert roundtrips with structural equality."""
@T.prim_func(s_tir=True)
def func(x: T.int32):
assert x > 0, ("TypeError", ["Expected ", "Tensor", " but got ", "int"])
script = func.script(show_meta=True)
roundtrip = tvm.script.from_source(script, check_well_formed=False)
tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
if __name__ == "__main__":
tvm.testing.main()
Reference in New Issue View Git Blame Copy Permalink