Investigate improving JIT heuristics with machine learning

[AndyAyersMS]

Overview

There is now fairly substantial literature on improving compiler optimization by leveraging machine learning (ML). A comprehensive survey compiled by Zheng Wang can be found here: https://github.com/zwang4/awesome-machine-learning-in-compilers.

Here's a closely relevant paper and github repo: MLGO: a Machine Learning Guided Compiler
Optimizations Framework. Like our efforts below it leverages a Policy Gradient algorithm (reinforcement Learning).

• https://github.com/google/ml-compiler-opt

Several years ago we attempted to leverage ML to create better inline heuristics. That experiment was largely a failure, at least as far as using ML to predict if an inline would improve performance. But one of the key speculations from that work was that lack of PGO was to blame. Now that we have PGO, it is a good time to revisit this area to see if it PGO was indeed the key missing ingredient.

Proposed Work

During the .NET 9 product cycle, we plan to investigate applying ML techniques to heuristics used by the JIT. The items below represent the initial steps of this investigation. This list will change and grow as the work progresses.

We also want to tackle a relatively simple problem, at least initially. Thus our initial effort will be to try and refine and improve the heuristic the JIT uses for Common Subexpression Elimination (aka CSE):

• Study the current CSE heuristic and try and understand what it is modelling. Identify apparent shortcomings and limitations.
  • See initial notes below.
• Develop tooling for varying which CSEs the JIT can perform, either in concert with the current heuristic or independent of it. Measure the impact variation of this on key JIT metrics (performance, code size, jit time, etc).
  • JIT: Allow experimenting with different CSE subsets #92918
  • Tool for exploring performance by varying JIT behavior jitutils#381
  • Update for new categories of tiered compilation jitutils#383
• See how well PerfScores can be used as a proxy for measuring actual performance. Ideally, we can leverage PerfScores to avoid needing to run benchmarks repeatedly. But if not, perhaps we can still rely on PerfScores to limit or prioritize the runs needed.
  • See notes below for some initial data.
• Try a number of different ML techniques for improving CSE heuristics, both "small step" (evaluating one CSE at a time) and "big step" (evaluate an entire set of CSEs). Identify the key predictive observations. Try both interpretable and black box models.
  • I have been working on a Policy Gradient approach. The implementation is split with part of it in the JIT and the rest in a driver program. Documentation is lacking, so ping me if you want to try this at home.
    • JIT part was introduced in JIT: initial support for reinforcement learning of CSE heuristic #96880
    • Driver part in Initial version of a tool for ML on CSEs jitutils#389
    • More work on stopping: JIT: stopping preference for ML CSE heuristic #98063
    • More work on the driver: Fix RLCSE for new metrics format jitutils#391
    • More features JIT: More CSE heuristics adjustments #98257
    • More work on the driver: RLCSE: fix MCMC and GatherFeatures, overwrite dumps jitutils#395
    • Streaming SPMI mode JIT: Streaming mode for SPMI #98440
    • And client updates to use that mode: RLCSE: streaming SPMI mode jitutils#397
    • Enable running in release with fixed parameter set: JIT: refactor CSE to allow running greedy ML heuristic in release #98729
    • actually try enabling (draft): JIT: temporarily enable RLCSEGreedy to see how it fares in CI #98776
• See how feasible it is to have a single architecture-agnostic heuristic (perhaps parameterized like the current heuristic with register info) or whether different heuristics are warranted for say arm64 and x64.
• Develop a common ML "infrastructure" that we can apply to other similar heuristics in the JIT.
• Set up a CI job to run the Policy Gradient
  • JIT: test the new cse policy in jit-experimental #98777 (running the policy, not the training)
• Set up a Perf lab experiment to run the current "best" greedy policy on benchmarks. Tis should be running now, will check on the data in a few days
  • Enable RLCSE experiment in the perf lab #98779
  • Add setup for RLCSE experiment performance#3975
    Also worth noting is Lee Culver's work to adapt the JIT CSE problem into a more standard gymnasium setting: JIT CSE Optimization - Add a gymnasium environment for reinforcement learning #101856

---

Update May 2024.

Given the work above we have been able to produce heuristics that can improve the aggregate perf score for methods via CSE, with about a 0.4% geomean improvement across all methods with CSE candidates. So far those results haven't translated into widespread improvements in our perf lab data. Exactly why is unclear ... data below would suggest one or more of the following:

