benchmarks.KernelBench.iree_compilation

IREE Kernel Compilation Script using iree-turbine

This script automates the compilation of PyTorch models (defined in 'level' directories) into optimized IREE artifacts (.vmfb) for various configurable targets (e.g., NVIDIA GPUs, ARM CPUs, RISC-V CPUs).

It performs the following steps: 1. Parses command-line arguments for input, output, and report directories, as well as kernel levels and target names to compile for. 2. Discovers all '*.py' kernel files in the specified 'level' directories (e.g., third_party/KernelBench/KernelBench/level1). 3. For each kernel file, it dynamically loads the 'Model' class, 'get_init_inputs()', and 'get_inputs()'. 4. For each specified target: a. Determines the correct precision (e.g., torch.float16 or torch.float32) for the target. b. Instantiates the model and inputs in that precision. c. Exports the PyTorch model to Torch MLIR using iree.turbine.aot.export() and saves it as 'kernel.mlir'. d. Compiles 'kernel.mlir' to 'artifact.vmfb' by invoking 'iree-compile' with the target-specific flags. e. Dumps intermediate MLIR stages into a 'compilation_phases' subdirectory. 5. Saves compiled artifacts in a structured output directory (e.g., benchmark/KernelBench/results/level1/A100/kernel_name/artifact.vmfb). 6. Logs all operations to 'iree_compile.log' in the report directory. 7. Generates detailed 'iree_compile_report.json' and 'iree_compile_report.md' summaries in the report directory, including success percentages per-level, per-target, and overall.