Improving Parallel Program Performance with LLM Optimizers via Agent-System Interfaces

Thank you for the thoughtful and constructive review.

Q1: How much domain expertise (about the benchmark domain and the LLM capabilities) is needed to expand the ASI and the prompt for the benchmark? More practically, do we need experts who write the mappers to design DSL encoding appropriate priors?

Our framework is designed to minimize the need for domain or system-level expertise when adding new benchmarks. Specifically, for all benchmarks in our evaluation, the agent begins from a randomly initialized mapper and improves it through generative optimization, without any handcrafted priors or expert-written strategies. To incorporate a new application into the benchmark, only two inputs are needed:

1. Application metadata, including task names and the list of data arguments each task accesses

2. Hardware specification, including the number of CPUs and GPUs per node and the total number of nodes.

These inputs are typically available from the application code and the machine, requiring no understanding of the DSL or the runtime system. As a result, application developers do not need to write or understand mappers, provide mapping hints, or possess any knowledge of the DSL.

That said, our system can also support optional injection of expert knowledge via the AutoGuide module. This module can provide customized interpretations of execution failures or application-specific heuristics (e.g., mapping large tasks to GPUs if the developer already knows which tasks are large). While our experiments do not use any such guidance, we believe this optional flexibility is valuable in practice, especially for developers who already have insight into performance bottlenecks and want to speed up the optimization process.

Q2: The authors can motivate the choice of benchmarks better, particularly for general readers unfamiliar with the domain knowledge of these tasks.

Thanks for the suggestion! We will improve the paper to provide more context and motivation for our benchmark selection. Among the 9 benchmarks, 6 are well-known matrix multiplication algorithms (Cannon, SUMMA, PUMMA, Johnson’s, Solomonik’s, and COSMA). Parallel matrix multiplication remains an active research topic due to its central role in high-performance computing and scientific simulations [1]. Furthermore, improving matrix multiplication performance has a broad impact, as it accelerates numerous downstream machine learning workloads [2,3]. The remaining 3 applications (Circuit, Stencil, and Pennant) represent diverse scientific computing workloads beyond matrix multiplication. Together, this benchmark suite offers both depth (through representative matrix-multiplication algorithms) and breadth (through diverse HPC workloads). We will expand the background and motivation accordingly in the revision.

References:

[1] Yadav et al. “DISTAL: The Distributed Tensor Algebra Compiler”. 2022

[2] Jangda et al. “Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads”. 2022

[3] Zheng et al. “TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives”. 2025

Thank you for the thoughtful and constructive feedback.

Q1: What is your DSL grammar?

In the revision, we will include a complete description of the DSL syntax in the Appendix, covering its constructs for task placement, memory allocation, layout specification, and index mapping, as shown below.

Terminals:
  TaskName, RegionName, var, int

Grammar Rules:
  Program         → Statement+
  Statement       → TaskMap | DataMap | DataLayout | FuncDef | IndexTaskMap TaskName var
  TaskMap         → Task TaskName Proc+
  DataMap         → Region TaskName RegionName Proc Memory+
  Proc            → CPU | GPU | OMP
  Memory          → SYSMEM | FBMEM | ZCMEM | SOCKMEM
  DataLayout      → Layout TaskName RegionName Proc Constraint+
  Constraint      → SOA | AOS | C_order | F_order | Align == int
  FuncDef         → def var(var+): FuncStmt+
  FuncStmt        → var = Expr | return Expr
  Expr            → var | var(Expr+) | Machine(Proc) | Expr.Expr | Expr Op Expr |
                    (Expr) | Expr[Expr] | *Expr | Expr ? Expr : Expr

Q2: Why only report the best point for the proposed approach rather than full optimization curve?

Please kindly note that we report both the best and average optimization trajectory over 10 iterations across 5 runs in Figure 4. Also, we want to clarify that reporting the best-performing mapper is appropriate in our context. Mapper optimization is an offline process, and in practice, it is standard to run the optimizer multiple times and deploy the best result. Once identified, the mapper can be reused across repeated executions on the same application, input, and hardware, incurring no further search cost.

Q3: Is the DSL expressive enough to cover all mapper code generation problems?

