# 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 # ruff: noqa: E501, F841 import numpy as np import pytest import tvm import tvm.testing from tvm import s_tir, te, tirx, topi from tvm.script import tirx as T def test_unique_name_complete_block(): A = te.placeholder((16, 16), name="A") B = te.compute((16, 16), lambda x, y: A[x, y] * 2, name="main") C = te.compute((16, 16), lambda x, y: B[x, y] + 1, name="main") func = te.create_prim_func([A, C]) s = tvm.s_tir.Schedule(func, debug_mask="all") assert isinstance(s.get_sref(s.get_sblock("main")), s_tir.schedule.StmtSRef) assert isinstance(s.get_sref(s.get_sblock("main_1")), s_tir.schedule.StmtSRef) def test_unique_name_reduction_block(): k1 = te.reduce_axis((0, 16), "k1") k2 = te.reduce_axis((0, 16), "k2") A = te.placeholder((16, 16), name="A") B = te.compute((16,), lambda i: te.sum(A[i, k1], axis=k1), name="sum") C = te.compute((), lambda: te.sum(B[k2], axis=k2), name="sum") func = te.create_prim_func([A, C]) s = tvm.s_tir.Schedule(func, debug_mask="all") assert isinstance(s.get_sref(s.get_sblock("sum")), s_tir.schedule.StmtSRef) assert isinstance(s.get_sref(s.get_sblock("sum_1")), s_tir.schedule.StmtSRef) def _check_workload(te_workload, tir_workload, index_dtype_override=None, do_simplify=False): func = te.create_prim_func(te_workload(), index_dtype_override) if do_simplify: simplify = tirx.transform.StmtSimplify() func = simplify(tvm.IRModule.from_expr(func))["main"] tir_workload = simplify(tvm.IRModule.from_expr(tir_workload))["main"] tvm.ir.assert_structural_equal(func, tir_workload) # make sure that we can create schedule from the func s = tvm.s_tir.Schedule(func, debug_mask="all") assert s def te_matmul(): k = te.reduce_axis((0, 128), "k") A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C") return [A, B, C] @T.prim_func(s_tir=True) def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(a, (128, 128)) B = T.match_buffer(b, (128, 128)) C = T.match_buffer(c, (128, 128)) for i0, j0, k0 in T.grid(128, 128, 128): with T.sblock(): i, j, k = T.axis.remap("SSR", [i0, j0, k0]) with T.init(): C[i, j] = 0.0 C[i, j] += A[i, k] * B[j, k] @T.prim_func(s_tir=True) def tir_matmul_int64( A: T.Buffer((T.int64(128), T.int64(128)), "float32"), B: T.Buffer((T.int64(128), T.int64(128)), "float32"), C: T.Buffer((T.int64(128), T.int64(128)), "float32"), ) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) for i0, j0, k0 in T.grid(T.int64(128), T.int64(128), T.int64(128)): with T.sblock(): i, j, k = T.axis.remap("SSR", [i0, j0, k0]) with T.init(): C[i, j] = 0.0 C[i, j] += A[i, k] * B[j, k] def test_matmul(): _check_workload(te_matmul, tir_matmul) def test_matmul_int64(): _check_workload(te_matmul, tir_matmul_int64, index_dtype_override="int64") def te_element_wise(): A = te.placeholder((128, 128), name="A") B = te.compute((128, 128), lambda x, y: A[x, y] * 2, name="B") C = te.compute((128, 128), lambda x, y: B[x, y] + 1, name="C") return [A, C] @T.prim_func(s_tir=True) def tir_element_wise(a: T.handle, c: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(a, (128, 128)) C = T.match_buffer(c, (128, 128)) B = T.sblock_alloc_buffer((128, 128)) for i0, j0 in T.grid(128, 128): with T.sblock(): i, j = T.axis.remap("SS", [i0, j0]) B[i, j] = A[i, j] * 2.0 for i0, j0 in T.grid(128, 128): with T.sblock(): i, j = T.axis.remap("SS", [i0, j0]) C[i, j] = B[i, j] + 1.0 def test_element_wise(): _check_workload(te_element_wise, tir_element_wise) def te_conv2d(): batch = 16 in_channel = 16 out_channel = 32 size = 14 kernel = 3 A = te.placeholder((batch, in_channel, size, size), name="A") W = te.placeholder((in_channel, kernel, kernel, out_channel), name="W") Apad = te.compute( (batch, in_channel, size + 2, size + 2), lambda nn, cc, yy, xx: tvm.tirx.if_then_else( tvm.tirx.all(yy >= 1, yy - 1 < size, xx >= 1, xx - 1 < size), A[nn, cc, yy - 1, xx - 1], 0.0, ), name="Apad", ) rc = te.reduce_axis((0, in_channel), name="rc") ry = te.reduce_axis((0, kernel), name="ry") rx = te.reduce_axis((0, kernel), name="rx") B = te.compute( (batch, out_channel, size, size), lambda nn, ff, yy, xx: te.sum( Apad[nn, rc, yy + ry, xx + rx] * W[rc, ry, rx, ff], axis=[rc, ry, rx] ), name="B", ) return [A, W, B] @T.prim_func(s_tir=True) def tir_conv2d(a: T.handle, w: T.handle, b: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(a, [16, 16, 14, 14]) W = T.match_buffer(w, [16, 3, 3, 32]) B = T.match_buffer(b, [16, 32, 14, 14]) Apad = T.sblock_alloc_buffer([16, 16, 16, 16]) for n, c, y, x in T.grid(16, 16, 16, 16): with T.sblock("Apad"): nn, cc, yy, xx = T.axis.remap("SSSS", [n, c, y, x]) Apad[nn, cc, yy, xx] = T.if_then_else( 1 <= yy and yy < 15 and 1 <= xx and xx < 15, A[nn, cc, yy - 1, xx - 1], 0.0, dtype="float32", ) for n, f, y, x, kc, ky, kx in T.grid(16, 32, 14, 14, 16, 3, 3): with T.sblock("B"): nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [n, f, y, x, kc, ky, kx]) with T.init(): B[nn, ff, yy, xx] = 0.0 B[nn, ff, yy, xx] += Apad[nn, rc, yy + ry, xx + rx] * W[rc, ry, rx, ff] def test_conv2d(): _check_workload(te_conv2d, tir_conv2d) def te_multi_output(): n = te.var("n") m = te.var("m") A0 = te.placeholder((m, n), name="A0") A1 = te.placeholder((m, n), name="A1") B0, B1 = te.compute((m, n), lambda i, j: (A0[i, j] + 2, A1[i, j] * 3), name="B") return [A0, A1, B0, B1] @T.prim_func(s_tir=True) def tir_multi_output(a0: T.handle, a1: T.handle, b0: T.handle, b1: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) m = T.int32() n = T.int32() A0 = T.match_buffer(a0, (m, n)) A1 = T.match_buffer(a1, (m, n)) B0 = T.match_buffer(b0, (m, n)) B1 = T.match_buffer(b1, (m, n)) for i0, i1 in T.grid(m, n): with T.sblock("B.v0"): i, j = T.axis.remap("SS", [i0, i1]) B0[i, j] = A0[i, j] + 2.0 with T.sblock("B.v1"): i, j = T.axis.remap("SS", [i0, i1]) B1[i, j] = A1[i, j] * 3.0 def test_multi_output(): _check_workload(te_multi_output, tir_multi_output) def te_extern(): A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") C = te.extern( (128, 128), [A, B], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], 0, 0 ), name="C", ) return [A, B, C] @T.prim_func(s_tir=True) def tir_extern(a: T.handle, b: T.handle, c: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) off1 = te.var("elem_offset") off2 = te.var("elem_offset_1") off3 = te.var("elem_offset_2") A = T.match_buffer(a, (128, 128), elem_offset=off1) B = T.match_buffer(b, (128, 128), elem_offset=off2) C = T.match_buffer(c, (128, 128), elem_offset=off3) # body with T.sblock("C"): T.reads() T.writes() T.evaluate( T.tvm_call_packed( "tvm.contrib.cblas.matmul", T.tvm_stack_make_array( A.data, T.tvm_stack_make_shape(128, 128, dtype="handle"), 0, 2, 0.0, off1, dtype="handle", ), T.tvm_stack_make_array( B.data, T.tvm_stack_make_shape(128, 128, dtype="handle"), 0, 2, 0.0, off2, dtype="handle", ), T.tvm_stack_make_array( C.data, T.tvm_stack_make_shape(128, 128, dtype="handle"), 0, 2, 0.0, off3, dtype="handle", ), 0, 0, dtype="int32", ) ) def test_extern(): _check_workload(te_extern, tir_extern) def te_extern_epilogue(): A = te.placeholder((4, 3), name="A") B = te.placeholder((3, 2), name="B") C = te.extern( (4, 2), [A, B], lambda ins, outs: tvm.tirx.call_packed("testing.echo", ins[0], ins[1], outs[0]), name="C", ) D = te.compute(C.shape, lambda i, j: C[i, j] + 1.0, name="D") return [A, B, D] @T.prim_func(s_tir=True) def tir_extern_epilogue(var_A: T.handle, var_B: T.handle, D: T.Buffer((4, 2), "float32")): T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(var_A, (4, 3), offset_factor=1) B = T.match_buffer(var_B, (3, 2), offset_factor=1) C = T.sblock_alloc_buffer((4, 2), elem_offset=0, offset_factor=1) with T.sblock("C"): T.reads() T.writes() T.call_packed("testing.echo", A, B, C) for i, j in T.grid(4, 2): with T.sblock("D"): vi, vj = T.axis.remap("SS", [i, j]) T.reads(C[vi, vj]) T.writes(D[vi, vj]) D[vi, vj] = C[vi, vj] + T.float32(1) def test_extern_epilogue(): _check_workload(te_extern_epilogue, tir_extern_epilogue) func = te.create_prim_func(te_extern_epilogue()).with_attr("global_symbol", "extern_epilogue") tvm.compile(func, target="llvm") def te_reordered_matmul(): k = te.reduce_axis((0, 128), "k") A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C") return [C, A, B] @T.prim_func(s_tir=True) def tir_reordered_matmul(c: T.handle, a: T.handle, b: T.handle) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(a, (128, 128)) B = T.match_buffer(b, (128, 128)) C = T.match_buffer(c, (128, 128)) for i0, j0, k0 in T.grid(128, 128, 128): with T.sblock(): i, j, k = T.axis.remap("SSR", [i0, j0, k0]) with T.init(): C[i, j] = 0.0 C[i, j] += A[i, k] * B[j, k] def test_arg_order(): _check_workload(te_reordered_matmul, tir_reordered_matmul) def te_scan(): m = te.var("m") n = te.var("n") X = te.placeholder((m, n), name="X") s_state = te.placeholder((m, n)) s_init = te.compute((1, n), lambda _, i: X[0, i]) s_update = te.compute((m, n), lambda t, i: s_state[t - 1, i] + X[t, i]) s_scan = tvm.te.scan(s_init, s_update, s_state, inputs=[X]) return [X, s_scan] def test_error_reporting(): try: te.create_prim_func(te_scan()) assert False except (TypeError, tvm.error.InternalError) as e: error_message = str(e) assert error_message.find("Unsupported Operation: te.ScanOp.") != -1 return assert False def test_constant(): M = 11 A = te.placeholder((M,), name="A") B = te.compute(tuple(), lambda: 2, name="B") # Manually craft ProducerLoad because `B[]` is not allowed. C = te.compute( (M,), lambda x: A[x] + tvm.tirx.expr.ProducerLoad(B, []), name="C", tag="broadcast" ) func = te.create_prim_func([C, A]) func = tvm.compile(func) a_np = np.random.uniform(size=(M,)).astype(A.dtype.dtype) c = tvm.runtime.tensor(np.zeros(M, dtype=C.dtype.dtype)) x = func(c, tvm.runtime.tensor(a_np)) tvm.testing.assert_allclose(a_np + 2, c.numpy()) @pytest.mark.parametrize("op_name", ["acos", "acosh", "asin", "asinh", "atanh"]) def test_topi_float_unary_rejects_integer_input(op_name): x = te.placeholder((1, 8), dtype="int16", name="x") op = getattr(topi, op_name) with pytest.raises( TypeError, match=rf"topi\.{op_name} only supports floating-point inputs, but got int16", ): op(x) @pytest.mark.parametrize("op_name", ["acos", "acosh", "asin", "asinh", "atanh"]) @pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) def test_topi_float_unary_accepts_float_input(op_name, dtype): x = te.placeholder((1, 8), dtype=dtype, name="x") op = getattr(topi, op_name) out = op(x) func = te.create_prim_func([x, out]).with_attr("target", tvm.target.Target("llvm")) mod = tvm.IRModule({"main": func}) compiled = tvm.build(mod, target="llvm") assert compiled is not None def test_data_dependent_access(): A = te.placeholder((10,), name="A") B = te.placeholder((10,), name="B", dtype="int32") C = te.compute((10,), lambda i: A[B[i]]) func = te.create_prim_func([C, A, B]) func = tvm.compile(func) a_np = np.random.uniform(size=(10,)).astype(A.dtype.dtype) b_np = np.arange(10, dtype=B.dtype.dtype) c = tvm.runtime.tensor(np.zeros(10, dtype=C.dtype.dtype)) func(c, tvm.runtime.tensor(a_np), tvm.runtime.tensor(b_np)) tvm.testing.assert_allclose(a_np[b_np], c.numpy()) def test_select_simplify(): placeholder = te.placeholder([1, 128, 10, 10, 4], dtype="float32") tensor = topi.nn.adaptive_pool(placeholder, [1, 1], "avg", "NCHW4c") result = te.create_prim_func([placeholder, tensor]) script_func = result.script() # There should be no Select assert script_func.find("Select") == -1 # There should be no undefined vars assert script_func.find("Var") == -1 def test_tensor_attr(): k = te.reduce_axis((0, 128), "k") A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") C = te.compute( (128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C", attrs={"layout_free_placeholders": [B]}, ) func = te.create_prim_func([A, B, C]) rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) @T.prim_func(s_tir=True) def expected_layout_attr( 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) @T.prim_func(s_tir=True) def expected_layout_attr_int64( A: T.Buffer((T.int64(128), T.int64(128)), "float32"), B: T.Buffer((T.int64(128), T.int64(128)), "float32"), D: T.Buffer((T.int64(128), T.int64(128)), "float32"), ): T.func_attr({"global_symbol": "main", "tirx.noalias": True, "layout_free_buffers": [1]}) C = T.sblock_alloc_buffer([T.int64(128), T.int64(128)], dtype="float32") for x, y, k in T.grid(T.int64(128), T.int64(128), T.int64(128)): with T.sblock("C"): v_x, v_y, v_k = T.axis.remap("SSR", [x, y, k]) T.reads(A[v_x, v_k], B[v_y, v_k]) T.writes(C[v_x, v_y]) with T.init(): C[v_x, v_y] = T.float32(0) C[v_x, v_y] = C[v_x, v_y] + A[v_x, v_k] * B[v_y, v_k] for x, y in T.grid(T.int64(128), T.int64(128)): with T.sblock("D"): T.sblock_attr({"layout_free_placeholders": [C]}) v_x, v_y = T.axis.remap("SS", [x, y]) T.reads(C[v_x, v_y]) T.writes(D[v_x, v_y]) D[v_x, v_y] = C[v_x, v_y] + T.float32(1) @pytest.mark.parametrize( "index_dtype_override, expected", [(None, expected_layout_attr), ("int64", expected_layout_attr_int64)], ) def test_tensor_layout_attr(index_dtype_override, expected): k = te.reduce_axis((0, 128), "k") A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") C = te.compute( (128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C", attrs={"layout_free_placeholders": [B]}, ) D = te.compute( (128, 128), lambda x, y: C[x, y] + 1, name="D", attrs={"layout_free_placeholders": [C]}, ) func = te.create_prim_func([A, B, D], index_dtype_override=index_dtype_override) tvm.ir.assert_structural_equal(func, expected) def te_argmax_idx_val(): def f_combine(x, y): lhs = tvm.tirx.Select((x[1] >= y[1]), x[0], y[0]) rhs = tvm.tirx.Select((x[1] >= y[1]), x[1], y[1]) return lhs, rhs def f_identity(dtype0: tvm.DataType, dtype1: tvm.DataType): return tvm.tirx.const(-1, dtype0), tvm.te.min_value(dtype1) argmax = te.comm_reducer(f_combine, f_identity, name="argmax") m = te.var("m") n = te.var("n") idx = te.placeholder((m, n), name="idx", dtype="int32") val = te.placeholder((m, n), name="val", dtype="float32") k = te.reduce_axis((0, n), "k") max_idx, max_val = te.compute( (m,), lambda i: argmax((idx[i, k], val[i, k]), axis=k), name="argmax" ) return [idx, val, max_idx, max_val] @T.prim_func(s_tir=True) def tir_argmax_idx_val( var_idx: T.handle, var_val: T.handle, var_argmax_v0: T.handle, var_argmax_v1: T.handle ) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) m = T.int32() n = T.int32() idx = T.match_buffer(var_idx, [m, n], dtype="int32") val = T.match_buffer(var_val, [m, n], dtype="float32") argmax_v0 = T.match_buffer(var_argmax_v0, [m], dtype="int32") argmax_v1 = T.match_buffer(var_argmax_v1, [m], dtype="float32") for i0, i1 in T.grid(m, n): with T.sblock("argmax"): i, k = T.axis.remap("SR", [i0, i1]) T.reads(val[i, k], idx[i, k]) T.writes(argmax_v0[i], argmax_v1[i]) with T.init(): argmax_v0[i] = T.int32(-1) argmax_v1[i] = T.min_value("float32") v_argmax_v0: T.let[T.int32] = T.Select( argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k] ) v_argmax_v1: T.let[T.float32] = T.Select( argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k] ) argmax_v0[i] = v_argmax_v0 argmax_v1[i] = v_argmax_v1 def te_argmax_val_idx(): def f_combine(x, y): lhs = tvm.tirx.Select((x[0] >= y[0]), x[0], y[0]) rhs = tvm.tirx.Select((x[0] >= y[0]), x[1], y[1]) return lhs, rhs def f_identity(dtype0: tvm.DataType, dtype1: tvm.DataType): return tvm.te.min_value(dtype0), tvm.tirx.const(-1, dtype1) argmax = te.comm_reducer(f_combine, f_identity, name="argmax") m = te.var("m") n = te.var("n") val = te.placeholder((m, n), name="val", dtype="float32") idx = te.placeholder((m, n), name="idx", dtype="int32") k = te.reduce_axis((0, n), "k") max_val, max_idx = te.compute( (m,), lambda i: argmax((val[i, k], idx[i, k]), axis=k), name="argmax" ) return [val, idx, max_val, max_idx] @T.prim_func(s_tir=True) def tir_argmax_val_idx( var_val: T.handle, var_idx: T.handle, var_argmax_v0: T.handle, var_argmax_v1: T.handle ) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) m = T.int32() n = T.int32() val = T.match_buffer(var_val, [m, n], dtype="float32") idx = T.match_buffer(var_idx, [m, n], dtype="int32") argmax_v0 = T.match_buffer(var_argmax_v0, [m], dtype="float32") argmax_v1 = T.match_buffer(var_argmax_v1, [m], dtype="int32") for i0, i1 in T.grid(m, n): with T.sblock("argmax"): i, k = T.axis.remap("SR", [i0, i1]) T.reads(val[i, k], idx[i, k]) T.writes(argmax_v0[i], argmax_v1[i]) with T.init(): argmax_v0[i] = T.min_value("float32") argmax_v1[i] = T.int32(-1) v_argmax_v0: T.let[T.float32] = T.Select( argmax_v0[i] >= val[i, k], argmax_v0[i], val[i, k] ) v_argmax_v1: T.let[T.int32] = T.Select( argmax_v0[i] >= val[i, k], argmax_v1[i], idx[i, k] ) argmax_v0[i] = v_argmax_v0 argmax_v1[i] = v_argmax_v1 def test_argmax_idx_val(): _check_workload(te_argmax_idx_val, tir_argmax_idx_val) def test_argmax_val_idx(): _check_workload(te_argmax_val_idx, tir_argmax_val_idx) def test_int64_indices(): n = te.var("n", "int64") A = te.placeholder((n,), name="A") B = te.compute(A.shape, lambda *i: A(*i) + 1, name="B") prim_func = te.create_prim_func([A, B]) loop = prim_func.body.block.body assert loop.loop_var.ty.dtype == "int64" assert loop.min.ty.dtype == "int64" assert loop.extent.ty.dtype == "int64" def test_zero_dim_add(): def te_func(): a = te.placeholder((), name="a", dtype="int32") b = te.placeholder((), name="b", dtype="int32") c = te.compute(a.shape, lambda *i: a(*i) + b(*i), name="c") return [a, b, c] @T.prim_func(s_tir=True) def expected( a: T.Buffer((), "int32"), b: T.Buffer((), "int32"), c: T.Buffer((), "int32"), ) -> None: T.func_attr({"global_symbol": "main", "tirx.noalias": True}) with T.sblock("root"): T.reads() T.writes() with T.sblock("c"): vi = T.axis.spatial(1, 0) T.reads(a[()], b[()]) T.writes(c[()]) c[()] = a[()] + b[()] _check_workload(te_func, expected) def te_reshape(): # The following is possible to be generated by TOPI. So we test this case. A = te.placeholder((tvm.tirx.IntImm("int64", 2), tvm.tirx.IntImm("int64", 4)), name="A") B = topi.reshape(A, (4, 2)) return [A, B] @T.prim_func(s_tir=True) def tir_reshape( A: T.Buffer((T.int64(2), T.int64(4)), "float32"), T_reshape: T.Buffer((T.int64(4), T.int64(2)), "float32"), ): T.func_attr({"global_symbol": "main", "tirx.noalias": True}) for i0, i1 in T.grid(T.int64(4), T.int64(2)): with T.sblock("T_reshape"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads( A[ (ax0 * T.int64(2) + ax1) % T.int64(8) // T.int64(4), (ax0 * T.int64(2) + ax1) % T.int64(4), ] ) T.writes(T_reshape[ax0, ax1]) T_reshape[ax0, ax1] = A[ (ax0 * T.int64(2) + ax1) % T.int64(8) // T.int64(4), (ax0 * T.int64(2) + ax1) % T.int64(4), ] def test_reshape(): _check_workload(te_reshape, tir_reshape, index_dtype_override="int64") def te_resize2d_symbolic(): oh = tirx.Var("oh", "int64") ow = tirx.Var("ow", "int64") roi = (0.0, 0.0, 0.0, 0.0) A = te.placeholder((2, 3, 128, 128), "float32", name="A") B = topi.image.resize2d( A, roi, size=(oh, ow), method="nearest_neighbor", coordinate_transformation_mode="asymmetric", rounding_method="round", ) return [A, B] @T.prim_func(s_tir=True) def tir_resize2d_symbolic( A: T.Buffer((T.int64(2), T.int64(3), T.int64(128), T.int64(128)), "float32"), var_resize: T.handle, ): T.func_attr({"global_symbol": "main", "tirx.noalias": True}) oh = T.int64() ow = T.int64() resize = T.match_buffer(var_resize, [T.int64(2), T.int64(3), oh, ow], dtype="float32") for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), oh, ow): with T.sblock("resize"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(A[v_i0, v_i1, T.int64(0) : T.int64(128), T.int64(0) : T.int64(128)]) T.writes(resize[v_i0, v_i1, v_i2, v_i3]) resize[v_i0, v_i1, v_i2, v_i3] = A[ v_i0, v_i1, T.max( T.min( T.Cast( "int64", T.round( T.float32(128) / T.Cast("float32", oh) * T.Cast("float32", v_i2), dtype="float32", ), ), T.int64(127), ), T.int64(0), ), T.max( T.min( T.Cast( "int64", T.round( T.float32(128) / T.Cast("float32", ow) * T.Cast("float32", v_i3), dtype="float32", ), ), T.int64(127), ), T.int64(0), ), ] def test_resize2d_symbolic(): _check_workload(te_resize2d_symbolic, tir_resize2d_symbolic, index_dtype_override="int64") def test_extern_with_explicit_buffer_access(): def te_extern(): A = te.placeholder((128, 128), name="A") B = te.placeholder((128, 128), name="B") P = te.placeholder((1,), name="P") C = te.extern( (128, 128), [A, B, P], lambda ins, outs: tvm.tirx.call_extern( "", "myfunc", ins[0].data, ins[1].data, outs[0].data, ins[2][0] ), name="C", ) return [A, B, P, C] @T.prim_func(s_tir=True) def tir_extern(var_A: T.handle, var_B: T.handle, var_P: T.handle, var_C: T.handle): T.func_attr({"global_symbol": "main", "tirx.noalias": True}) A = T.match_buffer(var_A, [128, 128], dtype="float32", offset_factor=1) B = T.match_buffer(var_B, [128, 128], dtype="float32", offset_factor=1) P = T.match_buffer(var_P, [1], dtype="float32", offset_factor=1) C = T.match_buffer(var_C, [128, 128], dtype="float32", offset_factor=1) with T.sblock("C"): T.reads() T.writes() T.call_extern("myfunc", A.data, B.data, C.data, P[0], dtype="") _check_workload(te_extern, tir_extern) def te_slice_with_var_input(): idx = te.var("idx", dtype="int64") m = te.var("m", dtype="int64") n = te.var("n", dtype="int64") tensor = te.placeholder((m, n), name="tensor") slice0 = te.compute((idx, n), lambda i, j: tensor[i, j], name="slice") return [tensor, idx, slice0] @T.prim_func(s_tir=True) def tir_slice_with_var_input(var_tensor: T.handle, idx: T.int64, var_slice: T.handle): T.func_attr({"tirx.noalias": True, "global_symbol": "main"}) m, n = T.int64(), T.int64() tensor = T.match_buffer(var_tensor, (m, n)) slice = T.match_buffer(var_slice, (idx, n)) # with T.sblock("root"): for i, j in T.grid(idx, n): with T.sblock("slice"): v_i = T.axis.spatial(idx, i) v_j = T.axis.spatial(n, j) T.reads(tensor[v_i, v_j]) T.writes(slice[v_i, v_j]) slice[v_i, v_j] = tensor[v_i, v_j] def test_with_var_input(): _check_workload(te_slice_with_var_input, tir_slice_with_var_input, index_dtype_override="int64") def test_loop_aware_initial_value(): """Test initial value aware of spatial iter position""" @T.prim_func(s_tir=True) def tir_workload(var_a: T.handle, var_b: T.handle, var_sum_red: T.handle): T.func_attr({"tirx.noalias": True, "global_symbol": "main"}) a = T.match_buffer(var_a, (5, 5)) b = T.match_buffer(var_b, (5,)) sum_red = T.match_buffer(var_sum_red, (5,)) for i, ax in T.grid(5, 5): with T.sblock("sum_red"): v_i, v_ax = T.axis.remap("SR", [i, ax]) T.reads(b[v_i], a[v_i, v_ax]) T.writes(sum_red[v_i]) with T.init(): sum_red[v_i] = b[v_i] sum_red[v_i] = sum_red[v_i] + a[v_i, v_ax] def te_workload(): data = te.placeholder((5, 5), "float32", "a") init = te.placeholder((5,), "float32", "b") ax = te.reduce_axis((0, 5), "ax") sum_red = te.compute( (5,), lambda i: te.comm_reducer( lambda x, y: x + y, lambda t: init[i], )(data[i, ax], axis=[ax]), name="sum_red", ) return [data, init, sum_red] _check_workload(te_workload, tir_workload) def test_loop_aware_reducer_combiner(): """Test combiner aware of spatial iter position""" @T.prim_func(s_tir=True) def tir_workload(var_a: T.handle, var_b: T.handle, var_sum_red: T.handle): T.func_attr({"tirx.noalias": True, "global_symbol": "main"}) a = T.match_buffer(var_a, (5, 5)) b = T.match_buffer(var_b, (5,)) sum_red = T.match_buffer(var_sum_red, (5,)) for i, ax in T.grid(5, 5): with T.sblock("sum_red"): v_i = T.axis.spatial(5, i) v_ax = T.axis.reduce(5, ax) T.reads(a[v_i, 0:5]) T.writes(sum_red[v_i]) with T.init(): sum_red[v_i] = T.float32(0.0) sum_red[v_i] = T.if_then_else( a[v_i, sum_red[v_i]] < a[v_i, v_ax], sum_red[v_i], T.Cast("float32", v_ax) ) def te_workload(): data = te.placeholder((5, 5), "float32", "a") init = te.placeholder((5,), "float32", "b") ax = te.reduce_axis((0, 5), "ax") sum_red = te.compute( (5,), lambda i: te.comm_reducer( lambda x, y: te.if_then_else(data[i, x] < y, x, ax), lambda _: te.const(0, "float32"), )(data[i, ax], axis=[ax]), name="sum_red", ) return [data, init, sum_red] _check_workload(te_workload, tir_workload) def test_adaptive_pooling_window(): @T.prim_func(s_tir=True) def tir_workload( x: T.Buffer((1, 1024, 16, 40), "float32"), adaptive_pool_avg: T.Buffer((1, 1024, 12, 30), "float32"), ): T.func_attr({"tirx.noalias": True, "global_symbol": "main"}) # fmt: off adaptive_pool_sum = T.sblock_alloc_buffer((1, 1024, 12, 30)) for ax0, ax1, ax2, ax3 in T.grid(1, 1024, 12, 30): with T.sblock("adaptive_pool_sum_1"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(x[v_ax0, v_ax1, v_ax2 * 16 // 12:v_ax2 * 16 // 12 + ((v_ax2 % 3 * 4 + 16) // 12 + 1), v_ax3 * 40 // 30:v_ax3 * 40 // 30 + ((v_ax3 % 3 * 10 + 40) // 30 + 1)]) T.writes(adaptive_pool_sum[v_ax0, v_ax1, v_ax2, v_ax3]) for rv0, rv1 in T.grid((v_ax2 % 3 * 4 + 16) // 12 + 1, (v_ax3 % 3 * 10 + 40) // 30 + 1): with T.sblock("adaptive_pool_sum"): v_ax0_1 = T.axis.spatial((v_ax0, v_ax0 + 1), v_ax0) v_ax1_1 = T.axis.spatial((v_ax1, v_ax1 + 1), v_ax1) v_ax2_1 = T.axis.spatial((v_ax2, v_ax2 + 1), v_ax2) v_ax3_1 = T.axis.spatial((v_ax3, v_ax3 + 1), v_ax3) v_rv0, v_rv1 = T.axis.remap("RR", [rv0, rv1]) T.reads(x[v_ax0_1, v_ax1_1, v_ax2_1 * 16 // 12 + v_rv0, v_ax3_1 * 40 // 30 + v_rv1]) T.writes(adaptive_pool_sum[v_ax0_1, v_ax1_1, v_ax2_1, v_ax3_1]) with T.init(): adaptive_pool_sum[v_ax0_1, v_ax1_1, v_ax2_1, v_ax3_1] = T.float32(0.0) adaptive_pool_sum[v_ax0_1, v_ax1_1, v_ax2_1, v_ax3_1] = adaptive_pool_sum[v_ax0_1, v_ax1_1, v_ax2_1, v_ax3_1] + x[v_ax0_1, v_ax1_1, v_ax2_1 * 16 // 12 + v_rv0, v_ax3_1 * 40 // 30 + v_rv1] for ax0, ax1, ax2, ax3 in T.grid(1, 1024, 12, 30): with T.sblock("adaptive_pool_avg"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(adaptive_pool_sum[v_ax0, v_ax1, v_ax2, v_ax3]) T.writes(adaptive_pool_avg[v_ax0, v_ax1, v_ax2, v_ax3]) T.sblock_attr({"schedule_rule": "meta_schedule.adaptive_pool_avg"}) adaptive_pool_avg[v_ax0, v_ax1, v_ax2, v_ax3] = adaptive_pool_sum[v_ax0, v_ax1, v_ax2, v_ax3] / (T.Cast("float32", (v_ax2 % 3 * 4 + 16) // 12 + 1) * T.Cast("float32", (v_ax3 % 3 * 10 + 40) // 30 + 1)) # fmt: on def te_workload(): x = te.placeholder([1, 1024, 16, 40], "float32", "x") y = topi.nn.adaptive_pool(x, [12, 30], pool_type="avg") f = te.create_prim_func([x, y]) return [x, y] _check_workload(te_workload, tir_workload) def test_global_pool(): # fix the issue-17938 data = te.placeholder((1, 1, 32, 32), dtype="int8", name="data") op_output = topi.nn.global_pool(data=data, pool_type="avg", layout="NCHW") f = te.create_prim_func([data, op_output]) assert f def test_nested_reduce_domain_dependency(): @T.prim_func(s_tir=True) def tir_workload( x: T.Buffer((8, 8, 8, 8, 8), "float32"), compute: T.Buffer((8, 8, 8), "float32") ): T.func_attr({"tirx.noalias": True, "global_symbol": "main"}) for i0, i1, i2 in T.grid(8, 8, 8): with T.sblock("compute_2"): v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2]) T.reads(x[v_i0, v_i1, v_i2, 0:v_i1, 0 : v_i1 - 1]) T.writes(compute[v_i0, v_i1, v_i2]) for rv in range(v_i1): with T.sblock("compute_1"): v_i0_1 = T.axis.spatial((v_i0, v_i0 + 1), v_i0) v_i1_1 = T.axis.spatial((v_i1, v_i1 + 1), v_i1) v_i2_1 = T.axis.spatial((v_i2, v_i2 + 1), v_i2) v_rv = T.axis.reduce(v_i1, rv) T.reads(x[v_i0_1, v_i1_1, v_i2_1, v_rv, 0:v_rv]) T.writes(compute[v_i0_1, v_i1_1, v_i2_1]) with T.init(): compute[v_i0_1, v_i1_1, v_i2_1] = T.float32(0.0) for rv_1 in range(v_rv): with T.sblock("compute"): v_i0_2 = T.axis.spatial((v_i0_1, v_i0_1 + 1), v_i0_1) v_i1_2 = T.axis.spatial((v_i1_1, v_i1_1 + 1), v_i1_1) v_i2_2 = T.axis.spatial((v_i2_1, v_i2_1 + 1), v_i2_1) v_rv_1 = T.axis.reduce((v_rv, v_rv + 1), v_rv) v_rv_2 = T.axis.reduce(v_rv, rv_1) T.reads(x[v_i0_2, v_i1_2, v_i2_2, v_rv_1, v_rv_2]) T.writes(compute[v_i0_2, v_i1_2, v_i2_2]) with T.init(): compute[v_i0_2, v_i1_2, v_i2_2] = T.float32(0.0) compute[v_i0_2, v_i1_2, v_i2_2] = ( compute[v_i0_2, v_i1_2, v_i2_2] + x[v_i0_2, v_i1_2, v_i2_2, v_rv_1, v_rv_2] ) def te_workload(): x = te.placeholder([8, 8, 8, 8, 8], "float32", "x") def fcompute(*axes): r1 = te.reduce_axis(tvm.ir.Range.from_min_extent(0, axes[1])) r2 = te.reduce_axis(tvm.ir.Range.from_min_extent(0, r1)) all_axes = [*axes, r1, r2] return te.sum(x(*all_axes), [r1, r2]) y = te.compute([8, 8, 8], fcompute) f = te.create_prim_func([x, y]) return [x, y] _check_workload(te_workload, tir_workload) if __name__ == "__main__": tvm.testing.main()