From ee48069c9870d87f43627c5aab4129e420bdba26 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 1/8] support `relu_.default` --- .../tvm/relax/frontend/torch/exported_program_translator.py | 1 + tests/python/relax/test_frontend_from_exported_program.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index cc9217c9f5f8..7fa39dc111fa 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -278,6 +278,7 @@ def create_convert_map( "neg.default": self._unary_op(relax.op.negative), "reciprocal.default": self._reciprocal, "relu.default": self._unary_op(relax.op.nn.relu), + "relu_.default": self._unary_op(relax.op.nn.relu), "round.default": self._round, "rsqrt.default": self._unary_op(relax.op.rsqrt), "selu.default": self._unary_op(relax.op.nn.selu), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 081b82b3c563..0c7a221f0f7e 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -511,6 +511,10 @@ class ReLU1(Module): def forward(self, input): return torch.nn.functional.relu(input) + class ReLU2(Module): + def forward(self, input): + return torch.ops.aten.relu_(input) + @tvm.script.ir_module class expected_relu: @R.function @@ -526,6 +530,7 @@ def main( verify_model(ReLU0(), example_args, {}, expected_relu) verify_model(ReLU1(), example_args, {}, expected_relu) + verify_model(ReLU2(), example_args, {}, expected_relu) # selu class Selu1(Module): From e472c49c900a2869f438fa3bf9ed7e8c751892b4 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 2/8] support `add_.Tensor` --- python/tvm/relax/frontend/torch/exported_program_translator.py | 1 + tests/python/relax/test_frontend_from_exported_program.py | 1 + 2 files changed, 2 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 7fa39dc111fa..bd336c895b67 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -297,6 +297,7 @@ def create_convert_map( "triu.default": self._tril_triu(relax.op.triu), # binary "add.Tensor": self._binary_op(relax.op.add, operator.add), + "add_.Tensor": self._binary_op(relax.op.add, operator.add), "div.Tensor": self._binary_op(relax.op.divide, operator.truediv), "eq.Scalar": self._binary_op(relax.op.equal, operator.eq), "eq.Tensor": self._binary_op(relax.op.equal, operator.eq), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 0c7a221f0f7e..0d3eaceed597 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -850,6 +850,7 @@ def main( operator_binary_1 = [ (operator.add, R.add), + (torch.ops.aten.add_, R.add), (operator.sub, R.subtract), (operator.mul, R.multiply), (operator.truediv, R.divide), From 57913842520d5c66854351f1c7198722d85bc9f9 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 3/8] support `silu_.default` --- .../tvm/relax/frontend/torch/exported_program_translator.py | 1 + tests/python/relax/test_frontend_from_exported_program.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index bd336c895b67..3053b5196b36 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -285,6 +285,7 @@ def create_convert_map( "sigmoid.default": self._unary_op(relax.op.sigmoid), "sign.default": self._unary_op(relax.op.sign), "silu.default": self._unary_op(relax.op.nn.silu), + "silu_.default": self._unary_op(relax.op.nn.silu), "sin.default": self._unary_op(relax.op.sin), "sinh.default": self._unary_op(relax.op.sinh), "softmax.int": self._softmax, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 0d3eaceed597..77031033e769 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -602,6 +602,10 @@ class SiLU2(Module): def forward(self, input): return torch.nn.functional.silu(input) + class SiLU3(Module): + def forward(self, input): + return torch.ops.aten.silu_(input) + @tvm.script.ir_module class expected_silu: @R.function @@ -617,6 +621,7 @@ def main( verify_model(SiLU(), example_args, {}, expected_silu) verify_model(SiLU2(), example_args, {}, expected_silu) + verify_model(SiLU3(), example_args, {}, expected_silu) # softmax test_softmax() From 9e1e194aa51b4db9229763063575bf406f4dbc7a Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 4/8] support `dropout_.default` --- .../tvm/relax/frontend/torch/exported_program_translator.py | 1 + tests/python/relax/test_frontend_from_exported_program.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 3053b5196b36..34bf2fe058b4 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -258,6 +258,7 @@ def create_convert_map( "cos.default": self._unary_op(relax.op.cos), "cosh.default": self._unary_op(relax.op.cosh), "dropout.default": lambda node: self.env[node.args[0]], + "dropout_.default": lambda node: self.env[node.args[0]], "elu.default": self._elu, "erf.default": self._unary_op(relax.op.erf), "exp.default": self._unary_op(relax.op.exp), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 77031033e769..02a311cc1402 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -254,6 +254,10 @@ class Dropout2(Module): def forward(self, input): return torch.dropout(input, 0.5, train=True) + class Dropout3(Module): + def forward(self, input): + return torch.ops.aten.dropout_(input, 0.5, train=True) + @tvm.script.ir_module class expected_dropout: @R.function @@ -268,6 +272,7 @@ def main( verify_model(Dropout1(), example_args, {}, expected_dropout) verify_model(Dropout2(), example_args, {}, expected_dropout) + verify_model(Dropout3(), example_args, {}, expected_dropout) # elu class Elu(Module): From c8010b802d5dd5b5375f0b0b3987cfd4f8dc15e8 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 5/8] support `hardswish_.default` --- .../tvm/relax/frontend/torch/exported_program_translator.py | 1 + tests/python/relax/test_frontend_from_exported_program.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 34bf2fe058b4..9f0ffccc38f3 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -266,6 +266,7 @@ def create_convert_map( "gelu.default": self._gelu, "hardsigmoid.default": self._hardsigmoid, "hardswish.default": self._hardswish, + "hardswish_.default": self._hardswish, "hardtanh.default": self._hardtanh, "isfinite.default": self._unary_op(relax.op.isfinite), "isinf.default": self._unary_op(relax.op.isinf), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 02a311cc1402..8b5ed12fa501 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -388,6 +388,10 @@ class Hardswish2(torch.nn.Module): def forward(self, input): return torch.nn.functional.hardswish(input) + class Hardswish3(torch.nn.Module): + def forward(self, input): + return torch.ops.aten.hardswish_(input) + @tvm.script.ir_module class expected1: @R.function @@ -407,6 +411,7 @@ def main( verify_model(Hardswish(), example_args, {}, expected1) verify_model(Hardswish2(), example_args, {}, expected1) + verify_model(Hardswish3(), example_args, {}, expected1) # hardtanh test_hardtanh() From edb01b0a88744f873e8823b0ffdf72e478637a34 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 6/8] support `hardtanh_.default` --- .../tvm/relax/frontend/torch/exported_program_translator.py | 5 +++-- tests/python/relax/test_frontend_from_exported_program.py | 5 +++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 9f0ffccc38f3..4c09a242c835 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -39,8 +39,8 @@ class ExportedProgramImporter(BaseFXGraphImporter): def _hardtanh(self, node: fx.Node) -> relax.Expr: args = self.retrieve_args(node) x = args[0] - min_val = node.args[1] if len(args) > 1 else node.kwargs("min_val", -1.0) - max_val = node.args[2] if len(args) > 2 else node.kwargs("max_val", 1.0) + min_val = node.args[1] if len(args) > 1 else node.kwargs.get("min_val", -1.0) + max_val = node.args[2] if len(args) > 2 else node.kwargs.get("max_val", 1.0) return self.block_builder.emit(relax.op.clip(x, min_val, max_val)) def _log2(self, node: fx.Node) -> relax.Var: @@ -268,6 +268,7 @@ def create_convert_map( "hardswish.default": self._hardswish, "hardswish_.default": self._hardswish, "hardtanh.default": self._hardtanh, + "hardtanh_.default": self._hardtanh, "isfinite.default": self._unary_op(relax.op.isfinite), "isinf.default": self._unary_op(relax.op.isinf), "isnan.default": self._unary_op(relax.op.isnan), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 8b5ed12fa501..4c5a08cb4eb2 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -656,6 +656,10 @@ class Hardtanh2(torch.nn.Module): def forward(self, input): return torch.nn.functional.hardtanh(input) + class Hardtanh3(torch.nn.Module): + def forward(self, input): + return torch.ops.aten.hardtanh_(input) + @tvm.script.ir_module class expected1: @R.function @@ -673,6 +677,7 @@ def main( example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),) verify_model(Hardtanh(), example_args, {}, expected1) verify_model(Hardtanh2(), example_args, {}, expected1) + verify_model(Hardtanh3(), example_args, {}, expected1) def test_leakyrelu(): From 57c0a4fafc9c07fd69def7f59222454ac3cc5d59 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 7/8] support `unflatten.int` --- .../torch/exported_program_translator.py | 14 ++++++++++ .../test_frontend_from_exported_program.py | 27 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 4c09a242c835..70dde0de45c6 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -216,6 +216,19 @@ def _slice(self, node: fx.Node) -> relax.Var: stride = [node.args[4] if len(node.args) > 4 else 1] return self.block_builder.emit(relax.op.strided_slice(x, axes, begin, end, stride)) + def _unflatten(self, node: fx.Node) -> relax.Var: + args = self.retrieve_args(node) + x = args[0] + dim = node.args[1] + sizes = node.args[2] + + x_shape = list(self.shape_of(x)) + if dim < 0: + dim += len(x_shape) + + new_shape = x_shape[:dim] + sizes + x_shape[dim + 1:] + return self.block_builder.emit(relax.op.reshape(x, new_shape)) + ########## Creation ########## def _one_hot(self, node: fx.Node) -> relax.Var: @@ -399,6 +412,7 @@ def create_convert_map( "tile.default": self._tile, "topk.default": self._topk, "transpose.int": self._transpose, + "unflatten.int": self._unflatten, "unsqueeze.default": lambda node: self.block_builder.emit( relax.op.expand_dims(self.env[node.args[0]], node.args[1]) ), diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 4c5a08cb4eb2..dd4ead9e593e 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -3629,6 +3629,33 @@ def main( verify_model(Select(), example_args, {}, Expected) +def test_unflatten(): + class Unflatten(Module): + def forward(self, input): + return torch.ops.aten.unflatten(input, 1, (3, 5)) + + class Unflatten1(Module): + def forward(self, input): + return torch.ops.aten.unflatten(input, -2, (3, 5)) + + @tvm.script.ir_module + class Expected: + @R.function + def main( + inp_0: R.Tensor((2, 15, 7), dtype="float32"), + ) -> R.Tuple(R.Tensor((2, 3, 5, 7), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((2, 3, 5, 7), dtype="float32") = R.reshape(inp_0, [2, 3, 5, 7]) + gv: R.Tuple(R.Tensor((2, 3, 5, 7), dtype="float32")) = (lv,) + R.output(gv) + return gv + + example_args = (torch.randn(2, 15, 7, dtype=torch.float32),) + + verify_model(Unflatten(), example_args, {}, Expected) + verify_model(Unflatten1(), example_args, {}, Expected) + + def test_gather(): class Gather0(Module): def forward(self, data, indices): From ccce4ebdfe06605466e1ebf493518507ad7d8eb9 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Mon, 7 Apr 2025 16:35:07 +0900 Subject: [PATCH 8/8] fix lint error --- python/tvm/relax/frontend/torch/exported_program_translator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 70dde0de45c6..c05858fd887e 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -226,7 +226,7 @@ def _unflatten(self, node: fx.Node) -> relax.Var: if dim < 0: dim += len(x_shape) - new_shape = x_shape[:dim] + sizes + x_shape[dim + 1:] + new_shape = x_shape[:dim] + sizes + x_shape[dim + 1 :] return self.block_builder.emit(relax.op.reshape(x, new_shape)) ########## Creation ##########