llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/build_tools/update_shape_lib.sh

35 lines
1.4 KiB
Bash
Raw Normal View History

Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
#!/bin/bash
# Updates auto-generated shape library files for the `torch` dialect.
Allow torch-mlir to support PyTorch extensions. (#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.
2022-06-14 05:51:30 +08:00
#
# 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
Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
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"
Allow torch-mlir to support PyTorch extensions. (#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.
2022-06-14 05:51:30 +08:00
pypath="${python_packages_dir}/torch_mlir"
Temporarily revert support for custom op extensions. (#944) The MacOS builders are having linking trouble with the extension library. Until it's fixed, all support for op extensions is disabled. It should be easy to restore once the issue is resolved.
2022-06-15 09:24:40 +08:00
# 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
Allow torch-mlir to support PyTorch extensions. (#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.
2022-06-14 05:51:30 +08:00
PYTHONPATH="${pypath}" python \
Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
-m torch_mlir.dialects.torch.importer.jit_ir.build_tools.shape_lib_gen \
--torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
Temporarily revert support for custom op extensions. (#944) The MacOS builders are having linking trouble with the extension library. Until it's fixed, all support for op extensions is disabled. It should be easy to restore once the issue is resolved.
2022-06-15 09:24:40 +08:00
# TODO: Add back to shape_lib_gen invocation once custom op support is back.
# --pytorch_op_extensions=${ext_module} \