diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 52122ce33369..ce1f284be6bc 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -616,21 +616,26 @@ def create_convert_map( nn.Flatten: self._flatten_module, ## call_function and call_method # unary + "abs": self._unary_op(relax.op.abs), "acos": self._unary_op(relax.op.acos), "acosh": self._unary_op(relax.op.acosh), "asin": self._unary_op(relax.op.asin), "asinh": self._unary_op(relax.op.asinh), "atan": self._unary_op(relax.op.atan), "atanh": self._unary_op(relax.op.atanh), + "ceil": self._unary_op(relax.op.ceil), "clamp": self._clamp, "cos": self._unary_op(relax.op.cos), "cosh": self._unary_op(relax.op.cosh), "dropout": lambda node: self.env[node.args[0]], + "erf": self._unary_op(relax.op.erf), "exp": self._unary_op(relax.op.exp), + "floor": self._unary_op(relax.op.floor), "gelu": self._gelu, "hardsigmoid": self._hardsigmoid, "hardswish": self._hardswish, "leaky_relu": self._leakyrelu, + "log": self._unary_op(relax.op.log), "log_softmax": self._log_softmax, "neg": self._unary_op(relax.op.negative), "relu": self._unary_op(relax.op.nn.relu), diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index d9857723b1f5..3c932f86c582 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -746,37 +746,6 @@ def main( verify_model(Einsum2(), [([5], "float32"), ([4], "float32")], {}, Expected2) -def test_relu(): - class ReLU0(Module): - def __init__(self): - super().__init__() - self.relu = torch.nn.ReLU() - - def forward(self, input): - return self.relu(input) - - class ReLU1(Module): - def forward(self, input): - return torch.nn.functional.relu(input) - - @tvm.script.ir_module - class expected: - @R.function - def main( - input_1: R.Tensor((10, 10), dtype="float32") - ) -> R.Tensor((10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((10, 10), dtype="float32") = R.nn.relu(input_1) - gv: R.Tensor((10, 10), dtype="float32") = lv - R.output(gv) - return gv - - input_info = [([10, 10], "float32")] - verify_model(ReLU0(), input_info, {}, expected) - verify_model(ReLU1(), input_info, {}, expected) - - @tvm.testing.requires_gpu def test_leakyrelu(): import torch @@ -814,30 +783,6 @@ def main( verify_model(LeakyReLU1(), input_info, {}, expected) -def test_relu6(): - class ReLU6(Module): - def __init__(self): - super().__init__() - self.relu6 = torch.nn.ReLU6() - - def forward(self, input): - return self.relu6(input) - - @tvm.script.ir_module - class expected: - @R.function - def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((10, 10), dtype="float32") = R.clip(input, 0, 6) - gv: R.Tensor((10, 10), dtype="float32") = lv - R.output(gv) - return gv - - input_info = [([10, 10], "float32")] - verify_model(ReLU6(), input_info, {}, expected) - - def test_maxpool2d(): input_info = [([1, 3, 10, 10], "float32")] @@ -1182,37 +1127,6 @@ def main( verify_model(model, input_info, binding, expected1) -def test_dropout(): - input_info = [([1, 3, 10, 10], "float32")] - - class Dropout1(Module): - def __init__(self): - super().__init__() - self.dropout = torch.nn.Dropout(0.5) - - def forward(self, input): - return self.dropout(input) - - class Dropout2(Module): - def forward(self, input): - return torch.dropout(input, 0.5, train=True) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1 - R.output(gv) - return gv - - verify_model(Dropout1(), input_info, {}, expected1) - verify_model(Dropout2(), input_info, {}, expected1) - - def test_stochastic_depth(): input_info = [([1, 3, 10, 10], "float32")] @@ -1519,109 +1433,6 @@ def main( verify_model(model, input_info, {}, expected1) -def test_silu(): - input_info = [([1, 3, 10, 10], "float32")] - - class SiLU(Module): - def __init__(self): - super().__init__() - self.silu = torch.nn.SiLU() - - def forward(self, input): - return self.silu(input) - - class SiLU2(Module): - def forward(self, input): - return torch.nn.functional.silu(input) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.silu(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(SiLU(), input_info, {}, expected1) - verify_model(SiLU2(), input_info, {}, expected1) - - -def test_hardsigmoid(): - input_info = [([1, 3, 10, 10], "float32")] - - class Hardsigmoid(torch.nn.Module): - def __init__(self): - super().__init__() - self.hs = torch.nn.Hardsigmoid() - - def forward(self, input): - return self.hs(input) - - class Hardsigmoid2(torch.nn.Module): - def forward(self, input): - return torch.nn.functional.hardsigmoid(input) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - inp_0: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32")) - lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6) - lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( - lv1, R.const(6, "float32") - ) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv2 - R.output(gv) - return gv - - verify_model(Hardsigmoid(), input_info, {}, expected1) - verify_model(Hardsigmoid2(), input_info, {}, expected1) - - -def test_hardswish(): - input_info = [([1, 3, 10, 10], "float32")] - - class Hardswish(torch.nn.Module): - def __init__(self): - super().__init__() - self.hs = torch.nn.Hardswish() - - def forward(self, input): - return self.hs(input) - - class Hardswish2(torch.nn.Module): - def forward(self, input): - return torch.nn.functional.hardswish(input) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - inp_0: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32")) - lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6) - lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( - lv1, R.const(6, "float32") - ) - lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(inp_0, lv2) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv3 - R.output(gv) - return gv - - verify_model(Hardswish(), input_info, {}, expected1) - verify_model(Hardswish2(), input_info, {}, expected1) - - def test_groupnorm(): import torch from torch.nn import Module @@ -1671,70 +1482,6 @@ def main( verify_model(model, input_info, binding, expected1) -def test_softmax(): - input_info = [([1, 3, 10, 10], "float32")] - - class Softmax(Module): - def __init__(self): - super().__init__() - self.sm = torch.nn.Softmax(dim=1) - - def forward(self, input): - return self.sm(input) - - class Softmax2(Module): - def forward(self, input): - return torch.nn.functional.softmax(input, dim=1) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.softmax(input_1, axis=1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Softmax(), input_info, {}, expected1) - verify_model(Softmax2(), input_info, {}, expected1) - - -def test_logsoftmax(): - input_info = [([1, 3, 10, 10], "float32")] - - class LogSoftmax(Module): - def __init__(self): - super().__init__() - self.lsm = torch.nn.LogSoftmax(dim=1) - - def forward(self, input): - return self.lsm(input) - - class LogSoftmax2(Module): - def forward(self, input): - return torch.nn.functional.log_softmax(input, dim=1) - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.log_softmax(input_1, axis=1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(LogSoftmax(), input_info, {}, expected1) - verify_model(LogSoftmax2(), input_info, {}, expected1) - - def test_binary(): input_info1 = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")] input_info2 = [([1, 3, 10, 10], "float32")] @@ -2183,288 +1930,306 @@ def main( verify_model(Slice2(), [([8, 16], "float32")], {}, expected2) -def test_unary(): +operator_basic_unary = [ + (torch.abs, R.abs), + (torch.acos, R.acos), + (torch.acosh, R.acosh), + (torch.asin, R.asin), + (torch.asinh, R.asinh), + (torch.atan, R.atan), + (torch.atanh, R.atanh), + (torch.ceil, R.ceil), + (torch.cos, R.cos), + (torch.cosh, R.cosh), + (torch.erf, R.erf), + (torch.exp, R.exp), + (torch.floor, R.floor), + (torch.log, R.log), + (torch.neg, R.negative), + (torch.round, R.round), + (torch.rsqrt, R.rsqrt), + (torch.sin, R.sin), + (torch.sinh, R.sinh), + (torch.sqrt, R.sqrt), + (torch.tan, R.tan), +] + + +@pytest.mark.parametrize("pytorch_op, relax_op", operator_basic_unary) +def test_basic_unary_ops(pytorch_op, relax_op): input_info = [([1, 3, 10, 10], "float32")] - # sin - class Sin(Module): + class Unary(Module): def forward(self, input): - return torch.sin(input) + return pytorch_op(input) @tvm.script.ir_module - class expected_sin: + class expected_unary: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sin(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(input_1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Sin(), input_info, {}, expected_sin) + verify_model(Unary(), input_info, {}, expected_unary) - # cos - class Cos(Module): - def forward(self, input): - return torch.cos(input) - - @tvm.script.ir_module - class expected_cos: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.cos(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - verify_model(Cos(), input_info, {}, expected_cos) +def test_extended_unary_ops(): + input_info = [([1, 3, 10, 10], "float32")] - # tan - class Tan(Module): + # clamp + class Clamp(Module): def forward(self, input): - return torch.tan(input) + return torch.clamp(input, min=0.1, max=0.5) @tvm.script.ir_module - class expected_tan: + class expected_clamp: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tan(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.1, 0.5) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Tan(), input_info, {}, expected_tan) - - # asin - class Asin(Module): - def forward(self, input): - return torch.asin(input) + verify_model(Clamp(), input_info, {}, expected_clamp) - @tvm.script.ir_module - class expected_asin: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.asin(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv + from tvm.relax.frontend.torch import from_fx - verify_model(Asin(), input_info, {}, expected_asin) + with pytest.raises( + ValueError, match="TVM only supports constant max value for torch.clamp/clip" + ): - # acos - class Acos(Module): - def forward(self, input): - return torch.acos(input) + class Clamp_Error(Module): + def forward(self, input): + return torch.clamp(input, min=0.5, max=None) - @tvm.script.ir_module - class expected_acos: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.acos(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv + gm = fx.symbolic_trace(Clamp_Error()) + from_fx(gm, input_info) - verify_model(Acos(), input_info, {}, expected_acos) + with pytest.raises( + ValueError, match="TVM only supports constant min value for torch.clamp/clip" + ): - # atan - class Atan(Module): - def forward(self, input): - return torch.atan(input) + class Clamp_Error(Module): + def forward(self, input): + return torch.clamp(input, min=input, max=input) - @tvm.script.ir_module - class expected_atan: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.atan(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv + gm = fx.symbolic_trace(Clamp_Error()) + from_fx(gm, input_info) - verify_model(Atan(), input_info, {}, expected_atan) + # dropout + class Dropout1(Module): + def __init__(self): + super().__init__() + self.dropout = torch.nn.Dropout(0.5) - # sinh - class Sinh(Module): def forward(self, input): - return torch.sinh(input) - - @tvm.script.ir_module - class expected_sinh: - @R.function - def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sinh(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Sinh(), input_info, {}, expected_sinh) + return self.dropout(input) - # cosh - class Cosh(Module): + class Dropout2(Module): def forward(self, input): - return torch.cosh(input) + return torch.dropout(input, 0.5, train=True) @tvm.script.ir_module - class expected_cosh: + class expected_dropout: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.cosh(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1 R.output(gv) return gv - verify_model(Cosh(), input_info, {}, expected_cosh) + verify_model(Dropout1(), input_info, {}, expected_dropout) + verify_model(Dropout2(), input_info, {}, expected_dropout) + + # gelu + class Gelu(Module): + def __init__(self): + super().__init__() + self.gelu = torch.nn.GELU() - # tanh - class Tanh(Module): def forward(self, input): - return torch.tanh(input) + return self.gelu(input) + + class Gelu2(Module): + def forward(self, input): + return torch.nn.functional.gelu(input) @tvm.script.ir_module - class expected_tanh: + class expected_gelu: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tanh(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.gelu(input_1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Tanh(), input_info, {}, expected_tanh) + verify_model(Gelu(), input_info, {}, expected_gelu) + verify_model(Gelu2(), input_info, {}, expected_gelu) + + # leaky_relu + test_leakyrelu() + + # hardsigmoid + class Hardsigmoid(torch.nn.Module): + def __init__(self): + super().__init__() + self.hs = torch.nn.Hardsigmoid() + + def forward(self, input): + return self.hs(input) - # asinh - class Asinh(Module): + class Hardsigmoid2(torch.nn.Module): def forward(self, input): - return torch.asinh(input) + return torch.nn.functional.hardsigmoid(input) @tvm.script.ir_module - class expected_asinh: + class expected_hardsigmoid: @R.function def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + inp_0: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.asinh(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32")) + lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6) + lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( + lv1, R.const(6, "float32") + ) + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv2 R.output(gv) return gv - verify_model(Asinh(), input_info, {}, expected_asinh) + verify_model(Hardsigmoid(), input_info, {}, expected_hardsigmoid) + verify_model(Hardsigmoid2(), input_info, {}, expected_hardsigmoid) + + # hardswish + class Hardswish(torch.nn.Module): + def __init__(self): + super().__init__() + self.hs = torch.nn.Hardswish() + + def forward(self, input): + return self.hs(input) - # acosh - class Acosh(Module): + class Hardswish2(torch.nn.Module): def forward(self, input): - return torch.acosh(input) + return torch.nn.functional.hardswish(input) @tvm.script.ir_module - class expected_acosh: + class expected_hardswish: @R.function def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + inp_0: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.acosh(input_1) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32")) + lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6) + lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( + lv1, R.const(6, "float32") + ) + lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(inp_0, lv2) + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv3 R.output(gv) return gv - verify_model(Acosh(), input_info, {}, expected_acosh) + verify_model(Hardswish(), input_info, {}, expected_hardswish) + verify_model(Hardswish2(), input_info, {}, expected_hardswish) + + # log_softmax + class LogSoftmax(Module): + def __init__(self): + super().__init__() + self.lsm = torch.nn.LogSoftmax(dim=1) + + def forward(self, input): + return self.lsm(input) - # atanh - class Atanh(Module): + class LogSoftmax2(Module): def forward(self, input): - return torch.atanh(input) + return torch.nn.functional.log_softmax(input, dim=1) @tvm.script.ir_module - class expected_atanh: + class expected_log_softmax: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.atanh(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.log_softmax(input_1, axis=1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Atanh(), input_info, {}, expected_atanh) + verify_model(LogSoftmax(), input_info, {}, expected_log_softmax) + verify_model(LogSoftmax2(), input_info, {}, expected_log_softmax) + + # relu + class ReLU0(Module): + def __init__(self): + super().__init__() + self.relu = torch.nn.ReLU() + + def forward(self, input): + return self.relu(input) - # exp - class Exp(Module): + class ReLU1(Module): def forward(self, input): - return torch.exp(input) + return torch.nn.functional.relu(input) @tvm.script.ir_module - class expected_exp: + class expected_relu: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Exp(), input_info, {}, expected_exp) + verify_model(ReLU0(), input_info, {}, expected_relu) + verify_model(ReLU1(), input_info, {}, expected_relu) + + # relu6 + class ReLU6(Module): + def __init__(self): + super().__init__() + self.relu6 = torch.nn.ReLU6() - # sqrt - class Sqrt(Module): def forward(self, input): - return torch.sqrt(input) + return self.relu6(input) @tvm.script.ir_module - class expected3: + class expected_relu6: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sqrt(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0, 6) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Sqrt(), input_info, {}, expected3) + verify_model(ReLU6(), input_info, {}, expected_relu6) # sigmoid class Sigmoid(Module): @@ -2480,7 +2245,7 @@ def forward(self, input): return torch.sigmoid(input) @tvm.script.ir_module - class expected4: + class expected_sigmoid: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") @@ -2492,65 +2257,68 @@ def main( R.output(gv) return gv - verify_model(Sigmoid(), input_info, {}, expected4) - verify_model(Sigmoid2(), input_info, {}, expected4) + verify_model(Sigmoid(), input_info, {}, expected_sigmoid) + verify_model(Sigmoid2(), input_info, {}, expected_sigmoid) + + # silu + class SiLU(Module): + def __init__(self): + super().__init__() + self.silu = torch.nn.SiLU() + + def forward(self, input): + return self.silu(input) - # round - class Round(Module): + class SiLU2(Module): def forward(self, input): - return torch.round(input) + return torch.nn.functional.silu(input) @tvm.script.ir_module - class expected5: + class expected_silu: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.round(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.silu(input_1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Round(), input_info, {}, expected5) + verify_model(SiLU(), input_info, {}, expected_silu) + verify_model(SiLU2(), input_info, {}, expected_silu) - -def test_gelu(): - input_info = [([1, 3, 10, 10], "float32")] - - class Gelu(Module): + # softmax + class Softmax(Module): def __init__(self): super().__init__() - self.gelu = torch.nn.GELU() + self.sm = torch.nn.Softmax(dim=1) def forward(self, input): - return self.gelu(input) + return self.sm(input) - class Gelu2(Module): + class Softmax2(Module): def forward(self, input): - return torch.nn.functional.gelu(input) + return torch.nn.functional.softmax(input, dim=1) @tvm.script.ir_module - class expected1: + class expected_softmax: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.gelu(input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.softmax(input_1, axis=1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Gelu(), input_info, {}, expected1) - verify_model(Gelu2(), input_info, {}, expected1) - - -def test_tanh(): - input_info = [([1, 3, 10, 10], "float32")] + verify_model(Softmax(), input_info, {}, expected_softmax) + verify_model(Softmax2(), input_info, {}, expected_softmax) + # tanh class Tanh(Module): def __init__(self): super().__init__() @@ -2564,7 +2332,7 @@ def forward(self, input): return torch.tanh(input) @tvm.script.ir_module - class expected1: + class expected_tanh: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") @@ -2576,55 +2344,58 @@ def main( R.output(gv) return gv - verify_model(Tanh(), input_info, {}, expected1) - verify_model(Tanh2(), input_info, {}, expected1) - + verify_model(Tanh(), input_info, {}, expected_tanh) + verify_model(Tanh2(), input_info, {}, expected_tanh) -def test_clamp(): - input_info = [([1, 3, 10, 10], "float32")] + # tril + class Tril(Module): + def forward(self, input): + return torch.tril(input, 1) - class Clamp(Module): + class InplaceTril(Module): def forward(self, input): - return torch.clamp(input, min=0.1, max=0.5) + input.tril_(1) + return input @tvm.script.ir_module - class expected1: + class expected_tril: @R.function def main( input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - # block 0 with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.1, 0.5) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tril(input_1, 1) gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv R.output(gv) return gv - verify_model(Clamp(), input_info, {}, expected1) - - from tvm.relax.frontend.torch import from_fx - - with pytest.raises( - ValueError, match="TVM only supports constant max value for torch.clamp/clip" - ): - - class Clamp_Error(Module): - def forward(self, input): - return torch.clamp(input, min=0.5, max=None) + verify_model(Tril(), input_info, {}, expected_tril) + verify_model(InplaceTril(), input_info, {}, expected_tril) - gm = fx.symbolic_trace(Clamp_Error()) - from_fx(gm, input_info) + # triu + class Triu(Module): + def forward(self, input): + return torch.triu(input, 1) - with pytest.raises( - ValueError, match="TVM only supports constant min value for torch.clamp/clip" - ): + class InplaceTriu(Module): + def forward(self, input): + input.triu_(1) + return input - class Clamp_Error(Module): - def forward(self, input): - return torch.clamp(input, min=input, max=input) + @tvm.script.ir_module + class expected_triu: + @R.function + def main( + input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.triu(input_1, 1) + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + R.output(gv) + return gv - gm = fx.symbolic_trace(Clamp_Error()) - from_fx(gm, input_info) + verify_model(Triu(), input_info, {}, expected_triu) + verify_model(InplaceTriu(), input_info, {}, expected_triu) def test_interpolate(): @@ -3085,64 +2856,6 @@ def forward(self, input): assert mod2["main"].body.blocks[0].bindings[0].value.data.dtype == "int64" -def test_tril(): - input_info = [([10, 10], "float32")] - - class Tril(Module): - def forward(self, input): - return torch.tril(input, 1) - - class InplaceTril(Module): - def forward(self, input): - input.tril_(1) - return input - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((10, 10), dtype="float32") - ) -> R.Tensor((10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((10, 10), dtype="float32") = R.tril(input_1, 1) - gv: R.Tensor((10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Tril(), input_info, {}, expected1) - verify_model(InplaceTril(), input_info, {}, expected1) - - -def test_triu(): - input_info = [([10, 10], "float32")] - - class Triu(Module): - def forward(self, input): - return torch.triu(input, 1) - - class InplaceTriu(Module): - def forward(self, input): - input.triu_(1) - return input - - @tvm.script.ir_module - class expected1: - @R.function - def main( - input_1: R.Tensor((10, 10), dtype="float32") - ) -> R.Tensor((10, 10), dtype="float32"): - # block 0 - with R.dataflow(): - lv: R.Tensor((10, 10), dtype="float32") = R.triu(input_1, 1) - gv: R.Tensor((10, 10), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Triu(), input_info, {}, expected1) - verify_model(InplaceTriu(), input_info, {}, expected1) - - def test_new_ones(): input_info = [([1, 2, 3], "float32")] @@ -3707,26 +3420,6 @@ def main( verify_model(MeanKeepDim(), [([256, 256], "float32")], {}, Expected2) -def test_rsqrt(): - class Rsqrt(Module): - def forward(self, input): - return torch.rsqrt(input) - - @I.ir_module - class Expected1: - @R.function - def main( - inp_0: R.Tensor((256, 256), dtype="float32") - ) -> R.Tensor((256, 256), dtype="float32"): - with R.dataflow(): - lv: R.Tensor((256, 256), dtype="float32") = R.rsqrt(inp_0) - gv: R.Tensor((256, 256), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Rsqrt(), [([256, 256], "float32")], {}, Expected1) - - def test_cat(): class Cat0(Module): def forward(self, x, y): @@ -3776,26 +3469,6 @@ def main( verify_model(Cat3(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1) -def test_neg(): - class Neg(Module): - def forward(self, input): - return -input - - @I.ir_module - class Expected1: - @R.function - def main( - inp_0: R.Tensor((256, 256), dtype="float32") - ) -> R.Tensor((256, 256), dtype="float32"): - with R.dataflow(): - lv: R.Tensor((256, 256), dtype="float32") = R.negative(inp_0) - gv: R.Tensor((256, 256), dtype="float32") = lv - R.output(gv) - return gv - - verify_model(Neg(), [([256, 256], "float32")], {}, Expected1) - - def test_max(): class Max(Module): def forward(self, x, y):