• poor correlation of perf score with actual perf
• few benchmarks whose scores critically depend on CSEs, or
• those benchmarks whose perf does critically depend on CSEs are all easy/obvious cases
• difficulty extracting small signal from noise (still lots of noisy benchmark results)

The training evaluations show that the best ML heuristic is obtaining about half of what could be gotten from an optimal heuristic. Some recent results:

Best/base: 0.9913 [optimizing for  score]
vs Base    0.9957 Better 440 Same 1476 Worse 84
vs Best    1.0044 Better 0 Same 1607 Worse 393

Params     0.6093, 0.7750,-0.5701,-0.6827, 0.5060,-0.0514,-1.3563, 0.4515,-2.0999, 0.0000,-1.0623,-1.3723, 0.0000,-0.6143,-0.8286,-0.2956,-1.1433, 0.3418, 1.3964,-0.0043,-0.4237, 0.6097,-1.9773,-0.3684, 0.7246

Collecting greedy policy data via SPMI... done (27583 ms)
Greedy/Base (score): 34965 methods, 8375 better, 24792 same, 1797 worse,  0.9960 geomean
Best:   76679 @  0.7731
Worst:  82000 @  1.4512

Greedy/Base (size): 34965 methods, 4628 better, 24694 same, 5642 worse,  1.0005 geomean
Best:   15489 @  0.7000
Worst:  82000 @  1.4205

 
We don't have any more active work planned on machine learning of heuristics for .NET 9. However, we expect to revisit this area as .NET 9 work winds down, so stay tuned for further development in the coming months.

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch
See info in area-owners.md if you want to be subscribed.

Issue Details

---

Overview

There is now fairly substantial literature on improving compiler optimization by leveraging machine learning (ML). A comprehensive survey compiled by Zheng Wang can be found here: https://github.com/zwang4/awesome-machine-learning-in-compilers.

Several years ago we attempted to leverage ML to create better inline heuristics. That experiment was largely a failure, at least as far as using ML to predict if an inline would improve performance. But one of the key speculations from that work was that lack of PGO was to blame. Now that we have PGO, it is a good time to revisit this area to see if it PGO was indeed the key missing ingredient.

Proposed Work

During the .NET 9 product cycle, we plan to investigate applying ML techniques to heuristics used by the JIT. The items below represent the initial steps of this investigation. This list will change and grow as the work progresses.

We also want to tackle a relatively simple problem, at least initially. Thus our initial effort will be to try and refine and improve the heuristic the JIT uses for Common Subexpression Elimination (aka CSE):

• Study the current CSE heuristic and try and understand what it is modelling. Identify apparent shortcomings and limitations.
• Develop tooling for varying which CSEs the JIT can perform, either in concert with the current heuristic or independent of it. Measure the impact variation of this on key JIT metrics (performance, code size, jit time, etc).
• See how well PerfScores can be used as a proxy for measuring actual performance. Ideally, we can leverage PerfScores to avoid needing to run benchmarks repeatedly. But if not, perhaps we can still rely on PerfScores to limit or prioritize the runs needed.
• Try a number of different ML techniques for improving CSE heuristics, both "small step" (evaluating one CSE at a time) and "big step" (evaluate an entire set of CSEs). Identify the key predictive observations. Try both interpretable and black box models.
• See how feasible it is to have a single architecture-agnostic heuristic (perhaps parameterized like the current heuristic with register info) or whether different heuristics are warranted for say arm64 and x64.
• Develop a common ML "infrastructure" that we can apply to other similar heuristics in the JIT.

---

cc @dotnet/jit-contrib

Author:	AndyAyersMS
Assignees:	AndyAyersMS
Labels:	area-CodeGen-coreclr
Milestone:	9.0.0

[AndyAyersMS]

Initial Notes on the Current JIT CSE Heuristic

1. Does the initial frame size computation hold up well? Can we justify the 3/2/1 formula in CSEHeuristic::Initialize?
   a. Disable CSE, see how well number of tracked locals in registers maps up with the prediction. If not, see if there is some way to make a better prediction.
   b. Then slowly enable CSE, and see if this continues to hold up well.
   c. See how well frame size matches estimates; if not good, see if we can improve this somehow

