Auto-tuning of thread configuration via TornadoExecutionPlan.withAutoTune()

[mikepapadim]

Summary

Introduce auto-tuning of the thread configuration (local work-group / thread-block size) in TornadoVM, exposed through a new fluent Execution Plan call:

TornadoExecutionResult result = new TornadoExecutionPlan(immutableTaskGraph)
        .withAutoTune()                 // <-- new
        .execute();

On the first execute() after withAutoTune(), TornadoVM sweeps a set of candidate local work-group sizes for each tunable task, measures the on-device kernel time, selects the fastest configuration per task, pins it into the effective GridScheduler, and reuses it for every subsequent execute() (zero tuning overhead thereafter).

First iteration scope: CUDA and OpenCL backends only. PTX, SPIR-V and Metal are out of scope for v1 (tasks on those devices are left at the default configuration; see Backend gating).

Motivation: the local work-group size is one of the highest-impact GPU performance knobs, and the optimal value is device- and kernel-specific. Today users either accept the driver/runtime default (OCLKernelScheduler / CUDA*Scheduler) or hand-tune a GridScheduler by trial and error. withAutoTune() automates that search.

Background — how thread configuration works today

• WorkerGrid (WorkerGrid1D/2D/3D, AbstractWorkerGrid) holds globalWork, localWork, numOfWorkgroups, globalOffset for a task. setLocalWork(x,y,z) also recomputes numOfWorkgroups = globalWork / localWork.
• GridScheduler maps "taskGraphName.taskName" → WorkerGrid, attached via TornadoExecutionPlan.withGridScheduler(gs).
• At launch, TornadoVMInterpreter (≈ line 1153) reads gridScheduler.get(task.getId()) and uses the WorkerGrid's local/global work to configure the launch for both the OpenCL and CUDA backends.
• When no WorkerGrid local size is set, the per-backend kernel schedulers (OCLKernelScheduler.checkLocalWorkGroupFitsOnDevice, CUDANVIDIAGPUScheduler, CUDAAMDScheduler, generic schedulers) pick a default. This is the baseline autotune competes against.
• Per-task on-device kernel time is already recorded: ProfilerType.TASK_KERNEL_TIME is set per task in TornadoVMInterpreter (≈ line 1321) from device events. This is the measurement signal autotune uses.
• Device limits are available via TornadoTargetDevice.getDeviceMaxWorkGroupSize() / getDeviceMaxWorkItemSizes() and TornadoDevice.getDeviceMaxWorkgroupDimensions().

Autotune therefore does not require new codegen or launch machinery — it drives the existing GridScheduler + profiler.

API design

tornado-api

New fluent method on TornadoExecutionPlan (mirrors the existing withX() → ExecutionPlanType builder pattern, e.g. WithGridScheduler):

public TornadoExecutionPlan withAutoTune();
public TornadoExecutionPlan withAutoTune(AutoTuneConfig config);

• Returns a new WithAutoTune plan node under api/plan/types/.
• Sets an autoTune flag + AutoTuneConfig on ExecutorFrame (new fields + getters/setters, alongside gridScheduler, profilerMode).

New config object (builder, sensible defaults so withAutoTune() works with no args):

public final class AutoTuneConfig {
    // defaults
    int    warmupIterations = 2;
    int    measuredIterations = 5;              // per candidate
    Metric metric = Metric.KERNEL_TIME;         // TASK_KERNEL_TIME (device events)
    Strategy strategy = Strategy.EXHAUSTIVE;     // v1
    long   timeBudgetMillis = 0;                 // 0 = unbounded
    boolean verifyOutput = false;                // optional correctness guard
    long[]  customCandidates1D = null;           // null = built-in table
    // enums: Metric { KERNEL_TIME, TOTAL_TIME }, Strategy { EXHAUSTIVE, /* future: COARSE_TO_FINE, HEURISTIC */ }
}

Semantics

• withAutoTune() without a GridScheduler: TornadoVM learns each task's global work size from a single probe run (task metadata / parallel range), synthesizes WorkerGrids, then tunes localWork.
• withAutoTune() with a GridScheduler: the user's globalWork is respected; autotune overrides only localWork. Any user-set localWork is added to the candidate set as a seed.
• Idempotent after the first tuned run: the winning config is cached and reused (no re-tuning) until the plan is reconfigured or the cache key changes.
• Composable with withProfiler, withWarmUpIterations, withDevice. Conflicts to define: withDefaultScheduler() after withAutoTune() clears the tuned grid.

Runtime design (tornado-runtime)

New orchestrator AutoTuner invoked by TornadoTaskGraph / TornadoVM on the first execute() when the frame's autoTune flag is set:

1. Select tunable tasks — for each task, resolve its target device backend (TornadoVMBackendType). Keep only OPENCL and CUDA tasks; skip others (leave default, log at info).
2. Resolve global work & device limits — from the attached WorkerGrid if present, else from a probe execution that captures the task's global range. Query getDeviceMaxWorkGroupSize() / getDeviceMaxWorkItemSizes().
3. Generate candidates (see below).
4. Benchmark — for each candidate: set WorkerGrid.setLocalWork(...), run warmupIterations (reuse existing warm-up path), then measuredIterations measured runs with the internal profiler forced on; record TASK_KERNEL_TIME; take the median (robust to jitter). Respect timeBudgetMillis.
5. Select the minimum-time candidate per task; write its localWork/numOfWorkgroups into the effective GridScheduler.
6. Cache the result keyed by (deviceIdentifier, taskId, kernelSignatureHash, globalWorkSize); reuse on subsequent executes.

Per-task tuning uses TASK_KERNEL_TIME, so multi-task graphs tune each task independently on the same launch. (Joint/global-work tuning is a Phase 2 item.)

Candidate search space (v1 = EXHAUSTIVE)

• 1D: warp/wavefront multiples {32, 64, 96, 128, 192, 256, 512, 1024} (NVIDIA warp 32; AMD wavefront 64), plus the driver default (localWork = null) as a baseline. Clamp to min(getDeviceMaxWorkGroupSize, getDeviceMaxWorkItemSizes[0]). Require the size to divide the global work size (correctness for the @Parallel loop API); non-dividing sizes are skipped in v1.
• 2D: {(8,8),(16,16),(16,8),(8,16),(32,8),(32,4),(4,32),(64,4)} with product ≤ max work-group size and each dim ≤ its max work-item size, dividing the respective global dims.
• 3D: {(4,4,4),(8,8,4),(8,4,4),(4,4,8)} under the same constraints.
• Pluggable AutoTuneSearchStrategy interface so future strategies (coarse-to-fine, occupancy-model-guided pruning, heuristic) drop in without API change.

Correctness

• @Parallel loop-API kernels: only divisor block sizes are used, so the iteration space is covered exactly.
• KernelContext kernels: the block size is semantically visible to user code (localIdx, barriers, local memory). v1 policy: only autotune KernelContext tasks when the user opts in (e.g. AutoTuneConfig.verifyOutput(true) or a dedicated flag) — otherwise leave them at the user-provided grid, because an arbitrary block size can change results. Default-safe.
• Optional verifyOutput: checksum the output of the first candidate and assert subsequent candidates match; discard any that diverge.

Backend gating (v1)

• Tunable only when the task's device backend is OPENCL or CUDA.
• If the whole plan targets an unsupported backend, withAutoTune() is a no-op with a one-time warning.
• Mixed plans: tune the CUDA/OpenCL tasks, leave the rest at default.

Caching

• v1: in-memory cache for the lifetime of the TornadoExecutionPlan.
• Phase 2: optional on-disk persistent cache (-Dtornado.autotune.cache=<path>, default e.g. ~/.tornadovm/autotune.json) so repeated JVM runs skip tuning; keyed by device identity + kernel signature + global size.

Files expected to change / add

• tornado-api: TornadoExecutionPlan.withAutoTune(), api/plan/types/WithAutoTune.java, AutoTuneConfig.java, ExecutorFrame (flag + config), docs.
• tornado-runtime: AutoTuner orchestrator + AutoTuneSearchStrategy (+ ExhaustiveStrategy), hook in TornadoTaskGraph/TornadoVM execute path, candidate generation using getDeviceMaxWorkGroupSize/getDeviceMaxWorkItemSizes, result cache, TASK_KERNEL_TIME readout.
• No backend codegen changes required (drives existing GridScheduler + interpreter launch path).

Testing & benchmarks

• Unit tests (tornado-unittests, gated to run on OpenCL/CUDA):
  • 1D: SAXPY / vector-add — assert autotuned kernel time ≤ default and output correct.
  • 2D: matrix multiply — assert a valid block is chosen and result matches the sequential reference.
  • No-op / gating: PTX/SPIR-V/Metal task → autotune skipped, execution still correct.
  • Cache: second execute() performs no re-tuning (assert tuning invoked once).
• WoW benchmark: matrix multiply / DFT autotuned vs default GridScheduler, reporting speedup on NVIDIA and Intel/AMD OpenCL devices.

Phasing

• Phase 1 (this issue, MVP): API (withAutoTune, AutoTuneConfig, WithAutoTune), AutoTuner with EXHAUSTIVE 1D/2D/3D block sweep, per-task TASK_KERNEL_TIME measurement, in-memory cache, CUDA + OpenCL gating, KernelContext opt-in safety, unit tests + one WoW benchmark.
• Phase 2 (future): persistent on-disk cache; heuristic / coarse-to-fine / occupancy-guided search; global-work (grid-size) tuning and multi-task joint tuning; PTX / SPIR-V / Metal support; integration with dynamic reconfiguration (DRMode).

Open questions

1. Default candidate set — ship the table above, or make it device-vendor-aware from day one?
2. Should KernelContext tasks be tunable by default (with verifyOutput) or strictly opt-in?
3. Metric default: median TASK_KERNEL_TIME (proposed) vs end-to-end.
4. Persist the tuned cache to disk in v1, or defer to Phase 2 (proposed)?

[mikepapadim]

Detailed design for #778