feat: add 2D FFT benchmark cases to all suites

- Smoke: +8 cases (16x16, 256x256 x c2c/z2z x fwd/inv)
- Typical: +24 cases (6 2D sizes x c2c/z2z x fwd/inv)
- Full: +72 cases (9 2D sizes incl Bluestein x c2c/z2z x fwd/inv x batch 1,4)
- Update run_benchmark fixture to detect 2D sizes (string with 'x')
  and pass --rank 2 --shape NxM to CLI

Author: Artlesbol

benchmark/utils/pytest_plugin.py

@@ -127,23 +127,29 @@ def invoke(*arguments: str, env: dict[str, str] | None = None, timeout: int = 24
 def run_benchmark(invoke_cli, bench_warmup, bench_iters):
     """Invoke the bench subcommand with configured warmup/iters."""
 
-    def _run(size: int, api: str = "c2c", direction: str = "forward", batch: int = 1):
-        return invoke_cli(
+    def _run(size, api: str = "c2c", direction: str = "forward", batch: int = 1):
+        is_2d = isinstance(size, str) and "x" in size
+        cmd = [
             "bench",
             "--api",
             api,
             "--direction",
             direction,
-            "--shape",
-            str(size),
+        ]
+        if is_2d:
+            cmd += ["--rank", "2", "--shape", size]
+        else:
+            cmd += ["--shape", str(size)]
+        cmd += [
             "--batch",
             str(batch),
             "--warmup",
             str(bench_warmup),
             "--iters",
             str(bench_iters),
             "--print-path",
-        )
+        ]
+        return invoke_cli(*cmd)
 
     return _run
 

benchmark/utils/suites.py

@@ -9,6 +9,9 @@
 # All APIs from ctest coverage
 ALL_APIS = ["c2c", "z2z", "r2c", "c2r", "d2z", "z2d"]
 
+# 2D complex APIs only (rank 2 does not support real transforms yet)
+ALL_APIS_2D = ["c2c", "z2z"]
+
 # Direction support per API
 API_DIRECTIONS = {
     "c2c": ["forward", "inverse"],
@@ -19,20 +22,36 @@
     "z2d": ["inverse"],
 }
 
+# 2D sizes: smooth (CT path), non-square, Bluestein (prime dimensions)
+ALL_SIZES_2D = ["16x16", "64x64", "128x256", "256x128", "256x256", "512x512"]
+BLUESTEIN_SIZES_2D = ["23x23", "997x16", "16x997"]
+
 # ── Suite definitions ──────────────────────────────────────────
 
-# Smoke (24 cases): 3 key sizes x batch=1 x all API/direction combos
-SMOKE = {"sizes": [16, 256, 2048], "batches": [1]}
+# Smoke (24 1D cases + 4 2D cases): key sizes x batch=1 x all API/direction combos
+SMOKE = {
+    "sizes": [16, 256, 2048],
+    "batches": [1],
+    "sizes_2d": ["16x16", "256x256"],
+    "batches_2d": [1],
+}
 
-# Typical (136 cases): all sizes x batch=1 + extra multi-batch on key sizes
+# Typical (136 1D + 24 2D cases)
 TYPICAL = {
     "sizes": ALL_SIZES,
     "batches": [1],
     "extra": [(256, 4), (256, 256), (2048, 4), (2048, 256)],
+    "sizes_2d": ALL_SIZES_2D,
+    "batches_2d": [1],
 }
 
-# Full (312 cases): all sizes x all batches x all API/direction combos
-FULL = {"sizes": ALL_SIZES, "batches": ALL_BATCHES}
+# Full (312 1D + 48 2D cases)
+FULL = {
+    "sizes": ALL_SIZES,
+    "batches": ALL_BATCHES,
+    "sizes_2d": ALL_SIZES_2D + BLUESTEIN_SIZES_2D,
+    "batches_2d": [1, 4],
+}
 
 SUITES = {"smoke": SMOKE, "typical": TYPICAL, "full": FULL}
 
@@ -49,11 +68,14 @@ def expand_params(suite: dict):
 
     Handles the 'extra' key for TYPICAL suite: additional (size, batch)
     pairs beyond the base sizes x batches cross product.
+
+    Also yields 2D entries from 'sizes_2d' x 'batches_2d' x ALL_APIS_2D.
+    2D sizes are strings like "256x256" (detected by the caller to set rank=2).
     """
     apis = ALL_APIS
     seen = set()
 
-    # Base: sizes x batches x apis
+    # Base 1D: sizes x batches x apis
     for size in suite["sizes"]:
         for batch in suite["batches"]:
             for api in apis:
@@ -63,11 +85,21 @@ def expand_params(suite: dict):
                         seen.add(key)
                         yield key
 
-    # Extra (size, batch) pairs applied to all API/direction combos
+    # Extra 1D (size, batch) pairs applied to all API/direction combos
     for size, batch in suite.get("extra", []):
         for api in apis:
             for direction in API_DIRECTIONS[api]:
                 key = (size, batch, api, direction)
                 if key not in seen:
                     seen.add(key)
                     yield key
+
+    # 2D: sizes_2d x batches_2d x ALL_APIS_2D (forward only)
+    for size_2d in suite.get("sizes_2d", []):
+        for batch in suite.get("batches_2d", [1]):
+            for api in ALL_APIS_2D:
+                for direction in API_DIRECTIONS[api]:
+                    key = (size_2d, batch, api, direction)
+                    if key not in seen:
+                        seen.add(key)
+                        yield key