2. Do the aggressive/moderate thresholds make sense? That is, should they be higher or lower? Based on something else?
   a. trickier to see how to model this; perhaps keep processing CSEs and see when spilling starts?
   (that is, do a local sensitivity analysis) -- might be counterbalanced by the promotion code.
   b. how about the overall approach of setting these thresholds based on IR observations? Seems sensible but maybe there are other ideas to try?
   c. at the end of CSEHeuristic::Initialize we increase the thresholds if they seem to low. Why?
   d. also doesn't that increase mix up weighted and unweighted cases? How does it make sense to compare unweighted ref counts with BB_UNITY_WEIGHT?
   e. should we try some “true” register pressure estimate, or is that hopeless? We also don’t know how the pressure points intersect with potential CSE temp liveness.

3. Is the assumption here that FP CSEs are rare true?

4. How reliable are CostSz and CostEx?
   a. probably a can of worms... how would we even go about validating these?

5. Should candidates be ranked by local benefit, or by global benefit, or some ratio of benefit / cost?
   a. That is, why doesn’t Compiler::optCSEcostCmpEx consider the csdUseWtCnt more prominently? Seems like the right figure of merit (when optimizing for speed is csdUseWtCnt * tree->CostEx()
   so a relatively “cheap” CSE that heavily used is preferable to a “costly” CSE that is not as heavily used.
   b. Should this factor in code size even in speed modes? Assuming CSEs generally decrease size a bit, we might want to bias the above towards the heavier ones slightly.
   c. Do the threshold increases in PerformCSE make sense? This gives a budget-like aspect to the ranking, so the order does matter. What if we ignored this?

6. Are we missing candidates?
   a. review bail outs in Compiler::optIsCSEcandidate.
   b. review limitations in Compiler::optValnumCSE_Locate

7. In ConsiderCandidates is the computation sensible?
   We have 3 models x 2 modes ({aggressive/moderate/conservative} x {size, speed}).
   Plus some stress modes (including random...)
   a. seems like we mix weighted/unweighted costs up here at times. eg extra_no_cost is unweighed (size) based, but we add it to no_cse_cost, which can be weighted (speed) based.
   b. various magic numbers and twisty decision logic
   c. various per-arch costs, some of them look dated

8. Is the "shared constant" CSE code sensible? Here we disregard the low bits of constants and at use
   points we rematerialize the missing bits via an add.
   a. So use cost is a bit higher... but this does not seem to be modelled by the heuristics ?
   b. Is the "best base" value computation reasonable? Seems like we're trying to use small immediates here?

[AndyAyersMS]

Perf vs Perf Score

Here's a chart of about 3000 measurements across 50 or so benchmarks, varying the set of CSEs we do in the hottest method. Each data point is the ratio of the perf vs the default JIT perf, vs the ratio of the perf score vs the default JIT perf score. Thus values above 1.0 on the y-axis are benchmarks that ran more slowly, and values to the right of 1.0 on the x-axis are benchmarks that were predicted to run more slowly.

We see most data in the "slower and predicted slower" (upper right) quadrant because CSEs generally improve perf, so doing less of them than we'd normally do will likely slow down the benchmark.

Linear fit has a correlation coefficient of about 0.3, so there is a fair amount of unexplained variation in perf. Also I have excluded one benchmark's data from this as it seems anomalous.

Ideally, we'd be able to get a better correlation out of perf scores, since these can be produced very quickly as we vary JIT behavior, on the order of 1000 experiments per second. If we actually run benchmarks then it takes around 30 seconds for one experiment, so there is a massive benefit to leveraging perf scores.

The main contributors to poor perf predictability are thought to be (in no particular order):

• errors in perf score per-block computations: wrong latency for instructions, lack of modelling dependence, other microarchitectural effects. Some of this we may be able to fix by invoking tools like llvm-mca to either give us accurate data or help us spot errors in our modelling. And that in turn requires making the JIT disasm be something readily parseable by standard tools.
• errors in block weights. Even though we have PGO data we may no longer have accurate block weights at the end of compilation. This is something we can improve, though fundamentally the JIT will be unable to correctly anticipate how weights may change because of optimizations, so some amount of error here is inevitable.
• "noisy" benchmarks. We know many benchmarks show considerable within-run or cross-run variation. Some of these sources of noise include: intel JCC errata, alignment sensitivity of data, other subtle microarchitectural effects (say branch collision)
• errors in PGO data. Note there is some deliberate randomness and some inherent randomness in PGO data. There may also be errors in the instrumentation, collection, or reconstruction processes. Also included here are potential phased behaviors by benchmarks (where PGO data does not reflect the actual benchmark block weights because the data was measured too early)

So the task here is to dig into these issues and see if and how perf scores can be improved.

Also I need to look into why MaximizeSchwarzCriterion has anomalous results; including data from this benchmark drops R^2 to essentially zero.

[ShreyasJejurkar]

This is a great initiative, just wondering do other popular (JIT-based) programming languages (e.g. - Java) have something like this in place? Would love to read their outcome. :)

