Skip to content

Commit

Permalink
Run torch backend pipeline
Browse files Browse the repository at this point in the history
  • Loading branch information
penguin-wwy committed Apr 23, 2024
1 parent 7fdd1e6 commit eb7b1a9
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion python/torch_mlir/fx.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,27 @@
)


def _simplify_lowering(verbose, output_type, torch_mod):
def _simplify_lowering(verbose, output_type, torch_mod, backend_legal_ops=None,
extra_library_file_name=None):
# TODO: pass backend_legal_ops/extra_library_file_name by caller
if backend_legal_ops is None:
backend_legal_ops = []
if extra_library_file_name is None:
extra_library_file_name = ""
option_string = ("{backend-legal-ops=" + ",".join(backend_legal_ops) +
" extra-library=" + extra_library_file_name + "}")
run_pipeline_with_repro_report(
torch_mod,
f"builtin.module(torch-simplification-pipeline)",
"Simplification pipeline for torch dialect",
enable_ir_printing=verbose,
)
run_pipeline_with_repro_report(
torch_mod,
f"builtin.module(torch-function-to-torch-backend-pipeline{option_string})",
"Lowering TorchFX IR -> Torch Backend IR",
enable_ir_printing=verbose,
)
return lower_mlir_module(verbose, output_type, torch_mod)


Expand Down

0 comments on commit eb7b1a9

Please sign in to comment.