wehub-resource-sync
546 lines
19 KiB
Python
546 lines
19 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-function-docstring,missing-module-docstring,invalid-name,pointless-string-statement
|
|
# ruff: noqa: E741, F401, F841
|
|
import sys
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import tvm.testing
|
|
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
|
|
from tvm.script import from_source
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def transformed_matmul_no_syntax_sugar(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 i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
|
|
with T.sblock("update"):
|
|
vi, vj = T.axis.remap("SS", [i0, i1])
|
|
vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
|
|
T.reads([C[vi, vj], A[vi, vk], B[vj, vk]])
|
|
T.writes([C[vi, vj], A[vi, vk]])
|
|
with T.init():
|
|
C[vi, vj] = 0.0
|
|
A[vi, vk] = A[vi, vk] + B[vj, vk]
|
|
C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def transformed_matmul_syntax_sugar(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 i0, i1, i2_outer, i2_inner_outer, i2_inner_inner in T.grid(128, 128, 4, 8, 4):
|
|
with T.sblock("update"):
|
|
vi, vj = T.axis.remap("SS", [i0, i1])
|
|
vk = T.axis.R(128, i2_outer * 32 + i2_inner_outer * 4 + i2_inner_inner)
|
|
T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
|
|
T.writes(C[vi, vj], A[vi, vk])
|
|
with T.init():
|
|
C[vi, vj] = 0.0
|
|
A[vi, vk] = A[vi, vk] + B[vj, vk]
|
|
C[vi, vj] = C[vi, vj] + (A[vi, vk] * B[vj, vk])
|
|
|
|
|
|
def test_reads_writes_syntax_sugar():
|
|
assert_structural_equal_ignore_global_symbol(
|
|
transformed_matmul_no_syntax_sugar, transformed_matmul_syntax_sugar
|
|
)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def loop_no_syntax_sugar(a: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128, 128, 128))
|
|
for i in T.serial(0, 128):
|
|
for j in T.parallel(0, 128):
|
|
for k in T.vectorized(0, 128):
|
|
for x in T.unroll(0, 128):
|
|
for y in T.thread_binding(0, 128, thread="threadIdx.x"):
|
|
for z in T.thread_binding(0, 128, thread="threadIdx.x"):
|
|
A[i, j, k, x] = A[i, j, k, x] * 2.0
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def loop_syntax_sugar(a: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128, 128, 128))
|
|
for i in T.serial(128):
|
|
for j in T.parallel(128):
|
|
for k in T.vectorized(128):
|
|
for x in T.unroll(128):
|
|
for y in T.thread_binding(128, "threadIdx.x"):
|
|
for z in T.thread_binding(128, thread="threadIdx.x"):
|
|
A[i, j, k, x] = A[i, j, k, x] * 2.0
|
|
|
|
|
|
def test_loop_syntax_sugar():
|
|
assert_structural_equal_ignore_global_symbol(loop_no_syntax_sugar, loop_syntax_sugar)
|
|
|
|
|
|
# match buffer - use kwargs
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise_handle(
|
|
a: T.handle,
|
|
b: T.handle,
|
|
) -> None:
|
|
A = T.match_buffer(a, (128, 128, 128, 128))
|
|
B = T.match_buffer(b, (128, 128, 128, 128))
|
|
for i, j, k, l in T.grid(128, 128, 128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
|
B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
|
|
|
|
|
|
# match buffer - use buffer with kwargs
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise_buffer_kwargs(
|
|
a: T.Buffer(shape=(128, 128, 128, 128), dtype="float32"),
|
|
b: T.Buffer(shape=(128, 128, 128, 128), dtype="float32"),
|
|
) -> None:
|
|
for i, j, k, l in T.grid(128, 128, 128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
|
b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
|
|
|
|
|
|
# match buffer - use buffer without kwargs
|
|
@T.prim_func(s_tir=True)
|
|
def elementwise_buffer_no_kwargs(
|
|
a: T.Buffer((128, 128, 128, 128), "float32"),
|
|
b: T.Buffer((128, 128, 128, 128), "float32"),
|
|
) -> None:
|
|
for i, j, k, l in T.grid(128, 128, 128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj, vk, vl = T.axis.remap("SSSS", [i, j, k, l])
|
|
b[vi, vj, vk, vl] = a[vi, vj, vk, vl] * 2.0
|
|
|
|
|
|
def test_match_buffer_syntax_sugar():
|
|
# with kwargs
|
|
assert_structural_equal_ignore_global_symbol(elementwise_handle, elementwise_buffer_kwargs)
|
|
# without kwargs
|
|
assert_structural_equal_ignore_global_symbol(elementwise_handle, elementwise_buffer_no_kwargs)
|
|
|
|
|
|
def test_match_buffer_1d():
|
|
@T.prim_func(s_tir=True)
|
|
def func_no_sugar(a: T.handle):
|
|
A = T.match_buffer(a, shape=(16,))
|
|
for i in T.serial(16):
|
|
A[i] = 0.0
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func_with_sugar(A: T.Buffer(16, "float32")):
|
|
for i in T.serial(16):
|
|
A[i] = 0.0
|
|
|
|
assert_structural_equal_ignore_global_symbol(func_no_sugar, func_with_sugar)
|
|
|
|
|
|
# dynamic shape gemm
|
|
@T.prim_func(s_tir=True)
|
|
def gemm_dyn_shape(a: T.handle, b: T.handle, c: T.handle):
|
|
N = T.int32()
|
|
M = T.int32()
|
|
K = T.int32()
|
|
A = T.match_buffer(a, (N, K), "float32")
|
|
B = T.match_buffer(b, (K, M), "float32")
|
|
C = T.match_buffer(c, (N, M), "float32")
|
|
for i, j, k in T.grid(N, M, K):
|
|
with T.sblock("gemm"):
|
|
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
|
with T.init():
|
|
C[vi, vj] = 0.0
|
|
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
|
|
|
|
|
def test_dynamic_shape_gemm():
|
|
gemm_dyn_shape_roundtrip = from_source(gemm_dyn_shape.script())
|
|
assert_structural_equal_ignore_global_symbol(gemm_dyn_shape, gemm_dyn_shape_roundtrip)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def match_buffer_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
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def match_buffer_int64_after_roundtrip(
|
|
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
C: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
) -> None:
|
|
B = T.sblock_alloc_buffer((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
|
|
|
|
|
|
def test_match_buffer_int64():
|
|
original = match_buffer_int64
|
|
after_roundtrip = match_buffer_int64_after_roundtrip
|
|
assert_structural_equal_ignore_global_symbol(original, after_roundtrip, True)
|
|
|
|
|
|
def test_match_buffer_region_has_implicit_shape_dtype():
|
|
@T.prim_func(s_tir=True)
|
|
def explicit_shape_dtype(A: T.Buffer((16, 64), "int32")):
|
|
with T.sblock():
|
|
B = T.match_buffer(A[8:16, 32:64], shape=(8, 32), dtype="int32")
|
|
T.evaluate(0)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit_shape_dtype(A: T.Buffer((16, 64), "int32")):
|
|
with T.sblock():
|
|
B = T.match_buffer(A[8:16, 32:64])
|
|
T.evaluate(0)
|
|
|
|
assert_structural_equal_ignore_global_symbol(explicit_shape_dtype, implicit_shape_dtype)
|
|
|
|
|
|
def test_match_buffer_input_requires_shape_arg():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func(a: T.handle):
|
|
A = T.match_buffer(a, dtype="int32")
|
|
T.evaluate(0)
|
|
|
|
|
|
def test_bind_bufferload_without_type_annotation():
|
|
# Variable assignment of Expr types uses the dtype of the
|
|
# Expr to determine the variable's dtype. Parsing of
|
|
# buf[indices] is done by generating a BufferSlice object, which
|
|
# handles both store and load cases. BufferSlice is not a
|
|
# Expr, and implements BufferSlice.dtype explicitly.
|
|
|
|
# Failure occurred during parsing of the tvmscript.
|
|
@T.prim_func(s_tir=True)
|
|
def func_without_type_annotation(A: T.Buffer((1,), "int32")):
|
|
x = A[0]
|
|
T.evaluate(x)
|
|
|
|
|
|
def test_bind_with_constant():
|
|
@T.prim_func(s_tir=True)
|
|
def constant_binds():
|
|
x = T.meta_var(1)
|
|
y = T.meta_var(42.0)
|
|
T.evaluate(T.cast(x, "float32") + y)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def constant_binds_wrapped():
|
|
x = T.meta_var(T.int32(1))
|
|
y = T.meta_var(T.float32(42.0))
|
|
T.evaluate(T.cast(x, "float32") + y)
|
|
|
|
assert_structural_equal_ignore_global_symbol(constant_binds, constant_binds_wrapped)
|
|
|
|
|
|
def test_func_call():
|
|
def shared_16x16_to_ldmatrix_32x8_layout(i, j):
|
|
thread_id = (i % 8) * 4 + (j % 8) // 2
|
|
return T.meta_var((thread_id, (j // 8) * 4 + (i // 8) * 2 + (j % 2)))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sync_m16n16k16_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
B = T.match_buffer(b, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
C = T.match_buffer(c, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
|
|
with T.sblock("root"):
|
|
T.reads(C[0:32, 0:8], A[0:32, 0:8], B[0:32, 0:8])
|
|
T.writes(C[0:32, 0:8])
|
|
for i, j, k in T.grid(16, 16, 16):
|
|
with T.sblock("C"):
|
|
i, j, k = T.axis.remap("SSR", [i, j, k])
|
|
thread_id_C, local_id_C = shared_16x16_to_ldmatrix_32x8_layout(i, j)
|
|
thread_id_A, local_id_A = shared_16x16_to_ldmatrix_32x8_layout(i, k)
|
|
thread_id_B, local_id_B = shared_16x16_to_ldmatrix_32x8_layout(k, j)
|
|
|
|
T.reads(
|
|
C[thread_id_C, local_id_C],
|
|
A[thread_id_A, local_id_A],
|
|
B[thread_id_B, local_id_B],
|
|
)
|
|
T.writes(C[thread_id_C, local_id_C])
|
|
|
|
C[thread_id_C, local_id_C] += (
|
|
A[thread_id_A, local_id_A] * B[thread_id_B, local_id_B]
|
|
)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sync_m16n16k16_desc_manual(a: T.handle, b: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
B = T.match_buffer(b, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
C = T.match_buffer(c, (32, 8), "float16", align=64, offset_factor=16, scope="warp")
|
|
|
|
with T.sblock("root"):
|
|
T.reads(C[0:32, 0:8], A[0:32, 0:8], B[0:32, 0:8])
|
|
T.writes(C[0:32, 0:8])
|
|
for i, j, k in T.grid(16, 16, 16):
|
|
with T.sblock("C"):
|
|
i, j, k = T.axis.remap("SSR", [i, j, k])
|
|
T.reads(
|
|
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2],
|
|
A[i % 8 * 4 + k % 8 // 2, k // 8 * 4 + i // 8 * 2 + k % 2],
|
|
B[k % 8 * 4 + j % 8 // 2, j // 8 * 4 + k // 8 * 2 + j % 2],
|
|
)
|
|
T.writes(C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2])
|
|
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2] = (
|
|
C[i % 8 * 4 + j % 8 // 2, j // 8 * 4 + i // 8 * 2 + j % 2]
|
|
+ A[i % 8 * 4 + k % 8 // 2, k // 8 * 4 + i // 8 * 2 + k % 2]
|
|
* B[k % 8 * 4 + j % 8 // 2, j // 8 * 4 + k // 8 * 2 + j % 2]
|
|
)
|
|
|
|
assert_structural_equal_ignore_global_symbol(
|
|
mma_sync_m16n16k16_desc, mma_sync_m16n16k16_desc_manual
|
|
)
|
|
|
|
# The following is an example of an error message from calling an invalid function
|
|
|
|
# error: Error occurred when invoking the function sqrt:
|
|
# loop of ufunc does not support argument 0 of type Var which has no callable sqrt method
|
|
# --> test_tvmscript_syntax_sugar.py:334:19
|
|
# |
|
|
# 334 | ind = sqrt(i)
|
|
# | ^^^^^^^
|
|
# note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.
|
|
|
|
# Uncomment to see the error above.
|
|
# def sqrt(x):
|
|
# import numpy as np
|
|
# return np.sqrt(x)
|
|
|
|
# @T.prim_func
|
|
# def loop(a: T.handle) -> None:
|
|
# A = T.match_buffer(a, (128,))
|
|
# for i in T.serial(128):
|
|
# ind = sqrt(i)
|
|
# A[i] = A[ind]
|
|
|
|
|
|
def test_int64_loop():
|
|
@T.prim_func(s_tir=True)
|
|
def int64_grid(
|
|
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
) -> None:
|
|
for i, j in T.grid(T.int64(128), T.int64(128)):
|
|
with T.sblock("C"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + 1.0
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def int64_grid_expanded(
|
|
A: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
B: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
) -> None:
|
|
for i in range(T.int64(0), T.int64(128)):
|
|
for j in range(T.int64(0), T.int64(128)):
|
|
with T.sblock("C"):
|
|
vi = T.axis.spatial(T.int64(128), i)
|
|
vj = T.axis.spatial(T.int64(128), j)
|
|
B[vi, vj] = A[vi, vj] + 1.0
|
|
|
|
assert_structural_equal_ignore_global_symbol(int64_grid, int64_grid_expanded)
|
|
|
|
|
|
def test_implicit_evaluate_assume():
|
|
@T.prim_func(s_tir=True)
|
|
def explicit(A: T.Buffer(1, "int32")):
|
|
T.evaluate(T.assume(A[0] == 5))
|
|
A[0] = 10
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit(A: T.Buffer(1, "int32")):
|
|
T.assume(A[0] == 5)
|
|
A[0] = 10
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_implicit_evaluate_call_extern():
|
|
@T.prim_func(s_tir=True)
|
|
def explicit(A: T.Buffer(1, "int32")):
|
|
T.evaluate(T.call_extern("extern_func", A.data, dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit(A: T.Buffer(1, "int32")):
|
|
T.call_extern("extern_func", A.data, dtype="int32")
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_preserve_trivial_let_binding():
|
|
"""Trivial `T.let[...]` annotations survive the parser as LetStmt and are not inlined.
|
|
|
|
In fork, bare `j = i` lowers to a local_scalar (AllocBuffer + BufferStore); the
|
|
LetStmt form is opt-in via `T.let[T.dtype]`. Both the explicit `T.bind(..., var=j)`
|
|
builder API and the `j: T.let[T.dtype]` annotation produce the same LetStmt IR.
|
|
"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit(i: T.int32):
|
|
j = T.int32()
|
|
T.bind(i, var=j)
|
|
T.evaluate(j)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit(i: T.int32):
|
|
j: T.let[T.int32] = i
|
|
T.evaluate(j)
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_preserve_trivial_let_binding_of_value():
|
|
"""Same as test_preserve_trivial_let_binding but with a constant RHS."""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit(i: T.int32):
|
|
j = T.int32()
|
|
T.bind(42, var=j)
|
|
T.evaluate(j)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit(i: T.int32):
|
|
j: T.let[T.int32] = 42
|
|
T.evaluate(j)
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_preserve_parameter_name():
|
|
@T.prim_func(s_tir=True)
|
|
def func(i: T.int32):
|
|
j = i
|
|
T.evaluate(j)
|
|
|
|
param_name = func.params[0].name
|
|
assert param_name == "i"
|
|
|
|
|
|
def test_preserve_variable_name():
|
|
"""Use variable name when generating tirx::Bind / AllocBuffer"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def func():
|
|
for i in T.serial(16):
|
|
j = i // 4
|
|
T.evaluate(j)
|
|
|
|
# In fork, bare `j = i // 4` lowers to AllocBuffer (local_scalar) in the for-body
|
|
# SeqStmt; the variable name lives on the underlying buffer.
|
|
var_name = func.body.body.seq[0].buffer.name
|
|
assert var_name == "j"
|
|
|
|
|
|
def test_boolean_constant():
|
|
"""Python booleans should become T.Bool objects"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit():
|
|
T.evaluate(T.bool(True))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit():
|
|
T.evaluate(True)
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_foldable_boolean_in_assert():
|
|
"""Foldable booleans T.Bool objects
|
|
|
|
The condition of an assert statement should be a boolean
|
|
expression. Previously, this test failed because the FFI does not
|
|
distinguish between integer primitives and boolean primitives.
|
|
"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit():
|
|
assert T.bool(False), "Message"
|
|
T.evaluate(0)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit():
|
|
assert 0 == 1, "Message"
|
|
T.evaluate(0)
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_return_statement():
|
|
"""A python `return` statement uses `T.ret`"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit():
|
|
T.evaluate(T.ret(5))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit():
|
|
return 5
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
def test_loop_jump_statement():
|
|
"""`break` and `continue` evaluates to TIR intrinsics"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def explicit():
|
|
for i in range(16):
|
|
if i % 2 == 0:
|
|
T.evaluate(T.continue_loop())
|
|
if i < 15:
|
|
T.evaluate(T.break_loop())
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def implicit():
|
|
for i in range(16):
|
|
if i % 2 == 0:
|
|
continue
|
|
if i < 15:
|
|
break
|
|
|
|
assert_structural_equal_ignore_global_symbol(implicit, explicit)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|