[ingress][pytorch] Basic KernelBench to MLIR conversion

[rolfmorel]

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 alongside the script.

56 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.

[rengolin]

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?

[rolfmorel]

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.

[rengolin]

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.

[rolfmorel]

The importer impasse has been resolved.

Whether the KernelBench submodule and converter script should live in this "ingress" directory is up to taste. I will defer to anyone who suggests a better path.

[rengolin]

do we want to mark error so to return non-zero at the end upon any such continue?

[rolfmorel]

I believe the uncaptured exception raised by the following line will terminate the whole script and exit with non-zero.

If we prefer to perform a more graceful exit, let me know.

My take is that an exception being raised here is truly unexpected and hence would provide us valuable info in case a user is able to include it in their report.

[rolfmorel]

I thought to leave the following here:

$ time uv run convert-kernel-bench-to-mlir.py                                                                    
Processing: level1/100_HingeLoss.py                                                                                                                         
Processing: level1/10_3D_tensor_matrix_multiplication.py                                                                                                    
Processing: level1/11_4D_tensor_matrix_multiplication.py                                                                                                    
Skipping: level1/12_Matmul_with_diagonal_matrices_.py
...
Processing: level2/96_ConvTranspose3d_Multiply_Max_GlobalAvgPool_Clamp.py
Skipping: level2/97_Matmul_BatchNorm_BiasAdd_Divide_Swish.py
Skipping: level2/98_Matmul_AvgPool_GELU_Scale_Max.py
Skipping: level2/99_Matmul_GELU_Softmax.py
Skipping: level2/9_Matmul_Subtract_Multiply_ReLU.py

real    6m15.501s
user    5m29.552s
sys     1m24.632s
$ ls -l cache/* | grep .mlir | wc -l
144

That is, even with the worst offenders filtered out, using vanilla torch-mlir to convert these 144 simple NNs is still terribly slow. I expect this is in no small part due to the huge dialect_resources: { builtin: { torch_tensor_...float...: "0x040... binary blobs that get tacked onto the IR. We need to find a way to get torch-mlir to do a more sensible thing for us.

EDIT: For completeness, here's how long it takes when not doing the deserialization and clean-up pass afterwards, i.e. completely vanilla torch_mlir.fx.export_and_import and dumping the IR straight to file:

real    5m50.468s
user    5m21.148s
sys     1m0.587s

[adam-smnk]

That is, even with the worst offenders filtered out, using vanilla torch-mlir to convert these 144 simple NNs is still terribly slow. I expect this is in no small part due to the huge dialect_resources: { builtin: { torch_tensor_...float...: "0x040... binary blobs that get tacked onto the IR. We need to find a way to get torch-mlir to do a more sensible thing for us.

Is there any option for torch-mlir to lower these constant into input args? That is prevent dumping constants at all.

[adam-smnk]

nit: It probably needs update --init when the dir is not present at all

[rolfmorel]

Tried it yesterday without any directory there. That is, started out clean, I ran the script and got the error whereupon I copied the (--init-less) command and ran that. After that I could run the script without error.

Might depend on git version though. If someone knows or encounters that this command isn't sufficient in all cases, please let me know and I will amend.

[adam-smnk]

I looks like the error only occurs for the first time before overall submodule initialization.
Now, if I only remove the cloned KernelBench dir, then a simple submodule update seems sufficient.
Might be down to some git caching?

However, I'm able to reproduce the error in a freshly cloned repo.
Running git submodule update ingress/KernelBench/KernelBench returns:

Submodule path 'ingress/KernelBench/KernelBench' not initialized
Maybe you want to use 'update --init'?

Not sure about best practices here. But not necessarily a blocker as it's easily fixable by following git's error msg.

[rolfmorel]

Is there any option for torch-mlir to lower these constant into input args? That is prevent dumping constants at all.

I tried to find it yesterday but had no luck. @Groverkss would you have some pointers?

[banach-space]

Thanks for working on this!

A bit tangential, but shouldn't we set-up some CI to make sure that the code uploaded to Lighthouse does work? That would be my preference and I am keen to help, just a bit blocked ATM 😅 (e.g. llvm/torch-mlir-release#22).

[rengolin]

Thanks for working on this!

A bit tangential, but shouldn't we set-up some CI to make sure that the code uploaded to Lighthouse does work? That would be my preference and I am keen to help, just a bit blocked ATM 😅 (e.g. llvm/torch-mlir-release#22).

Not tangential at all, we definitely should. At the very least to make sure we pull the right deps, etc.

@adam-smnk, it should be trivial to create a workflow file that runs this example. Then, when @rolfmorel merges the Kernel Bench, we can add that to the CI, too. It will be a bit loose in the beginning, but then we can start creating some test scripts later.

@banach-space, it would be great if you could check it works on Arm, too, after your PR to torch-mlir, so that we can have both architectures building early on.

[Groverkss]

Is there any option for torch-mlir to lower these constant into input args? That is prevent dumping constants at all.

I tried to find it yesterday but had no luck. @Groverkss would you have some pointers?

You can control how mutable tensors are lowered using FXImporterHooks: https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/extras/fx_importer.py#L470

See how IREE does it for parameters: https://github.com/iree-org/iree-turbine/blob/5e50318baf6d436f48413788e5c7ea11425d130f/iree/turbine/aot/support/procedural/exported_program.py#L268

I would ideally recommend against promoting constant arguements to input arguements, because that's not how these things work in the real world, and instead do something like what iree does with "util.global_load". Maybe you can use "ml_program.global_load". Promoting to constant arguements is an easy short term solution if you are only compliing but not running things. It's a fake solution.

This is also why i recommended against using import_and_export hook because you have no control over the fx graph representation now.