From f27840e744c1ffacbd580d8d376e0b3e77fcdd6e Mon Sep 17 00:00:00 2001 From: Shushi Hong <820958424@qq.com> Date: Wed, 2 Apr 2025 18:36:09 +0800 Subject: [PATCH 1/3] Update exported_program_translator.py --- .../relax/frontend/torch/exported_program_translator.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 97ccc6393cbb..62e98b88ed35 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -60,6 +60,10 @@ def _log1p(self, node: fx.Node) -> relax.Var: one = relax.const(1, x.struct_info.dtype) return self.block_builder.emit(relax.op.log(relax.op.add(x, one))) + def _reciprocal(self, node: fx.Node) -> relax.Var: + x = self.env[node.args[0]] + return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.struct_info.dtype), x)) + ########## Neural Network ########## def _batch_norm(self, node: fx.Node, training) -> relax.Var: @@ -272,6 +276,7 @@ def create_convert_map( "log1p.default": self._log1p, "log_softmax.int": self._log_softmax, "neg.default": self._unary_op(relax.op.negative), + "reciprocal.default": self._reciprocal, "relu.default": self._unary_op(relax.op.nn.relu), "round.default": self._round, "rsqrt.default": self._unary_op(relax.op.rsqrt), @@ -361,6 +366,7 @@ def create_convert_map( # search "argmax.default": self._argmax_argmin(relax.op.argmax), "argmin.default": self._argmax_argmin(relax.op.argmin), + "where.self": self._where, # tensor manipulation "cat.default": self._cat, "chunk.default": self._chunk, @@ -368,6 +374,7 @@ def create_convert_map( "concat.default": self._cat, "copy_.default": self._copy_, "cumsum.default": self._cumsum, + "cumprod.default": self._cumprod, "expand.default": self._expand, "expand_as.default": self._expand_as, "flip.default": self._flip, From 4ff0ea8ba60cc028abf955a4e655a74111a666f5 Mon Sep 17 00:00:00 2001 From: Shushi Hong <820958424@qq.com> Date: Wed, 2 Apr 2025 18:39:17 +0800 Subject: [PATCH 2/3] Update test_frontend_from_exported_program.py --- .../test_frontend_from_exported_program.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 98f0f1d9cac6..3112b20a72d4 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -477,6 +477,25 @@ def main( # log_softmax test_logsoftmax() + # reciprocal + class Reciprocal(Module): + def forward(self, input): + return torch.reciprocal(input) + + @tvm.script.ir_module + class expected_reciprocal: + @R.function + def main( + input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(R.const(1.0, "float32"), input_1) + gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,) + R.output(gv) + return gv + + verify_model(Reciprocal(), example_args, {}, expected_reciprocal) + # relu class ReLU0(Module): def __init__(self): @@ -3818,5 +3837,52 @@ def main( verify_model(Prod(), example_args, {}, Expected) +def test_cumprod(): + class Cumprod(Module): + def forward(self, x): + return torch.cumprod(x, 0) + + @tvm.script.ir_module + class Expected: + @R.function + def main( + inp_0: R.Tensor((5, 3), dtype="float32"), + ) -> R.Tuple(R.Tensor((5, 3), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((5, 3), dtype="float32") = R.cumprod(inp_0, axis=0, exclusive=False) + gv: R.Tuple(R.Tensor((5, 3), dtype="float32")) = (lv,) + R.output(gv) + return gv + + example_input = torch.randn(5, 3, dtype=torch.float32) + verify_model(Cumprod(), (example_input,), {}, Expected) + + +def test_where(): + class Where(Module): + def forward(self, condition, x, y): + return torch.where(condition, x, y) + + @tvm.script.ir_module + class Expected: + @R.function + def main( + inp_0: R.Tensor((5, 3), dtype="bool"), + inp_1: R.Tensor((5, 3), dtype="float32"), + inp_2: R.Tensor((5, 3), dtype="float32"), + ) -> R.Tuple(R.Tensor((5, 3), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((5, 3), dtype="float32") = R.where(inp_0, inp_1, inp_2) + gv: R.Tuple(R.Tensor((5, 3), dtype="float32")) = (lv,) + R.output(gv) + return gv + + condition = torch.randint(0, 2, (5, 3), dtype=torch.bool) + x = torch.randn(5, 3, dtype=torch.float32) + y = torch.randn(5, 3, dtype=torch.float32) + + verify_model(Where(), (condition, x, y), {}, Expected) + + if __name__ == "__main__": tvm.testing.main() From fbacfe31c930132de9584d560b34e723ac094677 Mon Sep 17 00:00:00 2001 From: Shushi Hong <820958424@qq.com> Date: Wed, 2 Apr 2025 22:19:24 +0800 Subject: [PATCH 3/3] Update test_frontend_from_exported_program.py --- tests/python/relax/test_frontend_from_exported_program.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 3112b20a72d4..e37ee0e40480 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -486,10 +486,12 @@ def forward(self, input): class expected_reciprocal: @R.function def main( - input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + input_1: R.Tensor((1, 3, 10, 10), dtype="float32") ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")): with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(R.const(1.0, "float32"), input_1) + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide( + R.const(1.0, "float32"), input_1 + ) gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,) R.output(gv) return gv