Repository: nod-ai/iree-kernel-benchmark
Pull request #50: Reorganise module layout to be more hermetic
Author: login: andfau-amd
Created: 2025-03-11T17:35:59Z | Merged: 2025-03-12T18:17:24Z

It was strange that the shared utilities had to be installed as a package, as they should be just an implementation detail. This commit moves all the modules into a new `iree-kernel-benchmark` directory, creating an implicit namespace package that can optionally be installed, makes the command-line tools be invoked in module style, and switches to using explicit relative imports for all internal dependencies.

Resolves #49.

Linked issue #49: Fix project python imports/dependencies
Currently the way the project is structured is really weird. We shouldn't be installing `common_tools` for example.
I'd expect that we launch `python -m gemmbench/gemm_bench` and it finds `common_tools` on its own.

You can compare this to how we use python modules in the tuner: https://github.com/nod-ai/shark-ai/tree/main/tuner.

Diff:
diff --git a/.github/workflows/run_bench.yml b/.github/workflows/run_bench.yml
index 1347552..1ccbe9c 100644
--- a/.github/workflows/run_bench.yml
+++ b/.github/workflows/run_bench.yml
@@ -29,45 +29,44 @@ jobs:
           pip install --upgrade pip
           pip install --pre --no-cache-dir --find-links https://iree.dev/pip-release-links.html iree-base-compiler iree-base-runtime --upgrade
           pip install -r requirements.txt
-          pip install --no-compile --pre --upgrade -e common_tools
           pip install iree-turbine@git+https://github.com/iree-org/iree-turbine.git@main
 
       - name: Convolutions
         run: |
           source bench_venv/bin/activate
-          python convbench/conv_bench.py
+          python -m iree-kernel-benchmark.convbench
 
       - name: TK Convolutions
         run: |
           source bench_venv/bin/activate
-          python convbench/conv_bench.py --tk
+          python -m iree-kernel-benchmark.convbench --tk
 
       - name: Attention
         run: |
           source bench_venv/bin/activate
-          python attentionbench/attention_bench.py
+          python -m iree-kernel-benchmark.attentionbench
 
       - name: TK GEMM
         run: |
           source bench_venv/bin/activate
-          python gemmbench/gemm_bench.py --tk
+          python -m iree-kernel-benchmark.gemmbench --tk
 
       - name: GEMM
         run: |
           source bench_venv/bin/activate
-          python gemmbench/gemm_bench.py
+          python -m iree-kernel-benchmark.gemmbench
 
       - name: Roofline Plots
         run: |
           source bench_venv/bin/activate
-          python convbench/conv_bench.py --roofline results/iree_conv.csv --plot results/iree_conv_i8.png --dtype i8
-          python convbench/conv_bench.py --roofline results/iree_conv.csv --plot results/iree_conv_f16.png --dtype f16
-          python convbench/conv_bench.py --roofline results/iree_conv_tk.csv --plot results/iree_conv_tk_f16.png --dtype f16
-          python convbench/conv_bench.py --roofline results/iree_attention.csv --plot results/iree_attention_fp16.png --dtype f16
-          python convbench/conv_bench.py --roofline results/iree_attention.csv --plot results/iree_attention_fp8.png --dtype f8E4M3FNUZ
-          python convbench/conv_bench.py --roofline results/iree_gemm.csv --plot results/iree_gemm.png
-          python convbench/conv_bench.py --roofline results/iree_gemm_tk.csv --plot results/iree_gemm_tk.png
-          python convbench/conv_bench.py --roofline results/iree_gemm.csv,results/iree_gemm_tk.csv,results/iree_attention.csv,results/iree_conv.csv,results/iree_conv_tk.csv --plot results/combined.png
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_conv.csv --plot results/iree_conv_i8.png --dtype i8
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_conv.csv --plot results/iree_conv_f16.png --dtype f16
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_conv_tk.csv --plot results/iree_conv_tk_f16.png --dtype f16
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_attention.csv --plot results/iree_attention_fp16.png --dtype f16
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_attention.csv --plot results/iree_attention_fp8.png --dtype f8E4M3FNUZ
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_gemm.csv --plot results/iree_gemm.png
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_gemm_tk.csv --plot results/iree_gemm_tk.png
+          python -m iree-kernel-benchmark.convbench --roofline results/iree_gemm.csv,results/iree_gemm_tk.csv,results/iree_attention.csv,results/iree_conv.csv,results/iree_conv_tk.csv --plot results/combined.png
 
       - name: Upload benchmark results
         uses: actions/upload-artifact@v4