[AndyAyersMS]

This is a great initiative, just wondering do other popular (JIT-based) programming languages (e.g. - Java) have something like this in place? Would love to read their outcome. :)

I think I may have read some papers by the Graal folks where they leveraged ML, but I don't have the references handy. Presumably they'd be somewhere in the big list of papers I referred to above.

[AndyAyersMS]

Have been digging into why the PerfScores have such poor correlation on MaximizeSchwarzCriterion and it seems like it is a combination of a few effects:

• Scalable profiles are notably approximate and will vary from run to run
• The exact timing of tiering up can vary, so the profile data will also vary from run to run
• Block weight reconstruction is ill-conditioned (small input changes can lead to large output changes)
• Methods that rely on OSR can end up with partial profiles. In particular for Tier0-instr methods that are called and never return we often do not have profile weight for the entry block since that is deduced from the returns.

Thus the block weights (in particular the large weights for innermost loop blocks) can vary quite a bit from one run to the next, and that seems to be the major factor. You can see this by just running the benchmark multiple times with the same jit configuration:

Benchmark,Index,Mask,NumCse,CodeSize,PerfScore,PerfScoreRatio,Perf,PerfRatio
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,0,ffffffff,17,1643,649125874.85,1.000,59311.6714,1.000
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,1,00000000,0,1533,80505239.17,0.124,61848.4732,1.043
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,2,00000001,1,1532,421090857.67,0.649,65608.5883,1.106
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,3,00000002,1,1528,1307172767.70,2.014,59236.0308,0.999
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,4,00000004,1,1584,108418141.52,0.167,65119.4767,1.098

Benchmark,Index,Mask,NumCse,CodeSize,PerfScore,PerfScoreRatio,Perf,PerfRatio
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,0,ffffffff,17,1643,100478799.14,1.000,61020.7827,1.000
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,1,00000000,0,1533,69888793.36,0.696,63267.6950,1.037
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,2,00000001,1,1530,69621450.93,0.693,66241.6050,1.086
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,3,00000002,1,1528,108183513.60,1.077,58789.5833,0.963
JetStream.TimeSeriesSegmentation.MaximizeSchwarzCriterion,4,00000004,1,1582,70447641.30,0.701,64906.9033,1.064

In particular the first line of each set above shows two "identical" runs with default JIT behavior, the first took 59,311us with PerfScore 649,125,874.85 and the second 61,020us with PerfScore 100,478,799.14. So perf varied by about 1.03x and PerfScore about 6x.

For the purposes of this investigation, we can just exclude methods that have OSR methods.

This variation in weights is possibly less of a problem in practice though it may be a contributor to the kind of fluctuating performance we noted in #87324. We should think about better strategies to counteract it:

• more aggressive blending? We currently use 0.99 for blend factor, perhaps something more like 0.90?
• Running a Tier-1 instr stage before we go to full on Tier-1, for methods that OSR'd, so that the root level blocks get a more accurate weight?
• adding extra probes based on patchpoint location that can recover missing counts. The interesting case is the one where the patchpoint helper transitions to the OSR method—the helper would need to increment the counter. And since this probe would not normally appear in the schema we'd have to figure out how to fit it in without disrupting anything else. We do something like this for throws already, but those are easy to anticipate.

[MichalPetryka]

• Running a Tier-1 instr stage before we go to full on Tier-1, for methods that OSR'd, so that the root level blocks get a more accurate weight?

Isn't this already the case for R2Red code?

[AndyAyersMS]

• Running a Tier-1 instr stage before we go to full on Tier-1, for methods that OSR'd, so that the root level blocks get a more accurate weight?

Isn't this already the case for R2Red code?

Yes -- the idea there was to avoid needing to go from R2R (optimized) -> Tier0 (unoptimized) to gather PGO data. Here the idea would be to go from Tier0+instr -> (OSR) -> Tier1+ instr -> Tier1, for methods that are called infrequently (but more than once) and contain hot loops. The second instrumented run is there to fill out the missing bits of profile in the parts of the method that currently only run under OSR.