HGEMM Optimization Problem

• HGEMM optimization problem is the challenge of scheduling and mapping half-precision matrix multiplication to maximize throughput and energy efficiency on modern hardware.
• It employs multi-objective strategies combining analytical models, machine-learned predictions, and reinforcement learning to navigate vast configuration spaces under hardware constraints.
• Automated frameworks including ML-guided search, RL with LLM enhancements, compiler-level tuning, and analytical code generation deliver significant runtime and energy improvements.

The HGEMM optimization problem refers to the determination of high-performance, energy-efficient scheduling and mapping strategies for half-precision General Matrix Multiply (HGEMM) kernels on advanced hardware architectures, notably including GPUs, FPGAs, and heterogenous SoCs. Recent research highlights multi-objective optimization techniques, combining analytical, machine-learned, and reinforcement-learning approaches to search vast configuration spaces, often targeting throughput (GFLOPS), energy efficiency (GFLOPS/W), and, in communication-bound domains, data movement cost. Several paradigms, including hypergraph partitioning for sparse-matrix kernels, compiler-level optimization tactics, and deep RL-guided kernel auto-tuning, represent the state-of-the-art.

1. Mathematical Formulation of the HGEMM Problem

HGEMM computes $C = \alpha AB + \beta C$, typically with $\alpha=1$, $\beta=0$, for $A \in \mathbb{R}^{M \times K}$, $B \in \mathbb{R}^{K \times N}$, $C \in \mathbb{R}^{M \times N}$ in half-precision (fp16). The optimization objective is commonly to maximize the achievable throughput:

$$\text{TFLOPS} = \frac{2MNK}{10^{12}t_{\text{custom}}}$$

where $t_{\text{custom}}$ is the elapsed execution time for a kernel evaluated on the target hardware. Other objective functions include energy efficiency $E = \text{TFLOPS}/\text{Power}$, and, in multi-objective settings, Pareto rankings over $(\text{TFLOPS}, E)$. The mapping $\theta$, specifying tiling, parallelism, and loop unrolling, must satisfy hardware constraints (register capacity, memory bandwidth, device-specific limits) (Papalamprou et al., 10 Nov 2025, Su et al., 2 Dec 2025, Zhang et al., 2019).

2. Design Space: Parameters and Constraints

The parameter space for HGEMM mapping spans tile sizes $(\text{BM}, \text{BN}, \text{BK})$, parallelization degrees (e.g., $P_m, P_n, P_k$), loop orderings, vectorization widths, and buffering/prefetching modes:

• In a modern FPGA SoC (Versal ACAP), choices include:
  • Degree of parallelism $P_d$ along $M,N,K$ (AIE assignment)
  • Tile size $B_d$ buffered in PL for data reuse
  • Layout and assignment indicators $x_{t,r}$ ("tile $t$ assigned to resource $r$")
• On CUDA GPUs, action vectors specify thread/block configuration, WMMA abstraction controls, shared-memory pipelining parameters ($n_{\text{stage}}$), swizzle, zero-padding, register buffering, prefetch, and epilogue copy strategy (Su et al., 2 Dec 2025). Hardware constraints include:
• Register capacity for accumulators, broadcasted values, temporaries (Veras et al., 2016)
• On-chip memory (AIE scratchpad, PL buffer, GPU shared mem)
• Memory bandwidth (DDR, L2, transport channels)
• Physical resource usage (DSPs, LUTs, BRAM, URAM, etc.)

A plausible implication is that the combinatorial growth in configuration space (e.g., >6000 candidate tuples for ACAP, $10^3$ kernel configurations for CUDA-L2) necessitates automated exploration driven by model-based or RL strategies.

3. Optimization Algorithms and Search Strategies

Several distinct approaches have been proposed and evaluated:

a. ML-Guided Design Space Exploration (ACAP)

Gradient-boosted trees (XGBoost) predict kernel latency, power, and resource usage using derived and raw features ($$N_{\text{AIE}}, \rho=\text{FLOP}/N_{\text{AIE}}, R_{P_d}=d/P_d, R_{B_d}=d/B_d$$). The online search iterates over tuples $(P_M, P_N, P_K, B_M, B_N, B_K)$, evaluating candidate mappings in milliseconds, and extracts the Pareto frontier for throughput/energy objectives (Papalamprou et al., 10 Nov 2025).

b. Reinforcement Learning with LLM Guidance (CUDA-L2)

The kernel tuning process is framed as a Markov Decision Process:

• State incorporates $(M,N,K)$, profiling metrics (NCU occupancy, memory throughput), and kernel history.
• Actions select kernel configuration vectors (tiling, pipelining, prefetch strategy, layout).
• Rewards quantify speedup over reference kernels, penalize numerical error and code length.
• LLMs (671B DeepSeek, SFT on CUDA code) generate kernel variants via contrastive prompting, guiding the RL agent towards non-obvious optimizations (zero-padding, swizzle, ping-pong buffering). Policy optimization uses GRPO; the empirical search covers thousands of configurations (Su et al., 2 Dec 2025).