diff --git a/README.md b/README.md
index 8d49ab3..1cc99de 100644
--- a/README.md
+++ b/README.md
@@ -12,9 +12,10 @@ Create a python environment and install the requirements for the project:
 python3.11 -m venv bench_venv
 source bench_venv/bin/activate
 pip install -r requirements.txt
-pip install --no-compile --pre --upgrade -e common_tools
 ```
 
+If you don't want to have to `cd` into the repo to use the tools, you may optionally run `pip install --no-compile --pre --upgrade -e .` to install the `iree-kernel-benchmark` package.
+
 If you plan to run the TK benchmarks, also install iree-turbine with
 ```
 pip install iree-turbine@git+https://github.com/iree-org/iree-turbine.git@main
@@ -28,31 +29,31 @@ Refer to the respective problems.py file in the folder to see which shapes are b
 ### Convolution Benchmarking
 
 ```
-python convbench/conv_bench.py
+python -m iree-kernel-benchmark.convbench
 ```
 
 ### TK Convolution Benchmarking
 
 ```
-python convbench/conv_bench.py --tk
+python -m iree-kernel-benchmark.convbench --tk
 ```
 
 ### GEMM Benchmarking
 
 ```
-python gemmbench/gemm_bench.py
+python -m iree-kernel-benchmark.gemmbench
 ```
 
 ### TK GEMM Benchmarking
 
 ```
-python gemmbench/gemm_bench.py --tk
+python -m iree-kernel-benchmark.gemmbench --tk
 ```
 
 ### Attention Benchmarking
 
 ```
-python attentionbench/attention_bench.py
+python -m iree-kernel-benchmark.attentionbench
 ```
 
 ### Roofline
diff --git a/attentionbench/attention_bench.py b/iree-kernel-benchmark/attentionbench/__main__.py
similarity index 97%
rename from attentionbench/attention_bench.py
rename to iree-kernel-benchmark/attentionbench/__main__.py
index 0b69d29..97ba43f 100644
--- a/attentionbench/attention_bench.py
+++ b/iree-kernel-benchmark/attentionbench/__main__.py
@@ -7,9 +7,9 @@
 import csv
 import argparse
 import sys
-from utils import *
-from attention_utils import *
-from problems import get_attention_configs
+from ..utils import *
+from .attention_utils import *
+from .problems import get_attention_configs
 
 
 def compile_attention(tag, config, kernel_dir, vmfb_dir):
diff --git a/attentionbench/attention_utils.py b/iree-kernel-benchmark/attentionbench/attention_utils.py
similarity index 99%
rename from attentionbench/attention_utils.py
rename to iree-kernel-benchmark/attentionbench/attention_utils.py
index b057759..ea28925 100644
--- a/attentionbench/attention_utils.py
+++ b/iree-kernel-benchmark/attentionbench/attention_utils.py
@@ -1,4 +1,4 @@
-from utils import *
+from ..utils import *
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
diff --git a/attentionbench/problems.py b/iree-kernel-benchmark/attentionbench/problems.py
similarity index 98%
rename from attentionbench/problems.py
rename to iree-kernel-benchmark/attentionbench/problems.py
index faca8e4..523e339 100644
--- a/attentionbench/problems.py
+++ b/iree-kernel-benchmark/attentionbench/problems.py
@@ -1,4 +1,4 @@
-from attention_utils import AttentionConfig
+from .attention_utils import AttentionConfig
 
 
 def llm_sweep(dtype: str) -> list[AttentionConfig]:
diff --git a/common_tools/kernel_stats.py b/iree-kernel-benchmark/common_tools/kernel_stats.py
similarity index 100%
rename from common_tools/kernel_stats.py
rename to iree-kernel-benchmark/common_tools/kernel_stats.py
diff --git a/common_tools/plot_roofline_percents.py b/iree-kernel-benchmark/common_tools/plot_roofline_percents.py
similarity index 100%
rename from common_tools/plot_roofline_percents.py
rename to iree-kernel-benchmark/common_tools/plot_roofline_percents.py
diff --git a/convbench/conv_bench.py b/iree-kernel-benchmark/convbench/__main__.py
similarity index 97%
rename from convbench/conv_bench.py
rename to iree-kernel-benchmark/convbench/__main__.py
index 8b35526..277d65b 100644
--- a/convbench/conv_bench.py
+++ b/iree-kernel-benchmark/convbench/__main__.py
@@ -7,11 +7,11 @@
 import csv
 import argparse
 import sys
-from utils import *
-from conv_utils import *
-from problems import get_conv_configs, get_tk_conv_configs, get_conv_test_configs
+from ..utils import *
+from .conv_utils import *
+from .problems import get_conv_configs, get_tk_conv_configs, get_conv_test_configs
 
-from wave_conv_utils import compile_wave_conv_config
+from .wave_conv_utils import compile_wave_conv_config
 
 def compile_conv_iree(tag, config, kernel_dir, vmfb_dir, extra_compiler_args):
     mlir_file, vmfb_file, dump_path = compile_conv_config(config, kernel_dir, vmfb_dir, extra_compiler_args)
diff --git a/convbench/conv_utils.py b/iree-kernel-benchmark/convbench/conv_utils.py
similarity index 99%
rename from convbench/conv_utils.py
rename to iree-kernel-benchmark/convbench/conv_utils.py
index 0cd8be4..8d6335d 100644
--- a/convbench/conv_utils.py
+++ b/iree-kernel-benchmark/convbench/conv_utils.py
@@ -1,4 +1,4 @@
-from utils import *
+from ..utils import *
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
diff --git a/convbench/miopen_conv2d_shapes.txt b/iree-kernel-benchmark/convbench/miopen_conv2d_shapes.txt
similarity index 100%
rename from convbench/miopen_conv2d_shapes.txt
rename to iree-kernel-benchmark/convbench/miopen_conv2d_shapes.txt
diff --git a/convbench/problems.py b/iree-kernel-benchmark/convbench/problems.py
similarity index 97%
rename from convbench/problems.py
rename to iree-kernel-benchmark/convbench/problems.py
index 19350ae..d474546 100644
--- a/convbench/problems.py
+++ b/iree-kernel-benchmark/convbench/problems.py
@@ -1,4 +1,4 @@
-from conv_utils import ConvConfig
+from .conv_utils import ConvConfig
 import os
 
 def unet_sweep(op: str, input_dtype: str, output_dtype: str) -> list[ConvConfig]:
@@ -70,7 +70,7 @@ def get_conv_configs() -> list[tuple[str, ConvConfig]]:
 
     # MIOpen bf16 convs
     cwd = os.getcwd()
-    shapes_file_path = os.path.join(cwd, "convbench/miopen_conv2d_shapes.txt")
+    shapes_file_path = os.path.join(cwd, "iree-kernel-benchmark/convbench/miopen_conv2d_shapes.txt")
     miopen_configs = []
     with open(shapes_file_path, 'r') as file:
         for line in file:
diff --git a/convbench/wave_conv_utils.py b/iree-kernel-benchmark/convbench/wave_conv_utils.py
similarity index 94%
rename from convbench/wave_conv_utils.py
rename to iree-kernel-benchmark/convbench/wave_conv_utils.py
index f268cd6..3489068 100644
--- a/convbench/wave_conv_utils.py
+++ b/iree-kernel-benchmark/convbench/wave_conv_utils.py
@@ -1,8 +1,8 @@
-from utils import *
+from ..utils import *
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
-from conv_utils import ConvConfig
+from .conv_utils import ConvConfig
 import traceback
 
 try:
diff --git a/gemmbench/gemm_bench.py b/iree-kernel-benchmark/gemmbench/__main__.py
similarity index 98%
rename from gemmbench/gemm_bench.py
rename to iree-kernel-benchmark/gemmbench/__main__.py
index d75f0de..1e05b03 100644
--- a/gemmbench/gemm_bench.py
+++ b/iree-kernel-benchmark/gemmbench/__main__.py
@@ -12,9 +12,9 @@
 from pathlib import Path
 import argparse
 import sys
-from utils import *
-from gemm_utils import *
-from problems import get_gemm_configs, get_tk_gemm_configs, get_matching_configs
+from ..utils import *
+from .gemm_utils import *
+from .problems import get_gemm_configs, get_tk_gemm_configs, get_matching_configs
 
 
 def compile_gemm(tag, config, kernel_dir, vmfb_dir, target, extra_compiler_args, tk, dump_dir=None):
diff --git a/gemmbench/gemm_utils.py b/iree-kernel-benchmark/gemmbench/gemm_utils.py
similarity index 99%
rename from gemmbench/gemm_utils.py
rename to iree-kernel-benchmark/gemmbench/gemm_utils.py
index 04d14d0..a6872a2 100644
--- a/gemmbench/gemm_utils.py
+++ b/iree-kernel-benchmark/gemmbench/gemm_utils.py
@@ -15,7 +15,7 @@
 else:
     TURBINE_AVAILABLE=True
 
-from utils import *
+from ..utils import *
 import os
 import traceback
 
diff --git a/gemmbench/problems.py b/iree-kernel-benchmark/gemmbench/problems.py
similarity index 99%
rename from gemmbench/problems.py
rename to iree-kernel-benchmark/gemmbench/problems.py
index 08bcb68..ac3f32f 100644
--- a/gemmbench/problems.py
+++ b/iree-kernel-benchmark/gemmbench/problems.py
@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from gemm_utils import GemmConfig
+from .gemm_utils import GemmConfig
 
 import re
 
diff --git a/common_tools/utils/__init__.py b/iree-kernel-benchmark/utils/__init__.py
similarity index 100%
rename from common_tools/utils/__init__.py
rename to iree-kernel-benchmark/utils/__init__.py
diff --git a/common_tools/utils/bench_utils.py b/iree-kernel-benchmark/utils/bench_utils.py
similarity index 100%
rename from common_tools/utils/bench_utils.py
rename to iree-kernel-benchmark/utils/bench_utils.py
diff --git a/common_tools/setup.py b/setup.py
similarity index 74%
rename from common_tools/setup.py
rename to setup.py
index 8006069..052f6c8 100644
--- a/common_tools/setup.py
+++ b/setup.py
@@ -4,15 +4,15 @@
 #### KERNEL BENCHMARK SUITE SETUP ####
 
 setup(
-    name=f"utils",
+    name=f"iree-kernel-benchmark",
     version="1.0.0",
     author="SHARK Authors",
     author_email="esaimana@amd.com",
     description="Kernel Perf Tools",
     packages=find_namespace_packages(
         include=[
-            "utils",
-            "utils.*",
+            "iree-kernel-benchmark.*",
         ],
     ),
+    package_dir={'': '.'},
 )


