Skip to content

Commit fa1141e

Browse files
rolfmorelrolfmorel
committed
[ingress][pytorch] Basic KernelBench to MLIR conversion
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.
1 parent 20c7e72 commit fa1141e

File tree

2 files changed

+186
-8
lines changed

2 files changed

+186
-8
lines changed

‎ingress/Torch-MLIR/convert-kernel-bench-to-mlir.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
#!/usr/bin/env python3
2+
3+
import importlib
4+
import importlib.util
5+
6+
from pathlib import Path
7+
8+
from mlir import ir, passmanager
9+
from torch_mlir import fx
10+
11+
kernels_as_pytorch_folder = Path(__file__).parent / "KernelBench" / "KernelBench"
12+
kernels_as_pytorch_level1 = kernels_as_pytorch_folder / "level1"
13+
kernels_as_pytorch_level2 = kernels_as_pytorch_folder / "level2"
14+
15+
kernels_as_mlir_folder = Path(__file__).parent / "cache"
16+
kernels_as_mlir_level1 = kernels_as_mlir_folder / "level1"
17+
kernels_as_mlir_level1.mkdir(parents=True, exist_ok=True)
18+
kernels_as_mlir_level2 = kernels_as_mlir_folder / "level2"
19+
kernels_as_mlir_level2.mkdir(parents=True, exist_ok=True)
20+
21+
level1, level2 = Path("level1"), Path("level2")
22+
ignore_list = [
23+
level1 / "12_Matmul_with_diagonal_matrices_.py", # torch.operator "torch.aten.diag"
24+
level1
25+
/ "34_InstanceNorm.py", # LLVM ERROR: SmallVector unable to grow. Requested capacity (93898875033000)
26+
level1
27+
/ "72_conv_transposed_3D_asymmetric_input_asymmetric_kernel___strided_padded_grouped_.py", # Bare exception during torch-backend-to-linalg-on-tensors-backend-pipeline
28+
level1
29+
/ "89_cumsum.py", # Dialect `tm_tensor' not found for custom op 'tm_tensor.scan'
30+
level1
31+
/ "90_cumprod.py", # Dialect `tm_tensor' not found for custom op 'tm_tensor.scan'
32+
level1
33+
/ "91_cumsum_reverse.py", # Dialect `tm_tensor' not found for custom op 'tm_tensor.scan'
34+
level1
35+
/ "92_cumsum_exclusive.py", # Dialect `tm_tensor' not found for custom op 'tm_tensor.scan'
36+
level1
37+
/ "93_masked_cumsum.py", # Dialect `tm_tensor' not found for custom op 'tm_tensor.scan'
38+
level1
39+
/ "95_CrossEntropyLoss.py", # Bare exception during torch-backend-to-linalg-on-tensors-backend-pipeline
40+
level1
41+
/ "96_HuberLoss.py", # Bare exception during torch-backend-to-linalg-on-tensors-backend-pipeline
42+
level1
43+
/ "97_ScaledDotProductAttention.py", # AssertionError: Torch not compiled with CUDA enabled
44+
level1
45+
/ "99_TripletMarginLoss.py", # Bare exception during torch-backend-to-linalg-on-tensors-backend-pipeline
46+
level2
47+
/ "17_Conv2d_InstanceNorm_Divide.py", # LLVM ERROR: SmallVector unable to grow. Requested capacity (94899412484104)
48+
level2
49+
/ "18_Matmul_Sum_Max_AvgPool_LogSumExp_LogSumExp.py", # error: failed to legalize operation 'torch.constant.int'
50+
level2
51+
/ "22_Matmul_Scale_ResidualAdd_Clamp_LogSumExp_Mish.py", # error: failed to legalize operation 'torch.constant.int'
52+
level2
53+
/ "28_BMM_InstanceNorm_Sum_ResidualAdd_Multiply.py", # LLVM ERROR: SmallVector unable to grow. Requested capacity (94899412484104)
54+
level2
55+
/ "42_ConvTranspose2d_GlobalAvgPool_BiasAdd_LogSumExp_Sum_Multiply.py", # error: failed to legalize operation 'torch.constant.int'
56+
level2
57+
/ "43_Conv3d_Max_LogSumExp_ReLU.py", # error: failed to legalize operation 'torch.constant.int'
58+
level2
59+
/ "45_Gemm_Sigmoid_LogSumExp.py", # error: failed to legalize operation 'torch.constant.int'
60+
level2
61+
/ "51_Gemm_Subtract_GlobalAvgPool_LogSumExp_GELU_ResidualAdd.py", # error: failed to legalize operation 'torch.constant.int'
62+
level2
63+
/ "52_Conv2d_Activation_BatchNorm.py", # failed to legalize operation 'torch.operator'
64+
level2 / "55_Matmul_MaxPool_Sum_Scale.py", # MLIR file too big: 16G
65+
level2 / "59_Matmul_Swish_Scaling.py", # MLIR file too big: 16G
66+
level2 / "56_Matmul_Sigmoid_Sum.py", # MLIR file too big: 16G
67+
level2 / "66_Matmul_Dropout_Softmax.py", # MLIR file too big: 4G
68+
level2 / "68_Matmul_Min_Subtract.py", # MLIR file too big: 4G
69+
level2 / "94_Gemm_BiasAdd_Hardtanh_Mish_GroupNorm.py", # MLIR file too big: 1G
70+
level2 / "33_Gemm_Scale_BatchNorm.py", # MLIR file too big: 1G
71+
level2 / "88_Gemm_GroupNorm_Swish_Multiply_Swish.py", # MLIR file too big: 1G
72+
level2 / "75_Gemm_GroupNorm_Min_BiasAdd.py", # MLIR file too big: 1G
73+
level2 / "84_Gemm_BatchNorm_Scaling_Softmax.py", # MLIR file too big: 1G
74+
level2 / "97_Matmul_BatchNorm_BiasAdd_Divide_Swish.py", # MLIR file too big: 1G
75+
level2 / "62_Matmul_GroupNorm_LeakyReLU_Sum.py", # MLIR file too big: 1G
76+
level2 / "30_Gemm_GroupNorm_Hardtanh.py", # MLIR file too big: 1G
77+
level2 / "95_Matmul_Add_Swish_Tanh_GELU_Hardtanh.py", # MLIR file too big: 1G
78+
level2 / "29_Matmul_Mish_Mish.py", # MLIR file too big: 1G
79+
level2 / "99_Matmul_GELU_Softmax.py", # MLIR file too big: 1G
80+
level2 / "98_Matmul_AvgPool_GELU_Scale_Max.py", # MLIR file too big: 1G
81+
level2 / "80_Gemm_Max_Subtract_GELU.py", # MLIR file too big: 1G
82+
level2 / "81_Gemm_Swish_Divide_Clamp_Tanh_Clamp.py", # MLIR file too big: 1G
83+
level2 / "12_Gemm_Multiply_LeakyReLU.py", # MLIR file too big: 1G
84+
level2 / "53_Gemm_Scaling_Hardtanh_GELU.py", # MLIR file too big: 1G
85+
level2 / "9_Matmul_Subtract_Multiply_ReLU.py", # MLIR file too big: 1G
86+
level2 / "70_Gemm_Sigmoid_Scaling_ResidualAdd.py", # MLIR file too big: 1G
87+
level2 / "86_Matmul_Divide_GELU.py", # MLIR file too big: 1G
88+
level2 / "63_Gemm_ReLU_Divide.py", # MLIR file too big: 1G
89+
level2 / "76_Gemm_Add_ReLU.py", # MLIR file too big: 1G
90+
level2 / "14_Gemm_Divide_Sum_Scaling.py", # MLIR file too big: 1G
91+
level2 / "39_Gemm_Scale_BatchNorm.py", # MLIR file too big: 256M
92+
level2 / "41_Gemm_BatchNorm_GELU_ReLU.py", # MLIR file too big: 256M
93+
level2 / "40_Matmul_Scaling_ResidualAdd.py", # MLIR file too big: 256M
94+
level2 / "37_Matmul_Swish_Sum_GroupNorm.py", # MLIR file too big: 64.3M
95+
level2
96+
/ "58_ConvTranspose3d_LogSumExp_HardSwish_Subtract_Clamp.py", # error: failed to legalize operation 'torch.constant.int'
97+
level2
98+
/ "64_Gemm_LogSumExp_LeakyReLU_LeakyReLU_GELU_GELU.py", # error: failed to legalize operation 'torch.constant.int'
99+
level2
100+
/ "79_Conv3d_Multiply_InstanceNorm_Clamp_Multiply_Max.py", # LLVM ERROR: SmallVector unable to grow. Requested capacity (94312016449768)
101+
level2
102+
/ "92_Conv2d_GroupNorm_Tanh_HardSwish_ResidualAdd_LogSumExp.py", # error: failed to legalize operation 'torch.constant.int'
103+
]
104+
105+
106+
ctx = ir.Context()
107+
pm = passmanager.PassManager(context=ctx)
108+
pm.add("linalg-specialize-generic-ops")
109+
110+
for pytorch_level, mlir_level in (
111+
(kernels_as_pytorch_level1, kernels_as_mlir_level1),
112+
(kernels_as_pytorch_level2, kernels_as_mlir_level2),
113+
):
114+
for kernel_pytorch_file in pytorch_level.iterdir():
115+
level_and_kernel = (
116+
Path(kernel_pytorch_file.parent.name) / kernel_pytorch_file.name
117+
)
118+
if level_and_kernel in ignore_list or not kernel_pytorch_file.is_file():
119+
print(
120+
f"Skipping: {kernel_pytorch_file.parent.name}/{kernel_pytorch_file.name}"
121+
)
122+
continue
123+
124+
module_name = kernel_pytorch_file.stem
125+
126+
kernel_as_mlir_path = mlir_level / (module_name + ".mlir")
127+
if kernel_as_mlir_path.exists():
128+
print(
129+
f"Already in cache: {kernel_pytorch_file.parent.name}/{kernel_pytorch_file.name}"
130+
)
131+
continue
132+
print(
133+
f"Processing: {kernel_pytorch_file.parent.name}/{kernel_pytorch_file.name}"
134+
)
135+
136+
module_spec = importlib.util.spec_from_file_location(
137+
module_name, kernel_pytorch_file
138+
)
139+
140+
if module_spec is None or module_spec.loader is None:
141+
print(f"Error: Could not create module spec for {kernel_pytorch_file}")
142+
continue
143+
144+
module = importlib.util.module_from_spec(module_spec)
145+
# Execute the module to load its contents
146+
module_spec.loader.exec_module(module)
147+
148+
if not all(
149+
hasattr(module, a) for a in ("Model", "get_inputs", "get_init_inputs")
150+
):
151+
print(f"Error: module in file {kernel_pytorch_file} not a proper benchmark")
152+
continue
153+
154+
# TODO: check hasattr(module, "in_features") etc and adjust to sizes that are more tractable for torch-mlir
155+
156+
try:
157+
m = fx.export_and_import(
158+
module.Model(*module.get_init_inputs()),
159+
*module.get_inputs(),
160+
output_type=fx.OutputType.LINALG_ON_TENSORS,
161+
)
162+
except Exception as e:
163+
print(f"Error: got the following error converting {kernel_pytorch_file}")
164+
raise e
165+
166+
before_clean_up = "//" + str(m)[:-1].replace("\n", "\n//") + "\n"
167+
# Cross boundary from torch-mlir's mlir to environment's mlir
168+
m = ir.Module.parse(str(m), context=ctx)
169+
# Run clean-up, e.g. linalg-"specialization" passes to raise within Linalg.
170+
try:
171+
pm.run(m.operation) # cleanup
172+
except Exception as e:
173+
print(f"Error: got the following error cleaning up {module_name}")
174+
raise e
175+
176+
with kernel_as_mlir_path.open("w") as f:
177+
print("// Torch-MLIR output:", file=f)
178+
print(before_clean_up, file=f)
179+
print("// MLIR output after clean-up:", file=f)
180+
print(m, file=f)

