llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/projects/pt1/examples/torchscript_stablehlo_backe...

17 lines
525 B
Python
Raw Normal View History

mhlo: migrate conversion to stablehlo (#1840) This patch replaces all MHLO operations with their StableHLO counterparts and adds a validation pass to ensure that no MHLO operations remain before translating all Stablehlo operations to the MHLO dialect for further lowering to the Linalg dialect. This patch also updates all lit tests so that they refer to the `convert-torch-to-stablehlo` pass and so that they check for StableHLO operations.
2023-02-02 21:29:47 +08:00
import torch
import torchvision.models as models
Rename torch_mlir.compile APIs and introduce FX based analogs (#2842) Link to related RFC: https://discourse.llvm.org/t/rfc-rename-torch-mlir-compile-apis-and-introduce-fx-based-analogs/76646 This commit updates the documentation, tests, CMake files, and API for the proposed changes in the RFC. There is a new torch_mlir/fx.py for user level APIs related to importing modules and a corresponding test for this path can be found at test/python/fx_importer/basic_test.py. --------- Co-authored-by: MaheshRavishankar <mravisha@amd.com>
2024-02-07 11:07:59 +08:00
from torch_mlir import torchscript
mhlo: migrate conversion to stablehlo (#1840) This patch replaces all MHLO operations with their StableHLO counterparts and adds a validation pass to ensure that no MHLO operations remain before translating all Stablehlo operations to the MHLO dialect for further lowering to the Linalg dialect. This patch also updates all lit tests so that they refer to the `convert-torch-to-stablehlo` pass and so that they check for StableHLO operations.
2023-02-02 21:29:47 +08:00
model = models.resnet18(pretrained=True)
model.eval()
[NFC reformat] Applies pre-commit formatting to Python files. (#3244) This is a large change because prior to this point, Python files in the project were not consistently formatted. This reformats them all with black defaults. Based on experience with prior projects, if you have a dev/long-term branch with Python patches, you can minimize merge conflicts prior to rebasing to include this commit by running `black` on your modified Python files, squashing, and then rebasing/merging.
2024-04-28 05:16:31 +08:00
data = torch.randn(2, 3, 200, 200)
mhlo: migrate conversion to stablehlo (#1840) This patch replaces all MHLO operations with their StableHLO counterparts and adds a validation pass to ensure that no MHLO operations remain before translating all Stablehlo operations to the MHLO dialect for further lowering to the Linalg dialect. This patch also updates all lit tests so that they refer to the `convert-torch-to-stablehlo` pass and so that they check for StableHLO operations.
2023-02-02 21:29:47 +08:00
out_stablehlo_mlir_path = "./resnet18_stablehlo.mlir"
[NFC reformat] Applies pre-commit formatting to Python files. (#3244) This is a large change because prior to this point, Python files in the project were not consistently formatted. This reformats them all with black defaults. Based on experience with prior projects, if you have a dev/long-term branch with Python patches, you can minimize merge conflicts prior to rebasing to include this commit by running `black` on your modified Python files, squashing, and then rebasing/merging.
2024-04-28 05:16:31 +08:00
module = torchscript.compile(
model, data, output_type=torchscript.OutputType.STABLEHLO, use_tracing=False
)
mhlo: migrate conversion to stablehlo (#1840) This patch replaces all MHLO operations with their StableHLO counterparts and adds a validation pass to ensure that no MHLO operations remain before translating all Stablehlo operations to the MHLO dialect for further lowering to the Linalg dialect. This patch also updates all lit tests so that they refer to the `convert-torch-to-stablehlo` pass and so that they check for StableHLO operations.
2023-02-02 21:29:47 +08:00
with open(out_stablehlo_mlir_path, "w", encoding="utf-8") as outf:
outf.write(str(module))
print(f"StableHLO IR of resent18 successfully written into {out_stablehlo_mlir_path}")