MAGE: Leveraging LLMs for Automated Mapper Generation in Parallel Programming

Thank you for your thoughtful review. We appreciate it!

Does the decision procedure described in Figure 4a depend on the program to be optimized? In other words, how general is the decision procedure?

The decision procedure in Figure 4a is designed to be general and can be applied to any application program. It takes application-specific metadata, such as the names of tasks and regions, which are extracted from the application source code. Based on this metadata, it generates a mapper tailored to the specific application. This makes the procedure adaptable to various task-based programming applications while maintaining a generalized framework.

Why not just expose the parameters to a tuning framework? Can you provide a direct comparison?

We added an experiment comparing our approach with OpenTuner, where we let OpenTuner decide whether to sample from a set of keywords/parameters. The results shown in Figures 5 and 6 in the updated paper indicate that OpenTuner performs far worse than our agentic LLM-based solution. OpenTuner cannot leverage existing documentation of the DSL or have any prior knowledge of the problem of mapping, while our LLM framework can read the documentation and leverage textual feedback with error explanations and mapping suggestions. This textual feedback enables the LLM to iteratively refine the mapper, addressing issues like GPU memory errors dynamically. In contrast, OpenTuner cannot interpret or act on such textual feedback and often proposes invalid mappers. The ability to form prior knowledge over the decision space, process runtime errors in text form and adapt mappings accordingly makes our approach significantly better.

Why were only three benchmarks used in the evaluation, and how were they implemented?

Our paper evaluates a total of 9 benchmarks. Section 5.2 focuses on 3 scientific applications, and Section 5.3 includes 6 different well-known parallel matrix multiplication algorithms. The simplest benchmark is Stencil, which forms an optimization space of $2^{38}$ possibilities, let alone the most complex application, Pennant, which involves 34 tasks and 11 regions.

The six matrix multiply algorithms that we benchmark have mappers that are as involved as the applications --- parallel matrix multiply algorithms are not simple – based on our estimates, each parallel matrix multiplication algorithm has approximately $10^9$ possible choices for index mapping.

The three scientific applications are developed by domain experts, and six well-known matrix multiplication algorithms are implemented in prior work by experts in performance engineering. The MM algorithms have mappings as complex as the applications – the expert-written mappings for these algorithms are significantly more elaborate than those for scientific applications, requiring roughly twice as many lines of code as the scientific applications.

Are they CUDA applications? What CUDA driver and CPU were used in the evaluation? Please provide details of your evaluation platforms.

All nine benchmarks are CUDA applications. The tasks for scientific applications have implementations (i.e., variants) in CUDA for GPU, OpenMP, and CPU to enable the search for processor selection. The matrix-multiplication implementations target GPUs, where the search focuses on different ways to partition and distribute the tensors over multiple GPUs. We conducted experiments using CUDA driver version 12.2 on a platform with two Intel 10-core E5-2640 v4 CPUs, 256 GB of main memory, and four NVIDIA Tesla P100 GPUs.

Additionally, why were well-established GPU benchmark suites like Parboil and Rodinia not considered?

We chose to focus on the Legion parallel programming framework because it exposes a rich set of correctness-agnostic mapping decisions, enabling performance tuning without affecting application correctness. Other well-established benchmarks do not inherently expose similar abstractions for mapping. Our goal is to demonstrate the feasibility and potential of LLMs for performance tuning in task-based programming systems. We hope this work inspires researchers to extend these ideas to other benchmarks and programming systems in the future.

We hope we have addressed your concerns and questions thoroughly, and we are happy to answer more questions if they arise. Thank you very much!

As the author/reviewer discussion period getting close to an end (in 4 days -- Nov 26), we are wondering if 1) Our rebuttal addresses some of your concerns about the paper; 2) there is anything else we can do in the next 4 days to change your opinion on the current rating?

Thank you for your time in drafting the author's response. Can you automatically generate the DSL program from the target program? This was not clear from the paper, and it appears that this requires user intervention. If this is an automated process that can be applicable to general cases, why not test your approach on well-established GPU benchmarks?

I am not convinced by the argument that parallel matrix multiply algorithms are complex - having a large number of tuning parameter combinations does not necessarily increase the complexity of the program. For example, the number of possible loop-unroll factors for a simple matrix multiplication program can also be large.

I also don't think the comparison to OpenTuner is fair. The number of search iterations given to OpenTuner is too small. As an evolutionary-based algorithm, OpenTuner will require at least thousands of search iterations. Furthermore, I am not sure applying OpenTuner to the DSL is a good idea. Why not just expose the search parameters to OpenTuner, apply the OpenTuner's chosen parameters to the target program, measure the resulting performance and give the measurement as feedback to OpenTuner? While OpenTuner requires a larger number of search iterations, it does not rely on crafting a DSL program and does not require access to the computational resources to run the LLMs.

