# 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: F401 import sys from collections import defaultdict import numpy import pytest import tvm.testing from tvm import tirx from tvm.s_tir.schedule.testing import verify_trace_roundtrip from tvm.script import tirx as T # pylint: disable=no-member,invalid-name,unused-variable @T.prim_func(s_tir=True) def elementwise(a: T.handle, b: T.handle) -> None: A = T.match_buffer(a, (128, 257, 1470)) B = T.match_buffer(b, (128, 257, 1470)) for i, j, k in T.grid(128, 257, 1470): with T.sblock("B"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) B[vi, vj, vk] = A[vi, vj, vk] * 2.0 @T.prim_func(s_tir=True) def tiled_conv2d_with_padding( inputs: T.Buffer((1, 224, 224, 3), "float32"), weight: T.Buffer((7, 7, 3, 64), "float32"), conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"), ) -> None: PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32") for i0, i1, i2, i3 in T.grid(1, 230, 230, 3): with T.sblock("PadInput"): i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1]) T.writes(PadInput[i0_1, i1_1, i2_1, i3_1]) PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else( 3 <= i1_1 and i1_1 < 227 and 3 <= i2_1 and i2_1 < 227, inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1], T.float32(0), dtype="float32", ) for ( i0_0, i1_0, i2_0, i3_0, i0_1_1, i1_1_1, i2_1_1, i3_1_1, i4_0, i5_0, i6_0, i0_2, i1_2, i2_2, i3_2, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3, i3_3, ) in T.grid(1, 1, 4, 1, 1, 2, 4, 1, 7, 7, 1, 1, 1, 1, 1, 1, 1, 3, 1, 56, 7, 64): with T.sblock("conv2d_nhwc"): n = T.axis.spatial(1, 0) h = T.axis.spatial(112, i1_1_1 * 56 + i1_3) w = T.axis.spatial(112, i2_0 * 28 + i2_1_1 * 7 + i2_3) co, rh, rw, rc = T.axis.remap("SRRR", [i3_3, i4_0, i5_0, i6_1]) T.reads( conv2d_nhwc[n, h, w, co], PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc], weight[rh, rw, rc, co], ) T.writes(conv2d_nhwc[n, h, w, co]) with T.init(): conv2d_nhwc[n, h, w, co] = T.float32(0) conv2d_nhwc[n, h, w, co] = ( conv2d_nhwc[n, h, w, co] + PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc] * weight[rh, rw, rc, co] ) # pylint: enable=no-member,invalid-name,unused-variable def test_sample_categorical(): """Test sample categorical sampling function""" n = 1000 sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all") counter = defaultdict(int) candidates = [5, 2, 7, 1] probs = [0.15, 0.55, 0.05, 0.25] for _ in range(n): v = sch.get(sch.sample_categorical(candidates, probs)) counter[v] += 1 for i, prob in enumerate(probs): assert (prob - 0.07) * n <= counter[candidates[i]] <= (prob + 0.07) * n verify_trace_roundtrip(sch, mod=elementwise) def test_sample_categorical_copy(): """Check the random variable sampling results after schedule copy""" n = 100 sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all") candidates = [1, 2, 3, 4] probs = [0.1, 0.2, 0.3, 0.4] rv_decisions = [] for _ in range(n): rv = sch.sample_categorical(candidates, probs) # pylint: disable=invalid-name rv_decisions.append((rv, sch.get(rv))) sch_copy = sch.copy() for rv, decision in rv_decisions: # pylint: disable=invalid-name decision_copy = sch_copy.get(rv) assert int(decision) == int(decision_copy) def test_sample_categorical_serialize(): """Check the random variable sampling results after schedule serialization""" n = 100 sch = tvm.s_tir.Schedule(elementwise, seed=42, debug_mask="all") candidates = [5, 6, 7, 8] probs = [0.23, 0.19, 0.37, 0.21] decisions = [] for _ in range(n): rv = sch.get(sch.sample_categorical(candidates, probs)) # pylint: disable=invalid-name decisions.append(rv) new_sch = verify_trace_roundtrip(sch, mod=elementwise) for i, new_inst in enumerate(new_sch.trace.insts): assert decisions[i] == candidates[new_sch.trace.decisions[new_inst]] def test_sample_perfect_tile_power_of_two(): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") i, _, _ = sch.get_loops(sch.get_sblock("B")) factors = sch.sample_perfect_tile(i, n=4) factors = [sch.get(i) for i in factors] prod = factors[0] * factors[1] * factors[2] * factors[3] assert prod == 128 verify_trace_roundtrip(sch, mod=elementwise) def test_sample_perfect_tile_prime(): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") _, i, _ = sch.get_loops(sch.get_sblock("B")) factors = sch.sample_perfect_tile(i, n=4) factors = [sch.get(i) for i in factors] prod = factors[0] * factors[1] * factors[2] * factors[3] assert prod == 257 verify_trace_roundtrip(sch, mod=elementwise) def test_sample_perfect_tile_composite(): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") _, _, i = sch.get_loops(sch.get_sblock("B")) factors = sch.sample_perfect_tile(i, n=4) factors = [sch.get(i) for i in factors] prod = factors[0] * factors[1] * factors[2] * factors[3] assert prod == 1470 verify_trace_roundtrip(sch, mod=elementwise) use_sugared_block = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True}) def test_sample_compute_location(use_sugared_block): n = 100 sch = tvm.s_tir.Schedule(tiled_conv2d_with_padding, seed=42, debug_mask="all") if use_sugared_block: pad_input = "PadInput" else: pad_input = sch.get_sblock("PadInput") decision_dict = dict() for _ in range(n): _ = sch.sample_compute_location(pad_input) # pylint: disable=invalid-name decision = sch.trace.decisions[sch.trace.insts[-1]] decision_dict[decision] = decision_dict[decision] + 1 if decision in decision_dict else 1 n_candidates = 8 expected_rate = 1.0 / n_candidates for _, cnt in decision_dict.items(): numpy.testing.assert_allclose(expected_rate, cnt / n, atol=0.04) def test_sample_perfect_tile_after_copy(): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") sch_copy = sch.copy() _, _, i = sch.get_loops(sch.get_sblock("B")) sch.sample_perfect_tile(i, n=4) _, _, i = sch_copy.get_loops(sch_copy.get_sblock("B")) # Hangs if ForkSeed is not invoked when copying a schedule sch_copy.sample_perfect_tile(i, n=4) def test_sample_perfect_tile_on_dynamic_loops(): """Currently dynamic loop is trivially tiled""" @T.prim_func(s_tir=True) def workload(a: T.handle) -> None: n = T.int32() A = T.match_buffer(a, (n, 1024)) for i, j in T.grid(n, 1024): with T.sblock("B"): vi, vj = T.axis.remap("SS", [i, j]) A[vi, vj] = 1.0 sch = tvm.s_tir.Schedule(workload, debug_mask="all") di, si = sch.get_loops(sch.get_sblock("B")) factors = sch.sample_perfect_tile(si, n=4) factors = [sch.get(i) for i in factors] prod = factors[0] * factors[1] * factors[2] * factors[3] assert prod == 1024 factors = sch.sample_perfect_tile(di, n=4) assert factors[0] is None factors = [sch.get(i) for i in factors[1:]] prod = factors[0] * factors[1] * factors[2] assert prod == 1 verify_trace_roundtrip(sch, mod=workload) if __name__ == "__main__": tvm.testing.main()