‎ingress/Torch-MLIR/install-virtualenv.sh‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,16 @@ else
99
DEVICE_TYPE=$(lspci | grep VGA)
1010
fi
1111

12-
13-
# Install torch-mlir inside a virtual environment
1412
echo "First ensure uv is installed"
15-
1613
python -m pip install uv --upgrade
1714

1815
echo "Preparing the virtual environment"
1916
python -m uv venv torch-mlir-venv --python 3.12
2017

21-
#echo "Preparing the virtual environment"
22-
#python3 -m venv torch-mlir-venv
2318
source torch-mlir-venv/bin/activate
24-
uv pip install --upgrade pip wheel
19+
20+
echo "Installing mlir-python-bindings and numpy"
21+
uv pip install numpy mlir-python-bindings -f https://makslevental.github.io/wheels
2522

2623
# GPU support ("AMD", "NVIDIA", "Intel")
2724
EXTRA_INDEX_URL=""
@@ -46,13 +43,14 @@ if [ $? != 0 ]; then
4643
exit 1
4744
fi
4845

49-
5046
echo "Installing torch-mlir"
51-
# This only seems to work on Ubuntu
5247
uv pip install --pre torch-mlir \
5348
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
5449
-f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
5550
if [ $? != 0 ]; then
5651
exit 1
5752
fi
5853

54+
# TODO(rolf): move to a more appropriate place
55+
echo "Obtaining KernelBench repo"
56+
git clone https://github.com/ScalingIntelligence/KernelBench.git

0 commit comments

Comments
 (0)