c. Compiler-Level Search Algorithms

Within the TVM framework, two search algorithms are prominent:

• Greedy Best-First Search (G-BFS): Uses a heuristic cost model, prioritizing local improvement and limited exploration radius per step ($\rho=5$ neighbors).
• Neighborhood Actor-Advantage Critic (N-A2C): Restricts RL actions to locally neighboring configurations ($k$-nearest), updating actor/critic networks on observed runtime cost. Both methods demonstrate drastic reductions in search space explored (0.1%) while achieving superior runtime savings over XGBoost and RNN-based tuners (Zhang et al., 2019).

d. Analytical Outer Product Code Generation

Analytical micro-kernel synthesis (outer-product decomposition) selects SIMD-tile shapes $(M_r,N_r)$, optimizing register usage and pipeline saturation. The instruction mix is statically scheduled considering hardware port allocations and latency, yielding near-peak performance for DLA kernels (Veras et al., 2016).

4. Cost Models and Pareto Optimization

Cost assessment across methodologies integrates empirical measurement and analytic estimates:

Model	Evaluated Metric	Features Used
ML (XGBoost)	Latency, Power, Usage	$M,N,K,P_*,B_*,N_{\text{AIE}},\rho,L/P_d,R/B_d$
Compiler-level	Execution Time	Tile sizes, loop order, vector width, overhead factor
Analytical	Cycles per micro-kernel	Outer-product shape, pipeline mix, register count
RL-Based	Reward (Speedup > baseline, code quality)	state/action, historic metrics, LLM-generated proposals

Pareto front extraction provides explicit trade-off curves. For Versal ACAP, throughput and energy-optimal solutions often diverge: maximal throughput uses maximal AIEs, while energy-optimized schedules reduce AIE count and increase PL tile sizes for data reuse (Papalamprou et al., 10 Nov 2025). CUDA-L2 demonstrates that RL-generated kernels outperform cuBLAS and cuBLASLt-AutoTuning across thousands of cases (up to +29% speedup in "server mode") (Su et al., 2 Dec 2025).

5. Insights, Challenges, and Best Practices

Empirical observations across research publications suggest:

• Analytical models can misestimate latency (by 10–20%) and ignore power, causing suboptimal scheduling (Papalamprou et al., 10 Nov 2025).
• Derived features (FLOP/AIE, size-to-tile ratios) are critical for generalizability of prediction models.
• RL agents, when informed by large-scale pretrained LLMs and accurate profiling, can explore and exploit non-obvious kernel optimizations not found in manual or heuristic methods (Su et al., 2 Dec 2025).
• Greedy and neighborhood RL search are highly efficient, yielding >24% cost savings while exploring minimal search fractions (Zhang et al., 2019).
• In outer-product codegen, explicit analysis of hardware pipeline constraints, register pressure, and memory hierarchy enables mapping that approaches the best human-tuned implementations (Veras et al., 2016).

Best-practice recommendations include restraining search-space complexity (limiting tiling depths, adhering to HW vector widths), precomputing valid divisors, batching measurement, exploiting asynchronous eval, and caching results to avoid redundant assessment (Zhang et al., 2019).

6. Communication-Optimal HGEMM and Sparse Kernels

For sparse-matrix multiplication (SpGEMM), the optimization problem generalizes to the minimization of communication volume under computational and memory-load balancing. Hypergraph models formalize this mapping:

• Vertices correspond to products and nonzeros.
• Nets (hyperedges) encode dependencies and communication requirements. Partitioning heuristics (e.g., PaToH, hMetis) minimize cut-net metrics subject to load-balance constraints, with fine-grained and coarsened models adapting to application needs. Empirical evidence shows choice of partition model impacts scalability and communication cost, with fine-grained approaches optimal for irregular sparsity (Ballard et al., 2016). This suggests algorithm design in HGEMM must adapt to data and architectural characteristics.

7. Experimental Benchmarks and Impact

Published frameworks have demonstrated strong performance relative to established libraries:

Approach	Throughput Gain	Energy Efficiency Gain	Fraction of Search Space Used
ACAP ML-DSE	1.23× geom. mean	1.25× geom. mean	~100% (6000 candidates, <2s eval)
CUDA-L2 RL+LLM	+11–29% vs cuBLAS/cuBLASLt	—	1000 configs, ~80% win rate
TVM G-BFS/N-A2C	-24% (G-BFS), -40% (N-A2C) runtime vs XGBoost/RNN	—	0.1%
Analytical Codegen	95% of peak vs. hand-tuned BLAS	—	Small enumerative search

Adaptable search mechanisms, integration with hardware profiling, and tight coupling of model prediction with empirical measurement are essential for scalable, portable, and high-performance HGEMM optimization (Papalamprou et al., 10 Nov 2025, Su et al., 2 Dec 2025, Zhang et al., 2019, Veras et al., 2016, Ballard et al., 2016).