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

33 lines
1.3 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
build: improve robustness of cmake and shell scripts (#1018) On my local machine, `unzip` didn't exist (producing a "command not found" error), but CMake ignored the error. Although the build did succeed (because it found a previously-built version of libtorch), it seems better to abort builds on such failures, so this patch checks the return code of all external process invocations. Along similar lines, this patch also updates the shell scripts in `build_tools` to extensively use double-quoting to prevent unintentional word splitting or globbing. Since some of the scripts execute `rm` while using shell variables, this patch also adds the preamble `set -u` to abort execution if an undefined variable is referenced, so that we reduce the chances of executing `rm -rf /` if the path expression happens to refer to an undefined variable.
2022-07-07 05:39:30 +08:00
set -euo 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
build: improve robustness of cmake and shell scripts (#1018) On my local machine, `unzip` didn't exist (producing a "command not found" error), but CMake ignored the error. Although the build did succeed (because it found a previously-built version of libtorch), it seems better to abort builds on such failures, so this patch checks the return code of all external process invocations. Along similar lines, this patch also updates the shell scripts in `build_tools` to extensively use double-quoting to prevent unintentional word splitting or globbing. Since some of the scripts execute `rm` while using shell variables, this patch also adds the preamble `set -u` to abort execution if an undefined variable is referenced, so that we reduce the chances of executing `rm -rf /` if the path expression happens to refer to an undefined variable.
2022-07-07 05:39:30 +08:00
src_dir="$(realpath "$(dirname "$0")"/..)"
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
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"
Re-enable custom op support
2022-08-03 04:35:56 +08:00
TORCH_MLIR_EXT_PYTHONPATH="${TORCH_MLIR_EXT_PYTHONPATH:-""}"
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"
Re-enable custom op support
2022-08-03 04:35:56 +08:00
if [ ! -z ${TORCH_MLIR_EXT_PYTHONPATH} ]; then
pypath="${pypath}:${TORCH_MLIR_EXT_PYTHONPATH}"
fi
TORCH_MLIR_EXT_MODULES="${TORCH_MLIR_EXT_MODULES:-""}"
if [ ! -z ${TORCH_MLIR_EXT_MODULES} ]; then
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 \
Bump the shape lib to match the upstream functions currently in PyTorch (#1236) Bumps the shape library: - Updates the function signature for aten.arange.start_step - upstream_shape_functions.mean_dim -> upstream_shape_functions.sum_mean_dim
2022-08-17 12:11:04 +08:00
--pytorch_op_extensions=${ext_module:-""} \
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
--torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"