forked from llvm/torch-mlir
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Allow torch-mlir to support PyTorch extensions. (llvm#895)
PyTorch allows new operators to be registered dynamically in modules. Torch-mlir already makes it fairly straightforward to add support for new operators, and this commit just extends that support to allow new PyTorch ops to come from a external module. This does *not* allow ops to be dynamically loaded into torch-mlir. Torch-mlir must still be compiled with support built-in. Add a `_torch_mlir_custom_op_example` subpackage to `torch_mlir` which registers an demonstration op. It will not be imported by default when importing torch_mlir. It's strictly for testing and documentation. Adds an end-to-end test for the `torch_mlir_custom_op_example::identity` op. With all these changes, we should now be actively testing PyTorch extension support with all future patches.
- Loading branch information
Showing
34 changed files
with
923 additions
and
65 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,13 +1,34 @@ | ||
#!/bin/bash | ||
# Updates auto-generated shape library files for the `torch` dialect. | ||
set -e | ||
# | ||
# Environment variables: | ||
# TORCH_MLIR_EXT_MODULES: comma-separated list of python module names | ||
# which register custom PyTorch operators upon being imported. | ||
# TORCH_MLIR_EXT_PYTHONPATH: colon-separated list of paths necessary | ||
# for importing PyTorch extensions specified in TORCH_MLIR_EXT_MODULES. | ||
# For more information on supporting custom operators, see: | ||
# ${TORCH_MLIR}/python/torch_mlir/_torch_mlir_custom_op_example/README.md | ||
|
||
set -eo pipefail | ||
|
||
src_dir="$(realpath $(dirname $0)/..)" | ||
build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")" | ||
torch_transforms_cpp_dir="${src_dir}/lib/Dialect/Torch/Transforms" | ||
python_packages_dir="${build_dir}/tools/torch-mlir/python_packages" | ||
|
||
#ninja -C "${build_dir}" | ||
PYTHONPATH="${python_packages_dir}/torch_mlir" python \ | ||
pypath="${python_packages_dir}/torch_mlir" | ||
# TODO: Re-enable once custom op support is back. | ||
#if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then | ||
# pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" | ||
#fi | ||
#ext_module="torch_mlir._torch_mlir_custom_op_example" | ||
#if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then | ||
# ext_module="${ext_module},${TORCH_MLIR_EXT_MODULES} " | ||
#fi | ||
|
||
PYTHONPATH="${pypath}" python \ | ||
-m torch_mlir.dialects.torch.importer.jit_ir.build_tools.shape_lib_gen \ | ||
--torch_transforms_cpp_dir="${torch_transforms_cpp_dir}" | ||
|
||
# TODO: Add back to shape_lib_gen invocation once custom op support is back. | ||
# --pytorch_op_extensions=${ext_module} \ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,14 +1,35 @@ | ||
#!/bin/bash | ||
# Updates auto-generated ODS files for the `torch` dialect. | ||
set -euo pipefail | ||
# | ||
# Environment variables: | ||
# TORCH_MLIR_EXT_MODULES: comma-separated list of python module names | ||
# which register custom PyTorch operators upon being imported. | ||
# TORCH_MLIR_EXT_PYTHONPATH: colon-separated list of paths necessary | ||
# for importing PyTorch extensions specified in TORCH_MLIR_EXT_MODULES. | ||
# For more information on supporting custom operators, see: | ||
# ${TORCH_MLIR}/python/torch_mlir/_torch_mlir_custom_op_example/README.md | ||
|
||
set -eo pipefail | ||
|
||
src_dir="$(realpath $(dirname $0)/..)" | ||
build_dir="$(realpath "${TORCH_MLIR_BUILD_DIR:-$src_dir/build}")" | ||
torch_ir_include_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR" | ||
python_packages_dir="${build_dir}/tools/torch-mlir/python_packages" | ||
|
||
#ninja -C "${build_dir}" | ||
PYTHONPATH="${python_packages_dir}/torch_mlir" python \ | ||
pypath="${python_packages_dir}/torch_mlir" | ||
# TODO: Re-enable once custom op support is back. | ||
#if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then | ||
# pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}" | ||
#fi | ||
#ext_module="torch_mlir._torch_mlir_custom_op_example" | ||
#if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then | ||
# ext_module="${ext_module},${TORCH_MLIR_EXT_MODULES}" | ||
#fi | ||
|
||
PYTHONPATH="${pypath}" python \ | ||
-m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen \ | ||
--torch_ir_include_dir="${torch_ir_include_dir}" \ | ||
--debug_registry_dump="${torch_ir_include_dir}/JITOperatorRegistryDump.txt" | ||
|
||
# TODO: Add back to torch_ods_gen invocation once custom op support is back. | ||
# --pytorch_op_extensions="${ext_module}" \ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.