Given the concerns discussed above, my main reservations still stand and I will keep my original score. However, I do appreciate the author's effort in responding to my feedback and adding the new experimental results. Overall, I feel this is not the right problem for LLMs.

Thank you for your thoughtful review. We appreciate it!

The work does not introduce new techniques; incorporating an LLM into an existing reinforcement learning framework does not appear to be a novel contribution.

Thank you for engaging in a discussion with us. We respectfully disagree. Using LLM as optimizers (i.e., “generative optimization”) is a new field of study. Many discrete optimization problems have the potential to be solved by generative optimization (see applications in mixed integer programming [1]). This is the first paper to show that by using such a technique, we can learn implementations of mappers that surpass previous expert implementations – which has never been shown before. Our paper demonstrates the powerfulness of generative optimization and encourages future work in this area.

[1] Ali AhmadiTeshnizi, Wenzhi Gao, and Madeleine Udell. Optimus: Optimization modeling using mip solvers and large language models. arXiv preprint arXiv:2310.06116, 2023.

Evaluations are limited to simple benchmarks

Please kindly note that the benchmarks are not simple. We have utilized 9 real-world benchmarks in total. Section 5.2 focuses on 3 scientific applications, while Section 5.3 examines 6 well-known parallel matrix multiplication algorithms. Even the simplest scientific application, Stencil, is non-trivial: it comprises two tasks and twelve data arguments. Each task and data argument offers two placement options, and each data argument has four additional layout choices, leading to an optimization space of $2^{38}$ possibilities, let alone the most complex application, Pennant, which involves 34 tasks and 11 regions. The six matrix multiplication algorithms we benchmark are also intricate. Parallel matrix multiplication algorithms are inherently complex; our estimates indicate that each algorithm has approximately $10^9$ possible choices for index mapping. Moreover, the expert-written mappings for these algorithms are significantly more elaborate than those for scientific applications, requiring roughly twice as many lines of code as the scientific applications.

The correctness of task mappings is not discussed.

"Mappers do not change the correctness of an application's output; they only affect its performance." How do you verify that the LLM-generated mapper is correct? Are there any formal guarantees? Are communication operations (e.g., all_reduce) handled afterward?

The mapping is correct by construction: the mapper expressed in the DSL simply cannot give a mapping directive that results in a wrong execution result, because the underlying runtime ensures correct execution. For example, if two task calls t1(A) and then t2(A) are mapped to different processors of the machine, the runtime will automatically insert the copy of A from where t1 executed to where t2 will execute. All collective operations will be properly handled by the runtime.

Why is the grammar of the DSL structured as shown in Figure 3(a)?

The DSL can succinctly describe the search space (e.g., whether the decision is a per-task, per-region, or per-processor decision) in a syntactically reasonable way, enabling LLMs to generate interpretable programs rather than using discrete numerical numbers to represent each choice. This design simplifies reasoning and ensures code clarity. We are happy to add more explanations of the ideas behind the DSL design in the revised version.

Does the underlying hardware affect mapping strategies? For example, if A100 GPUs were used instead of P100s, would the RL pipeline need to be rerun?

Yes, mapping strategies are machine-dependent. Changing the hardware would require rerunning the agentic pipeline to adapt to the new environment. This context-specific nature of mapping underscores the importance of having a fully automated solution to avoid labor-intensive manual tuning.

Can your approach apply to distributed LLM training, and how do the generated task mapping strategies compare with Alpa?

The "index mapping" search in our method is analogous to the search performed in Alpa for data parallelism and operator parallelism, but our framework does not currently cover pipeline parallelism, which is one aspect of the search space in Alpa for deep learning.

Conversely, Alpa focuses on distributed deep learning workloads, where all computations are typically mapped to GPUs, and thus does not address processor kind selection or memory placement—both of which are essential in our search space for scientific applications. These additional dimensions, such as CPU vs. GPU selection and memory placement and layout optimization, are critical for workloads beyond deep learning.

Our current benchmarks focus on scientific applications and matrix multiplication algorithms. While we do not yet have benchmarks specifically for distributed LLM training, this is an exciting direction for future exploration.

We have added citations [A][C] to the paper (they are highlighted in red). Alpa [B] has already been cited and discussed in Section 3 (we applied red highlight to it as well).

Some early RL-based task mappers [C] are not discussed. Could you explain why including LLMs produces better results? Could your method handle the experiments presented in [C]?

Thank you for the suggestions! We cited it in the revised version. The intuition behind the advantage of LLM-based methods lies in their ability to read through existing documents on the problem domain to form a better prior over search space and can incorporate textual feedback such as error explanations and suggestions. Traditional RL techniques are much less sample efficient, often resulting in 100x or more search budgets. Our mapping search is general and could be applied to any computation graphs, including those similar to those ML workloads in [C], as part of future work.