Reviews:
user: login: kuhar
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T02:57:10Z
user: login: kuhar
body: > (resolves #49)

This should be a part of the PR description in order for the GitHub automations to trigger:  https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/using-keywords-in-issues-and-pull-requests
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-12T02:58:13Z
user: login: andfau-amd
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T12:01:10Z
user: login: kuhar
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T14:24:59Z
user: login: kuhar
body: LGTM. Thanks for cleaning this up.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-12T16:32:03Z

Review comments:
diff_hunk: @@ -29,7 +29,7 @@ jobs:
           pip install --upgrade pip
           pip install --pre --no-cache-dir --find-links https://iree.dev/pip-release-links.html iree-base-compiler iree-base-runtime --upgrade
           pip install -r requirements.txt
-          pip install --no-compile --pre --upgrade -e common_tools
+          pip install --no-compile --pre --upgrade -e .
path: .github/workflows/run_bench.yml
user: login: kuhar
body: Why do we prefer to install this instead of running python on nested modules directly? We don't have to do that in the tuner.
created_at: 2025-03-12T02:56:13Z
author_association: MEMBER
diff_hunk: @@ -29,7 +29,7 @@ jobs:
           pip install --upgrade pip
           pip install --pre --no-cache-dir --find-links https://iree.dev/pip-release-links.html iree-base-compiler iree-base-runtime --upgrade
           pip install -r requirements.txt
-          pip install --no-compile --pre --upgrade -e common_tools
+          pip install --no-compile --pre --upgrade -e .
path: .github/workflows/run_bench.yml
user: login: andfau-amd
body: I don't really know if there's a benefit to installing versus not installing, but as far as I know my changes aren't incompatible with choosing not to install. Would you prefer that I remove this line from the CI and from the README? (I guess I'd want to retain installing as an option though.)
created_at: 2025-03-12T12:01:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -29,7 +29,7 @@ jobs:
           pip install --upgrade pip
           pip install --pre --no-cache-dir --find-links https://iree.dev/pip-release-links.html iree-base-compiler iree-base-runtime --upgrade
           pip install -r requirements.txt
-          pip install --no-compile --pre --upgrade -e common_tools
+          pip install --no-compile --pre --upgrade -e .
path: .github/workflows/run_bench.yml
user: login: kuhar
body: In general, the fewer setup steps the better, so I'd prefer not to have it
created_at: 2025-03-12T14:24:59Z
author_association: MEMBER

PR comments:
user: login: andfau-amd
created_at: 2025-03-11T18:01:53Z
body: CI failed due to:

> ERROR: Could not install packages due to an OSError: [Errno 28] No space left on device

😅 Well, I guess that's not my fault. I did update the YAML at least.
author_association: CONTRIBUTOR
user: login: kuhar
created_at: 2025-03-12T03:00:37Z
body: Adding @makslevental for python module setup
author_association: MEMBER
user: login: andfau-amd
created_at: 2025-03-12T12:04:51Z
body: > This should be a part of the PR description in order for the GitHub automations to trigger

It being part of the commit message also counts:
![image](https://github.com/user-attachments/assets/cd70bb00-3cc0-493f-b59b-b9f117961ef3)

author_association: CONTRIBUTOR
user: login: kuhar
created_at: 2025-03-12T14:50:15Z
body: > It being part of the commit message also counts:

Thanks, TIL. Regardless, let's not start a new convention and have it down the PR description like in other projects, incl., IREE and llvm/mlir.
author_association: MEMBER
user: login: saienduri
created_at: 2025-03-13T17:17:31Z
body: I moved the runner to a cluster btw, so you shouldn't be seeing the space issues hopefully going forward (https://github.com/nod-ai/iree-kernel-benchmark/pull/51)
author_association: CONTRIBUTOR
user: login: andfau-amd
created_at: 2025-03-13T17:19:30Z
body: Thanks! Glad to see the CI passes now, I was fairly confident in the changes but it was unfortunate to not know for sure if they would pass.
author_association: CONTRIBUTOR