| # 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: F821 |
| """Tests for ``@T.jit`` compile-time specialization.""" |
| |
| from __future__ import annotations |
| |
| import typing |
| |
| import pytest |
| |
| import tvm |
| from tvm.ir import assert_structural_equal |
| from tvm.script import tirx as T |
| from tvm.script.tirx import tile as Tx |
| |
| |
| def test_int_constexpr_specializes_loop_bound(): |
| @T.jit(private=True) |
| def add( |
| A: T.Buffer((N,), "int32"), |
| B: T.Buffer((N,), "int32"), |
| C: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| for i in range(N): |
| C[i] = A[i] + B[i] |
| |
| @T.prim_func(private=True) |
| def expected( |
| A: T.Buffer((128,), "int32"), |
| B: T.Buffer((128,), "int32"), |
| C: T.Buffer((128,), "int32"), |
| ): |
| for i in range(128): |
| C[i] = A[i] + B[i] |
| |
| assert_structural_equal(add.specialize(N=128), expected, map_free_vars=True) |
| |
| |
| def test_constexpr_in_2d_buffer_shape(): |
| @T.jit(private=True) |
| def matadd( |
| A: T.Buffer((M, K), "int32"), |
| B: T.Buffer((M, K), "int32"), |
| C: T.Buffer((M, K), "int32"), |
| *, |
| M: T.constexpr, |
| K: T.constexpr, |
| ): |
| for m in range(M): |
| for k in range(K): |
| C[m, k] = A[m, k] + B[m, k] |
| |
| @T.prim_func(private=True) |
| def expected( |
| A: T.Buffer((4, 8), "int32"), |
| B: T.Buffer((4, 8), "int32"), |
| C: T.Buffer((4, 8), "int32"), |
| ): |
| for m in range(4): |
| for k in range(8): |
| C[m, k] = A[m, k] + B[m, k] |
| |
| assert_structural_equal(matadd.specialize(M=4, K=8), expected, map_free_vars=True) |
| |
| |
| def test_constexpr_in_body_expression(): |
| @T.jit(private=True) |
| def scaled_copy( |
| A: T.Buffer((N,), "int32"), |
| B: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| SCALE: T.constexpr, |
| ): |
| for i in range(N): |
| B[i] = A[i] * SCALE |
| |
| @T.prim_func(private=True) |
| def expected( |
| A: T.Buffer((16,), "int32"), |
| B: T.Buffer((16,), "int32"), |
| ): |
| for i in range(16): |
| B[i] = A[i] * 3 |
| |
| assert_structural_equal(scaled_copy.specialize(N=16, SCALE=3), expected, map_free_vars=True) |
| |
| |
| def test_specialize_cache_returns_same_instance(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| for i in range(N): |
| A[i] = 0 |
| |
| a = k.specialize(N=8) |
| b = k.specialize(N=8) |
| assert a is b |
| |
| |
| def test_specialize_different_args_produce_different_funcs(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| for i in range(N): |
| A[i] = 0 |
| |
| assert k.specialize(N=8) is not k.specialize(N=16) |
| |
| |
| def test_specialize_missing_constexpr_raises(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| SCALE: T.constexpr, |
| ): |
| for i in range(N): |
| A[i] = SCALE |
| |
| with pytest.raises(TypeError, match="missing"): |
| k.specialize(N=8) |
| |
| |
| def test_specialize_extra_kwarg_raises(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| for i in range(N): |
| A[i] = 0 |
| |
| with pytest.raises(TypeError, match="unexpected"): |
| k.specialize(N=8, BOGUS=42) |
| |
| |
| def test_jit_kernel_with_nested_inline_helper(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| @T.inline |
| def double(x): |
| A[x] = A[x] * 2 |
| |
| for i in range(N): |
| double(i) |
| |
| @T.prim_func(private=True) |
| def expected( |
| A: T.Buffer((4,), "int32"), |
| ): |
| for i in range(4): |
| A[i] = A[i] * 2 |
| |
| assert_structural_equal(k.specialize(N=4), expected, map_free_vars=True) |
| |
| |
| def test_constexpr_default_value(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| SCALE: T.constexpr = 7, |
| ): |
| for i in range(N): |
| A[i] = SCALE |
| |
| @T.prim_func(private=True) |
| def expected( |
| A: T.Buffer((8,), "int32"), |
| ): |
| for i in range(8): |
| A[i] = 7 |
| |
| assert_structural_equal(k.specialize(N=8), expected, map_free_vars=True) |
| # Override the default |
| overridden = k.specialize(N=8, SCALE=99) |
| assert k.specialize(N=8) is not overridden |
| |
| |
| def test_specialize_returns_primfunc(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((N,), "int32"), |
| *, |
| N: T.constexpr, |
| ): |
| for i in range(N): |
| A[i] = 0 |
| |
| spec = k.specialize(N=8) |
| assert isinstance(spec, tvm.tirx.PrimFunc) |
| # Specialized PrimFunc has only the runtime params (constexpr stripped). |
| assert len(spec.params) == 1 |
| |
| |
| def test_constexpr_specializes_nested_selector_condition(): |
| @T.jit(private=True) |
| def k( |
| A: T.Buffer((8,), "float16"), |
| B: T.Buffer((8,), "float16"), |
| C: T.Buffer((8,), "float16"), |
| flag: T.int32, |
| *, |
| LIMIT: T.constexpr, |
| ): |
| Tx.copy_async( |
| C[:], |
| A[:], |
| dispatch="tma_explicit", |
| mbar=C.data, |
| src_selector=[(flag < LIMIT, B)], |
| ) |
| |
| specialized = k.specialize(LIMIT=4) |
| op_call = specialized.body |
| condition, candidate = op_call.config["src_selector"][0] |
| assert isinstance(condition, tvm.tirx.LT) |
| assert int(condition.b) == 4 |
| assert candidate.same_as(specialized.params[1]) |
| |
| |
| def test_optional_param_present_and_absent_ir(): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), out_h: T.handle): |
| out = T.match_buffer(out_h, (1,), "int32") |
| if a is not None: |
| A = T.match_buffer(a, (1,), "int32") |
| out[0] = A[0] |
| else: |
| out[0] = -1 |
| |
| @T.prim_func(private=True) |
| def expected_present(a: T.handle, out_h: T.handle): |
| A = T.match_buffer(a, (1,), "int32") |
| out = T.match_buffer(out_h, (1,), "int32") |
| out[0] = A[0] |
| |
| @T.prim_func(private=True) |
| def expected_absent(out_h: T.handle): |
| out = T.match_buffer(out_h, (1,), "int32") |
| out[0] = -1 |
| |
| present = kernel.specialize() |
| absent = kernel.specialize(a=None) |
| assert_structural_equal(present, expected_present, map_free_vars=True) |
| assert_structural_equal(absent, expected_absent, map_free_vars=True) |
| assert [param.name for param in present.params] == ["A", "out"] |
| assert [param.name for param in absent.params] == ["out"] |
| assert all(tvm.tirx.is_buffer_var(param) for param in present.params) |
| assert all(tvm.tirx.is_buffer_var(param) for param in absent.params) |
| |
| |
| def test_optional_specialization_cache_includes_presence(): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), out_h: T.handle): |
| out = T.match_buffer(out_h, (1,), "int32") |
| if a is not None: |
| A = T.match_buffer(a, (1,), "int32") |
| out[0] = A[0] |
| else: |
| out[0] = 0 |
| |
| present = kernel.specialize() |
| absent = kernel.specialize(a=None) |
| assert present is kernel.specialize() |
| assert absent is kernel.specialize(a=None) |
| assert present is not absent |
| |
| |
| def test_multiple_optional_params_preserve_runtime_order(): |
| @T.jit(private=True) |
| def kernel( |
| first_h: T.handle, |
| a: T.Optional(T.handle), |
| scale: T.int32, |
| b: T.Optional(T.handle), |
| out_h: T.handle, |
| ): |
| first = T.match_buffer(first_h, (1,), "int32") |
| out = T.match_buffer(out_h, (1,), "int32") |
| out[0] = first[0] * scale |
| if a is not None: |
| A = T.match_buffer(a, (1,), "int32") |
| out[0] = out[0] + A[0] |
| if b is not None: |
| B = T.match_buffer(b, (1,), "int32") |
| out[0] = out[0] + B[0] |
| |
| assert [param.name for param in kernel.specialize().params] == [ |
| "first", |
| "A", |
| "scale", |
| "B", |
| "out", |
| ] |
| assert [param.name for param in kernel.specialize(a=None).params] == [ |
| "first", |
| "scale", |
| "B", |
| "out", |
| ] |
| assert [param.name for param in kernel.specialize(a=None, b=None).params] == [ |
| "first", |
| "scale", |
| "out", |
| ] |
| |
| |
| def test_optional_only_accepts_none_at_specialization_time(): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), out_h: T.handle): |
| if a is not None: |
| T.match_buffer(a, (1,), "int32") |
| T.match_buffer(out_h, (1,), "int32") |
| |
| with pytest.raises(TypeError, match="only accept None"): |
| kernel.specialize(a=object()) |
| |
| |
| def test_only_explicit_t_optional_is_specializable(): |
| @T.jit(private=True) |
| def typing_optional(a: typing.Optional[T.handle], out_h: T.handle): # noqa: UP045 |
| T.evaluate(0) |
| |
| @T.jit(private=True) |
| def union_optional(a: T.handle | None, out_h: T.handle): |
| T.evaluate(0) |
| |
| with pytest.raises(TypeError, match="unexpected"): |
| typing_optional.specialize(a=None) |
| with pytest.raises(TypeError, match="unexpected"): |
| union_optional.specialize(a=None) |
| |
| |
| def test_t_optional_is_restricted_to_jit(): |
| with pytest.raises(tvm.error.DiagnosticError, match="only supported by @T.jit"): |
| |
| @T.prim_func(private=True) |
| def invalid(a: T.Optional(T.handle)): |
| T.evaluate(0) |
| |
| |
| def test_compile_time_if_binding_uses_python_scope(): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), out_h: T.handle): |
| if a is None: |
| selected = T.match_buffer(out_h, (1,), "int32") |
| else: |
| selected = T.match_buffer(a, (1,), "int32") |
| selected[0] = 1 |
| |
| present = kernel.specialize() |
| absent = kernel.specialize(a=None) |
| assert len(present.params) == 2 |
| assert len(absent.params) == 1 |
| assert sum(tvm.tirx.is_buffer_var(param) for param in present.params) == 1 |
| assert sum(tvm.tirx.is_buffer_var(param) for param in absent.params) == 1 |
| |
| |
| def test_compile_time_bool_ops_and_if_expression_short_circuit(): |
| def fail_if_evaluated(): |
| raise RuntimeError("dead expression was evaluated") |
| |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), out_h: T.handle): |
| out = T.match_buffer(out_h, (1,), "int32") |
| if a is None or fail_if_evaluated(): |
| out[0] = 1 |
| if a is not None and fail_if_evaluated(): |
| out[0] = 2 |
| out[0] = 3 if a is None else fail_if_evaluated() |
| |
| absent = kernel.specialize(a=None) |
| assert [param.name for param in absent.params] == ["out"] |
| |
| |
| def test_runtime_tir_if_cannot_guard_absent_optional_param(): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle), flag: T.int32): |
| if flag != 0: |
| T.match_buffer(a, (1,), "int32") |
| |
| with pytest.raises(tvm.error.DiagnosticError, match="match_buffer"): |
| kernel.specialize(a=None) |
| |
| |
| @pytest.mark.parametrize( |
| ("operation", "source_text"), |
| [ |
| ("subscript", "a[10]"), |
| ("attribute", "a.ptr_to"), |
| ("match_buffer", "T.match_buffer"), |
| ], |
| ) |
| def test_unguarded_absent_optional_param_reports_source(operation, source_text): |
| @T.jit(private=True) |
| def kernel(a: T.Optional(T.handle)): |
| if operation == "subscript": |
| a[10] |
| elif operation == "attribute": |
| a.ptr_to([0]) |
| else: |
| T.match_buffer(a, (1,), "int32") |
| |
| with pytest.raises(tvm.error.DiagnosticError) as exc_info: |
| kernel.specialize(a=None) |
| assert source_text in str(exc_info.value) |
| |
| |
| def test_present_optional_param_still_rejects_ffi_none(): |
| @T.jit |
| def kernel(a: T.Optional(T.handle)): |
| A = T.match_buffer(a, (1,), "int32") |
| A[0] = 0 |
| |
| executable = tvm.compile(kernel.specialize(), target="llvm", tir_pipeline="tirx") |
| with pytest.raises(TypeError, match="expected Tensor"): |
| executable(None) |
| |
| |
| if __name__ == "__main__": |
| pytest.main([__file__, "-v"]) |