We added an experiment comparing our approach with OpenTuner (traditional RL-based method), where we let OpenTuner decide whether to sample from a set of keywords/parameters. The results shown in Figures 5 and 6 in the updated paper indicate that OpenTuner performs far worse than our agentic LLM-based solution. OpenTuner cannot leverage existing documentation of the DSL or have any prior knowledge of the problem of mapping, while our LLM framework can read the documentation and leverage textual feedback with error explanations and mapping suggestions. This textual feedback enables the LLM to iteratively refine the mapper, addressing issues like GPU memory errors dynamically. In contrast, OpenTuner cannot interpret or act on such textual feedback and often proposes invalid mappers. The ability to form prior knowledge over the decision space, process runtime errors in text form and adapt mappings accordingly makes our approach significantly better.

We hope we have addressed your concerns and questions thoroughly, and we are happy to answer more questions if they arise. Thank you very much!

Using a template with fill-in-the-blank prompts in C++; DSL’s effectiveness or necessity

Our goal is not to prove that LLMs cannot generate low-level C++ code correctly, but to highlight that the DSL significantly simplifies the task so that we can let LLMs further explore the search space beyond just generating code given a specified strategy.

We added a new experiment using fill-in-the-blank prompts in C++. This approach achieved a 20% success rate (passing the tests for the two memory placement problems out of the ten mapping problems), better than earlier C++ baselines but still far behind the DSL-based generation.

While giving more C++ examples can help LLM with the generation, there is a token limit and context window constraint. In general, code generation for real-world software is known to be a challenging and unsolved problem for LLMs [2,3].

[2] Jimenez et. al. SWE-Bench: Can Language Models Resolve Real-World Github Issues?

[3] Du et. al. Classeval: A Manually-Crafted Benchmark for Evaluating LLMs on Class-Level Code Generation.

To illustrate the effectiveness of DSL, we added Table A1 (lines of code comparison between DSL and C++) to the paper’s appendix. As shown below, for the 9 benchmarks used in our evaluation, the average C++ mapper is 406 lines long, while the average DSL mapper is just 29 lines, a 14x reduction. The DSL abstracts low-level details, making it much easier for LLMs to generate and reason about compared to the complexity of writing hundreds of C++ lines with intricate APIs.

We also updated Figure 3 to better illustrate the code complexity between DSL and C++.

Can the two errors shown in Table 1 for your DSL be avoided by improving the prompt?

Yes, both errors were syntax-related. We improved our prompt by explicitly emphasizing DSL syntax rules and common pitfalls. With these enhancements, the LLM successfully generates correct DSL mappers on the first trial.

Is your DSL general enough to apply parallel strategies to other frameworks, like TaskFlow [A]? How much effort is required to support frameworks other than Legion? Would your DSL need adjustments?

Many parallel programming frameworks, including TaskFlow, StarPU, Chapel, X10, and Charm++, designed similar mapping controls. Our DSL covers features like data layout selection, which are not universally supported (e.g., TaskFlow lacks this feature). Adapting the DSL for other frameworks would depend on the number of correctness-agnostic choices exposed by each framework. Adjustments would primarily involve aligning the DSL with the specific controls and abstractions of the target framework. But we believe the final result would be similar — Legion has the most elaborate mapping API, and other systems would have fewer choices to search.

How many GPU nodes were used?

We ran all experiments on one node with 4 GPUs.

As the author/reviewer discussion period getting close to an end (in 4 days -- Nov 26), we are wondering if 1) Our rebuttal addresses some of your concerns about the paper; 2) there is anything else we can do in the next 4 days to change your opinion on the current rating?

Thanks for the thoughtful review! We have uploaded a new rebuttal version. We address your concerns point by point below:

IIUC, the parallel programming is at task/processors level instead of kernel… The use of "RL" tends to lead people to think of gradient-descent based RL algo instead of agentic based RL flow… We have changed our writing in the paper. Please see the updated paper draft where we highlighted the changes in red. Ablation study of feedback: add few-shot prompting before adopting an agentic flow

Thanks for suggesting a new experiment! We conducted the experiments with zero-shot and few-shot prompting (without feedback). The performance of these approaches is much worse than our agent-based solution. We have included these results in Figure 7 of our revised version to highlight the importance of our feedback-driven solution.

Is it possible for LLM to discover new GEMM algorithms?

Thank you for the question! Our results demonstrate that LLMs can find better implementations of existing GEMM algorithms. While the algorithms themselves are typically published as pseudo-code in papers, implementing them on distributed machines involves many decisions not well-specified in the pseudo-code. Our LLM-based solution improves the expert implementation of the existing GEMM algorithms.

