[ingress][pytorch] Basic KernelBench to MLIR conversion #5
Conversation
Basic as can be torch-mlir converter for the level1 and level2 KernelBench kernels. The `convert-kernel-bench-to-mlir.py` script does the conversion and dumps the results in the `cache/level1` and `cache/level2` folders. Relies on pre-packaged mlir wheels and mlir-torch, as this PR considers dealing with versioning and packaging an orthogonal matter to getting ingress up and running. About ~55 of the 200 kernels are filtered out as they either crash torch-mlir or yield very big .mlir files. This ignore_list is meant to be amended as these issues get addressed, e.g. by altering init_inputs on a per kernel basis. The conversion script sticks to outputting just linalg for now. As it does this, it does do some basic post-processing of torch-mlir's output, namely it runs the -linalg-specialize-generic-ops pass.
| from mlir import ir, passmanager | ||
| from torch_mlir import fx | ||
|
|
||
| kernels_as_pytorch_folder = Path(__file__).parent / "KernelBench" / "KernelBench" |
There was a problem hiding this comment.
since this depends on where the git was cloned in the bash script, perhaps that last step (clone) could be done in this script as well?
There was a problem hiding this comment.
I am not sure.
Doing a git clone in either script feels unclean. I also don't like the idea of it being a submodule as that then seems to imply you have to clone KernelBench to do anything useful with lighthouse. It seems to me KernelBench will be just one source of ingress compute graphs of interest, with it potentially making sense to allow users/CI to opt-in to which paths they want to run tests with. What's the right mechanism for that? I am not sure.
There was a problem hiding this comment.
KernelBench is NOT an ingress. Torch-MLIR is.
We now have three PRs that work with FX importer, none using the other. We should have one FX importer script that is used by others.
| Path(kernel_pytorch_file.parent.name) / kernel_pytorch_file.name | ||
| ) | ||
| if level_and_kernel in ignore_list or not kernel_pytorch_file.is_file(): | ||
| print( |
| if not all( | ||
| hasattr(module, a) for a in ("Model", "get_inputs", "get_init_inputs") | ||
| ): | ||
| print(f"Error: module in file {kernel_pytorch_file} not a proper benchmark") |
There was a problem hiding this comment.
do we want to mark error so to return non-zero at the end upon any such continue?
| print(f"Error: got the following error converting {kernel_pytorch_file}") | ||
| raise e | ||
|
|
||
| before_clean_up = "//" + str(m)[:-1].replace("\n", "\n//") + "\n" |
There was a problem hiding this comment.
This is to include the pre-cleanup output of torch-mlir in the final file, though as commented-out code.
It helped me see what the specialization pass (& potentially other clean-up) is accomplishing.
I don't have strong opinion on whether this should stay.
2 participants
Basic as can be torch-mlir converter for the level1 and level2 KernelBench kernels. The
convert-kernel-bench-to-mlir.pyscript does the conversion and dumps the results in thecache/level1andcache/level2folders.Relies on pre-packaged mlir wheels and mlir-torch, as this PR considers dealing with versioning and packaging an orthogonal matter to getting ingress up and running.
About ~55 of the 200 kernels are filtered out as they either crash torch-mlir or yield very big .mlir files. This ignore_list is meant to be amended as these issues get addressed, e.g. by altering init_inputs on a per kernel basis.
The conversion script sticks to outputting just linalg for now. As it does this, it does do some basic post-processing of torch-mlir's output, namely it runs the -linalg-specialize-generic-ops pass.