a83e106f92
This is a first step towards reworking the scalarize-shapes pass which has been integral to our ONNX frontend path detangling shape computations. ## Purpose: 1. Restrict the scope of the pass to only apply to op sequences which are used to compute shapes. 2. Make the pass more efficient by applying patterns in an appropriate order for scalarization propagation. 3. Report failed scalarization patterns for easier debugging (Not yet implemented). I can't seem to find a good path for this right now to capture the right diagnostics. I'd like to defer this addition to a later patch so we can add some high-value patterns to this pass in the meantime. With these changes, some reworking of the conversions themselves will be necessary. 1. The removal of the SqueezeDim fold pattern was an appropriate fix to avoid folding a pattern that may be needed to propagate further. The reversal of pattern application order uncovered this bug. The addition of rank 0 item logic was added to replace the functionality needed from the squeeze dim pattern. 2. Rework getListFromTensor to modify a `SmallVector<OpFoldResult>` to allow processing value tensor literals without immediately materializing the ints. This should factor out a significant portion of code that was used in specific cases to handle constants. ## RFC 1: Currently, we are going to add all prim list of int ops to the worklist. Can anyone identify problems with uniformly anchoring on prim lists of ints? E.g. Does there exist a Torch Op satisfying all of the following conditions: 1. Accepts a list of constant ints, LIST, as an input 2. The role of LIST is **not** shape related. All the examples I can think of are indeed shape related: padding ints passed to a pad op, kernel size ints passed to a conv op, size ints passed to a view op, etc. 4. The LIST is not gotten entirely from scalars already. If there does not exist a torch op satisfying all three of those conditions, I think it will be safe to "anchor" on prim lists of ints. ### Conclusion for RFC 1: I just scanned through the `GeneratedTorchOps.td` and `TorchOps.td` for all references of `AnyTorchListOfTorchIntType` and verified this will not be problematic to apply in any of those cases. ## RFC 2: What should I use to report failed scalarization? Like my dumb idea was just to walk back through the func op after applying the passes and check if anything in the worklist is still a tensor. If so, emit/log a warning. It certainly works, since you can just look at the warnings and start debugging from the last printed warning upwards, but there has to be a better way to handle this without walking back through the func.func op. ### Conclusion for RFC 2: I tried a few things without much success. The fundamental problem is that identifying the cause of a failed scalarization could be myriad: 1. We could be missing a pattern for an op entirely: E.g., a pattern we need is scalarizing rank0 arithmetic ops (e.g. AtenMulTensorOp -> AtenMulIntOp). 2. We could fail a scalarization pattern because it should fold instead. This is specifically the case for rank0 where.self ops. These ops MUST fold, or we need to have custom lowering logic for the rank 0 case. 3. Walking through the func op a second time and emiting a warning for ops that have tensor result types seems to give locations that are inconsistent or hard to track in the converted IR. Doing this on IR that doesn't apply any patterns seems to give decent information, but it's still dramatically insufficient considering how complex these patterns can get, and still takes manually reading IR to try and figure out what is really blocking the simplification. I'd like to skip out on fleshing out the error reporting for now and come back to it after iterating a few time on the patterns. |
||
|---|---|---|
| .github | ||
| build_tools | ||
| docs | ||
| externals | ||
| include | ||
| lib | ||
| projects | ||
| python | ||
| test | ||
| tools | ||
| utils/bazel | ||
| .clang-format | ||
| .git-blame-ignore-revs | ||
| .gitignore | ||
| .gitmodules | ||
| .pre-commit-config.yaml | ||
| .yamllint.yml | ||
| CITATION.cff | ||
| CMakeLists.txt | ||
| LICENSE | ||
| README.md | ||
| build-requirements.txt | ||
| pyproject.toml | ||
| pytorch-hash.txt | ||
| pytorch-requirements.txt | ||
| requirements.txt | ||
| setup.py | ||
| test-requirements.txt | ||
| torchvision-requirements.txt | ||
| whl-requirements.txt | ||
README.md
The Torch-MLIR Project
The Torch-MLIR project aims to provide first class compiler support from the PyTorch ecosystem to the MLIR ecosystem.
This project is participating in the LLVM Incubator process: as such, it is not part of any official LLVM release. While incubation status is not necessarily a reflection of the completeness or stability of the code, it does indicate that the project is not yet endorsed as a component of LLVM.
PyTorch PyTorch is an open source machine learning framework that facilitates the seamless transition from research and prototyping to production-level deployment.
MLIR The MLIR project offers a novel approach for building extensible and reusable compiler architectures, which address the issue of software fragmentation, reduce the cost of developing domain-specific compilers, improve compilation for heterogeneous hardware, and promote compatibility between existing compilers.
Torch-MLIR Several vendors have adopted MLIR as the middle layer in their systems, enabling them to map frameworks such as PyTorch, JAX, and TensorFlow into MLIR and subsequently lower them to their target hardware. We have observed half a dozen custom lowerings from PyTorch to MLIR, making it easier for hardware vendors to focus on their unique value, rather than needing to implement yet another PyTorch frontend for MLIR. The ultimate aim is to be similar to the current hardware vendors adding LLVM target support, rather than each one implementing Clang or a C++ frontend.
All the roads from PyTorch to Torch MLIR Dialect
We have few paths to lower down to the Torch MLIR Dialect.
- ONNX as the entry points.
- Fx as the entry points
Project Communication
#torch-mlirchannel on the LLVM Discord - this is the most active communication channel- Github issues here
torch-mlirsection of LLVM Discourse
Install torch-mlir snapshot
At the time of writing, we release pre-built snapshots of torch-mlir for Python 3.11 and Python 3.10.
If you have supported Python version, the following commands initialize a virtual environment.
python3.11 -m venv mlir_venv
source mlir_venv/bin/activate
Or, if you want to switch over multiple versions of Python using conda, you can create a conda environment with Python 3.11.
conda create -n torch-mlir python=3.11
conda activate torch-mlir
python -m pip install --upgrade pip
Then, we can install torch-mlir with the corresponding torch and torchvision nightlies.
pip install --pre torch-mlir torchvision \
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
Using torch-mlir
Torch-MLIR is primarily a project that is integrated into compilers to bridge them to PyTorch and ONNX. If contemplating a new integration, it may be helpful to refer to existing downstreams:
While most of the project is exercised via testing paths, there are some ways that an end user can directly use the APIs without further integration:
FxImporter ResNet18
# Get the latest example if you haven't checked out the code
wget https://raw.githubusercontent.com/llvm/torch-mlir/main/projects/pt1/examples/fximporter_resnet18.py
# Run ResNet18 as a standalone script.
python projects/pt1/examples/fximporter_resnet18.py
# Output
load image from https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg
...
PyTorch prediction
[('Labrador retriever', 70.65674591064453), ('golden retriever', 4.988346099853516), ('Saluki, gazelle hound', 4.477451324462891)]
torch-mlir prediction
[('Labrador retriever', 70.6567153930664), ('golden retriever', 4.988325119018555), ('Saluki, gazelle hound', 4.477458477020264)]
Repository Layout
The project follows the conventions of typical MLIR-based projects:
include/torch-mlir,libstructure for C++ MLIR compiler dialects/passes.testfor holding test code.toolsfortorch-mlir-optand such.pythontop level directory for Python code
Developers
If you would like to develop and build torch-mlir from source please look at Development Notes