Index mapping is basically "sharding" in distributed ML frameworks (e.g., Tensor Parallelisms). Is it possible to generalize the ideas to automate sharding in a distributed system?

Correct. Our "index mapping" is analogous to "sharding" in ML frameworks. While this paper focuses on scientific workloads, the idea of using LLMs to decide sharding is indeed an exciting direction to explore further, especially for distributed machine learning workloads.

We hope we have addressed your concerns and questions thoroughly, and we are happy to answer more questions if they arise. Thank you very much!

Why do error explanations significantly impact final throughput results?

Error explanations help LLMs interpret and fix mistakes more effectively in subsequent iterations. Raw error messages alone (without explanation) lead to repeated mistakes because the raw error messages can be hard to interpret (e.g., LLMs cannot understand what this raw error message mean: “Assertion failed: stride does not match expected value.”, but the error explanation “Memory layout is unexpected” is more informative). We added more examples of enhanced feedback in our revised version. We show different kinds of feedback below, with more examples in the appendix (Table A2) of our updated paper.

We hope we have addressed your concerns and questions thoroughly, and we are happy to answer more questions if they arise. Thank you very much!

Thank you for your thoughtful review! We uploaded a revised version of the paper. We address your concerns point by point as follows:

Limited Experimental Scope, only 3 applications

Our paper evaluates a total of 9 benchmarks. Section 5.2 focuses on 3 real-world scientific applications, while Section 5.3 includes 6 different well-known parallel matrix multiplication algorithms. The six matrix multiply algorithms that we benchmark have mappers that are as involved as the applications --- parallel matrix multiply algorithms are not simple. Based on our estimates, each parallel matrix multiplication algorithm has approximately 10^9 possible choices for index mapping.

Insufficient Discussion on Mapping Strategy Selection in Section 4.2.

Decisions such as "task selection" or "memory allocation," there is no exploration of how these decisions may affect final performance

The performance impact of different mapping strategies is discussed in Section 3 and Section 4.1, along with the design of our DSL.

For task selection, we are selecting which processor (GPU or CPU) to process the task. This is determined by factors such as GPU memory and kernel execution time. It is a per-task decision. This choice depends on factors such as task size, GPU memory, and kernel launch latency. For example, tiny tasks that require very little computation may prefer to run on CPUs due to the GPU kernel launch overhead, whereas tasks with large memory footprints might prefer to be assigned to OpenMP or CPU when GPU memory is insufficient.

For memory placement, data can be placed in GPU’s FrameBuffer for faster access, or to ZeroCopy memory for shared access between CPU and GPU (to avoid data communication cost between CPU and GPU), or in CPU system memory for large data (because GPU memory is usually more limited). Each choice introduces a trade-off between memory access speed, memory usage, and transfer overhead.

Lack of Discussion on Failure Cases

Apologies for not emphasizing this sufficiently. We analyze failure cases in Section 5.1, particularly focusing on compilation issues and test failures. In our revised version, we highlight this discussion more clearly and structure it better. Please see the revised version in detail. Here is the summarized failure case analysis:

Compilation errors in C++ are caused by LLM’s failure to understand contextual dependencies in real-world software such as hallucinating non-existent variables. When code compiles but fails the test case, this happens because LLMs fail to understand how to coordinate different APIs together to achieve certain functionality. A detailed example is provided in the paper. For DSL, the failure case stems from incorrect syntax.

How does the DSL specifically address the complexities of C++ APIs and low-level details?

We included compiler implementation details in the appendix (Sections A2, A3, A4) and provided the complete compiler code in the supplementary material.

With such a high level of abstraction, will using the DSL result in the loss of low-level C++ details?

DSL is more restricted than what can be done at the C++ level, but it still can express better mappers than the hand-written C++ designed by experts. In our experiment results, all the mappers that are found to be better than expert-written C++ mappers are expressible in the DSL.

In Figure 7(c), the second best mode line and the best mode line exhibit alternating performance increases over 10 iterations. Is there a need for more iterations to ensure Full Feedback Mode consistently outperforms others?

Thank you for pointing this out! We extended experiments to run up to 15 iterations for Figure 7(c). See the updated Figure 7(c) in the revised version of the paper. The result showed after 10 iterations, the performance of two modes converged.

We have expected this result. We highlight that the Full Feedback Mode is most useful to provide efficient learning signals early on when errors occur, by suggesting how the LLMs can correct the execution error. It does not provide useful suggestions on how to improve the performance of the mapper. We show this phenomenon in both Figure 7(a) and 7(b), where the Full feedback mode (System + Explain + Suggest) gained higher normalized throughput much faster than other modes.

As we approach the end of the discussion period, we were wondering if there are any remaining concerns or questions about our paper that we could address?