Our DSL is designed to express a wide range of high-performance mapping strategies, including all of the most important decisions. While there may be cases where certain optimizations are not directly expressible, we have not encountered any. Despite being more constrained than general-purpose C++, the DSL has been proven to be effective: all mappers discovered by our agent that outperform expert-written C++ implementations are expressible within the current DSL.

Q4: Server specifications and application information are mentioned as inputs, but details are missing.

We will clarify these details in the revised version. The server specification includes the number of GPUs and CPUs, and whether the OpenMP runtime is enabled. The application specification includes the list of tasks defined in the application and the data regions accessed by each task. These inputs define the structured search space explored by the agent during optimization. We will also include an example of such input specifications in the revised version for completeness.

Q5: The authors present Code Generation Success Rates (Section 5.2). However, in the main evaluation (Section 5.1), only the performance results are reported.

In the main evaluation (Section 5.1), we focus on measuring end-to-end performance, which includes both the correctness and performance of generated mappers. If the generated code has any syntax or runtime issues, its throughput is recorded as 0. Therefore, the performance numbers in Section 5.1 implicitly reflect code generation success, i.e., incorrect mappers yield zero performance.

Section 5.2 isolates the code generation aspect to better analyze the effects of the DSL on LLM generation success. This section does not aim to evaluate performance directly but rather investigates how often the LLM produces syntactically and semantically correct mappers given natural language descriptions. It complements the main results by demonstrating that using the DSL significantly improves generation success compared to C++, which underpins the performance improvements seen in Section 5.1.

Q6: Are all random mappers correct?

No, not all random mappers are correct. For each application, we generate 10 random mappers by sampling from the full DSL-defined search space, totaling 90 mappers across 9 applications. Among them, 74 (82.2%) raise runtime errors due to invalid mapping decisions. The runtime system enforces correctness by rejecting such mappers during execution, resulting in a throughput of zero.

Thank you for your thoughtful review! We address the concerns below:

Q1: Performance metrics could be extended

We appreciate your suggestion and welcome the opportunity to clarify our evaluation methodology and expand the reported statistics.

In our setting, reporting the best result across multiple runs is appropriate, as the best mapper is the one that is desired by the user. Mapper search is an offline optimization process, and it is feasible to run the optimizer multiple times to select the highest-performing mapper. Once identified, this mapper can be reused without incurring additional search cost, as the deployment scenario (application, input, and hardware) remains fixed.

That said, we agree that additional statistics provide a more complete picture of performance variations. In the revised version, we will include the mean, standard deviation, worst, median, and best normalized throughput across five runs for each benchmark. The extended results are as follows:

Our Framework (Normalized Throughput)

We additionally report OpenTuner results.

OpenTuner (Normalized Throughput)

Our method achieves relatively stable performance across most benchmarks. The higher variance and occasional 0.00× worst-case throughput observed in SUMMA, PUMMA, and Solomonik are due to invalid mapper configurations in the search space (e.g., violating cuBLAS layout constraints). The runtime enforces correctness by rejecting such configurations during execution. While the generative optimizer typically learns to avoid these cases through the AutoGuide mechanism, occasional failures within the 10-iteration budget are still possible. In practice, such failures can be mitigated by repeating the optimization and selecting the best-performing mapper.

In contrast, OpenTuner, despite running the same number of iterations, fails to generate valid mappers for 8 out of 9 benchmarks. This highlights the difficulty of exploring the search space using traditional reinforcement learning methods.

Q2: Experiments on multi-node systems

Thank you for raising this point. Our system supports multi-node execution, as this capability is already provided by the underlying runtime system. There is no fundamental technical limitation in our approach that prevents generalization to multi-node systems. The only technical issue we encountered in scaling up our experiments is an engineering issue, not a limitation of our method. Specifically, the cluster we currently use permits interactive execution on a single node, but multi-node runs require submitting jobs through the SLURM scheduler and waiting for resource allocation and job execution. While this can be addressed through SLURM-specific customization, it is orthogonal to the core contributions and novelty of our system. For users with direct multi-node access (e.g., without queueing), our system runs seamlessly without modification. We plan to extend our evaluation to larger-scale experiments as part of future work.