# 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. import tvm import tvm.testing from tvm import relax from tvm.relax.frontend import detach_params from tvm.relax.frontend.common import autopad from tvm.script import ir as I from tvm.script import tirx as T from tvm.script.parser import relax as R def test_detach_params(): @R.function def func(x: R.Tensor((2, 3), "float32")): return x param = tvm.runtime.empty((3,), "float32") mod = tvm.IRModule({"func": func.with_attr("params", [param])}) detached_mod, detached_params = detach_params(mod) tvm.ir.assert_structural_equal(detached_mod, tvm.IRModule({"func": func})) assert len(detached_params) == 1 assert "func" in detached_params assert isinstance(detached_params["func"], list) assert len(detached_params["func"]) == 1 tvm.testing.assert_allclose(detached_params["func"][0].numpy(), param.numpy()) class TestAutopad: def _test_autopad(self, pad_type, expected): bb = relax.BlockBuilder() input_shape = (1, 1, 4, 4) x = relax.Var("x", relax.TensorType(input_shape, "float32")) with bb.function("main", [x]): with bb.dataflow(): result = autopad( bb, x, strides=[2, 2], kernel_shape=[3, 3], dilations=(1, 1), pad_type=pad_type, deconv=False, mode="SAME_UPPER", pad_value=0.0, ) out = bb.emit_output(result) bb.emit_func_output(out) tvm.ir.assert_structural_equal(bb.get(), expected) def test_constant(self): @I.ir_module(s_tir=True) class expected: @T.prim_func(private=True, s_tir=True) def pad( x: T.Buffer((T.int64(1), T.int64(1), T.int64(4), T.int64(4)), "float32"), PadInput: T.Buffer((T.int64(1), T.int64(1), T.int64(5), T.int64(5)), "float32"), ): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(5), T.int64(5)): with T.sblock("PadInput"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(x[v_i0, v_i1, v_i2, v_i3]) T.writes(PadInput[v_i0, v_i1, v_i2, v_i3]) PadInput[v_i0, v_i1, v_i2, v_i3] = T.if_then_else( T.int64(0) <= v_i2 and v_i2 < T.int64(4) and T.int64(0) <= v_i3 and v_i3 < T.int64(4), x[v_i0, v_i1, v_i2, v_i3], T.float32(0.0), ) @R.function def main(x: R.Tensor((1, 1, 4, 4), dtype="float32")) -> R.Tensor( (1, 1, 5, 5), dtype="float32" ): cls = expected with R.dataflow(): lv = R.call_tir(cls.pad, (x,), out_ty=R.Tensor((1, 1, 5, 5), dtype="float32")) gv: R.Tensor((1, 1, 5, 5), dtype="float32") = lv R.output(gv) return gv self._test_autopad("constant", expected) def test_edge(self): @I.ir_module(s_tir=True) class expected: @T.prim_func(private=True, s_tir=True) def replicate_pad( x: T.Buffer((T.int64(1), T.int64(1), T.int64(4), T.int64(4)), "float32"), ReplicatePadInput: T.Buffer( (T.int64(1), T.int64(1), T.int64(5), T.int64(5)), "float32" ), ): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(5), T.int64(5)): with T.sblock("ReplicatePadInput"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads( x[ T.int64(0), T.int64(0), T.int64(0) : T.int64(4), T.int64(0) : T.int64(4), ] ) T.writes(ReplicatePadInput[v_i0, v_i1, v_i2, v_i3]) ReplicatePadInput[v_i0, v_i1, v_i2, v_i3] = x[ T.if_then_else( v_i0 < T.int64(0), T.int64(0), T.if_then_else(T.int64(1) <= v_i0, T.int64(0), v_i0), ), T.if_then_else( v_i1 < T.int64(0), T.int64(0), T.if_then_else(T.int64(1) <= v_i1, T.int64(0), v_i1), ), T.if_then_else( v_i2 < T.int64(0), T.int64(0), T.if_then_else(T.int64(4) <= v_i2, T.int64(3), v_i2), ), T.if_then_else( v_i3 < T.int64(0), T.int64(0), T.if_then_else(T.int64(4) <= v_i3, T.int64(3), v_i3), ), ] @R.function def main(x: R.Tensor((1, 1, 4, 4), dtype="float32")) -> R.Tensor( (1, 1, 5, 5), dtype="float32" ): cls = expected with R.dataflow(): lv = R.call_tir( cls.replicate_pad, (x,), out_ty=R.Tensor((1, 1, 5, 5), dtype="float32") ) gv: R.Tensor((1, 1, 5, 5), dtype="float32") = lv R.output(gv) return gv self._test_autopad("edge", expected) def test_reflect(self): @I.ir_module(s_tir=True) class expected: @T.prim_func(private=True, s_tir=True) def mirror_pad( x: T.Buffer((T.int64(1), T.int64(1), T.int64(4), T.int64(4)), "float32"), MirrorPadInput: T.Buffer( (T.int64(1), T.int64(1), T.int64(5), T.int64(5)), "float32" ), ): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(5), T.int64(5)): with T.sblock("MirrorPadInput"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(x[v_i0, v_i1, T.int64(0) : T.int64(4), T.int64(0) : T.int64(4)]) T.writes(MirrorPadInput[v_i0, v_i1, v_i2, v_i3]) MirrorPadInput[v_i0, v_i1, v_i2, v_i3] = x[ v_i0, v_i1, T.if_then_else( T.int64(4) <= v_i2, T.int64(6) - v_i2, T.if_then_else(v_i2 < T.int64(0), v_i2 * T.int64(-1), v_i2), ), T.if_then_else( T.int64(4) <= v_i3, T.int64(6) - v_i3, T.if_then_else(v_i3 < T.int64(0), v_i3 * T.int64(-1), v_i3), ), ] @R.function def main(x: R.Tensor((1, 1, 4, 4), dtype="float32")) -> R.Tensor( (1, 1, 5, 5), dtype="float32" ): cls = expected with R.dataflow(): lv = R.call_tir( cls.mirror_pad, (x,), out_ty=R.Tensor((1, 1, 5, 5), dtype="float32") ) gv: R.Tensor((1, 1, 5, 5), dtype="float32") = lv R.output(gv) return gv self._test_autopad("reflect", expected) if __name__ == "__main__": tvm.testing.main()