YJIT: Profile and report memory used by YJIT Rust

[k0kubun]

Current memory usage

On railsbench, with M1 Mac and stats build of ruby@0949cd7, yjit_alloc_size is 9.0MB (code_region_size is 6.9MB).

Heap profiling result

Using yjit-dhat branch, the profiling result was:

• 3.1MB (33.9%): Block by gen_single_block
• 2.2MB (23.9%): Branch by make_branch_entry
• 1.9MB (21.5%): Vec by get_or_create_version_list, etc.
• 0.7MB (7.6%): BranchStub by set_branch_target
• 0.5MB (5.8%): HashMap by assume_method_lookup_stable, etc.
• 0.2MB (2.3%): IseqPayload by get_or_create_iseq_payload
• 0.2MB (2.1%): BranchTarget by set_branch_target

Struct counts

Using yjit-mem-stats branch, the breakdown of structs was:

# of IseqPayload: 2190
# of Block: 15296
# of Branch: 26968
# of BranchTarget: 28161
# of BranchTarget::Stub: 11836
# of BranchTarget::Block: 16325

size of IseqPayload: 96
size of Block: 176
size of Branch: 56
size of BranchTarget: 16
size of BranchStub: 58
size of Context: 38

Note that #s might be missing some unreachable structs on the heap, but then it might also mean they are currently leaking.

Known ideas

The following ideas are ordered based on the estimated impact (the earlier, the more impactful).

1. Eliminate extra blocks/branches created by defer_compilation

Related: #462

defer_count was 6883. If we just replace the previous block with a combined block, we could reduce Blocks and Branches by 6883, which would save more memory than any other idea.

2. Hash consing for Context

The idea was proposed by Maxime.

Context could be a series of "cons" where each cons is a pair of a diff and a pointer to the previous Context, which roots from the default Context. If you replace Context (38 bytes) in Block and BranchStub with a Box of cons (8 bytes), it would save 30 bytes from 15296 blocks and 11836 branch stubs, which could give 0.81MB reduction at most (9% of all).

Background

I already tried deduplicating contexts, but the hash table used for deduplication ended up being large, even with shrink_to_fit. The reduction was too limited despite the complexity.

However, thanks to the patch, I was able to easily collect the distribution of duplicated contexts https://gist.github.com/k0kubun/c1a5854e011db90d7e6efd73a6eda3b3. Many duplicated contexts share the same local_types, temp_types, and temp_mapping, so we could save those 32 bytes if we optimize the data structure for it.

3. Replace CodePtr with write_pos

If we allow only up to 4GB in --yjit-exec-mem-size, Option<CodePtr> could be changed to Option<NonZeroU32> by storing a write_pos instead of a pointer. (well, write_pos needs to be interpreted as a 1-origin value then 😅)

4. Combine targets and shape in Branch

In Branch, we could combine targets and shape into a single enum. Each variant should have either Box<[BranchTarget; 2]> or Box<BranchTarget>. Not a lot of reduction, but it’d still save 8 byte for every case.

Notes

For each Vec, we might be able to save some space by eliminating the capacity if it's no longer modified, e.g using into_boxed_slice.

---

cc: @maximecb @XrXr

[k0kubun]

An experiment for the second point ruby#7279

[k0kubun]

I tried deduplicating Context with HashMap<Context, Rc<Context>> again k0kubun@5f80d41, but again, it didn't give a significant difference (changed only 0.2%):

before

compiled_iseq_count:           2,200
compiled_block_count:         20,715
compiled_branch_count:        36,100

yjit_alloc_size:           9,170,201
live_context_size:         1,021,212
live_context_count:           26,874

after

compiled_iseq_count:           2,200
compiled_block_count:         20,778
compiled_branch_count:        36,227

yjit_alloc_size:           9,152,465
live_context_size:           216,176
live_context_count:           27,022
context_table_size:          442,368
context_table_count:           6,035

• live_context_size: The size used for storing a reference Rc<Context>, which is part of Block and BranchStub
• live_context_count: The number of context references
• context_table_size: The estimated size of the deduplication HashMap
• context_table_count: The number of entries in the deduplication HashMap

---

In this change, size_of::<Block>() is reduced from 176 to 144 (-32), and size_of::<BranchStub>() is reduced from 58 to 26 (-26).

Because compiled_block_count could count freed blocks as well, I prepared live_block_count and it was 15,340. So there were 11682 branch stubs.

32 * 15340 + 26 * 11682 is 794612. This is offset by the HashMap size 442368 and Contexts referenced by Rc, which would consume 6035 * 38 = 229330. So it's supposed to give a 122914 decrease in memory usage, however, it only gave a 17736 decrease. Allocating individual 38 bytes on heap wouldn't be that storage-efficient, so the number is not that surprising.

---

I think we could also try having a vector of unique Contexts per iseq and store an index to it instead of Rc.

[maximecb]

Thank you for putting together this experiment.

If I understand correctly, before, there were 26,874 context objects. With the table, there are 6,035 entries in the table. That's the number of deduplicated contexts. So there is about a 5x reduction in the number of contexts.

Maybe the problem is that for HashMap<Context, Rc<Context>>, there's a Context stored once in the key, and another one allocated on the heap for the value. Then, the HashMap also has some storage overhead (~2x) because hash maps can't be too full, otherwise it creates too many collisions.

I think we could also try having a vector of unique Contexts per iseq and store an index to it instead of Rc.

Feel free to run that experiment if you want to. I think the potential issue there is that we would need to look for existing contexts in the hash map, which might scale poorly for larger ISEQs because it's effectively a linear O(n) search?

This is not an easy problem to solve hmm.

[k0kubun]

Maybe the problem is that for HashMap<Context, Rc>, there's a Context stored once in the key, and another one allocated on the heap for the value. Then, the HashMap also has some storage overhead (~2x) because hash maps can't be too full, otherwise it creates too many collisions.

Yeah, I figured out how to address this and was about to comment on that, actually. It was as easy as just changing it to HashMap<Rc<Context>, Rc<Context>>, and everything seems to be working as intended.

Updated numbers:

compiled_iseq_count:           2,200
compiled_block_count:         20,777
compiled_branch_count:        36,220
...
yjit_alloc_size:           8,887,953
live_context_size:           216,056
live_context_count:           27,007
context_table_size:          442,368
context_table_count:           6,036

A 282248 decrease seems like a good-enough number for this deduplication idea. I've just rebased my Packed Context branch from latest master and measured it again, and yjit_alloc_size was 8,902,985, i.e. only a 267216 decrease.

So the deduplication idea seems fairly efficient. The only concern on the deduplication idea is GC, but maybe it's fine to lazily sweep them on code_gc (using Rc's strong count)?

[k0kubun]

It turns out that you could get a key from HashSet (I thought you could only query contains), so it's compatible with that idea (and a nicer construct for readability). Now it's a 336052 decrease.

yjit_alloc_size:           8,834,149

I think this is good enough for now. I'll try making a PR with this idea and Rc-based GC on code_gc. (edit: ruby#7332)

[k0kubun]

We have open issues or pull requests for everything discussed in this issue, so I'll close this. I'll probably profile the same thing again after Alan mergers his Rc removal.

[maximecb]

Alan and I rebased the yjit-dhat branch on master and viewed the profile for railsbench using the DHAT viewer: https://nnethercote.github.io/dh_view/dh_view.html
Filter option: at t-gmax (bytes)

./run_once.sh --yjit-stats benchmarks/railsbench/benchmark.rb

Total alloc size:
  At t-gmax: 8,485,146 bytes (100%) in 175,820 blocks (100%), avg size 48.26 bytes

Blocks:
   At t-gmax: 2,368,512 bytes (27.91%) in 18,504 blocks (10.52%), avg size 128 bytes

Branches:
 At t-gmax: 1,656,720 bytes (19.52%) in 34,515 blocks (19.63%), avg size 48 bytes

Assumptions table:
 At t-gmax: 691,728 bytes (8.15%) in 5,550 blocks (3.16%), avg size 124.64 bytes

Branch targets:
 At t-gmax: 556,840 bytes (6.56%) in 13,921 blocks (7.92%), avg size 40 bytes

Iseq Payloads:
 At t-gmax: 257,040 bytes (3.03%) in 2,142 blocks (1.22%), avg size 120 bytes

The good news is that the yjit alloc size is smaller than previously, however the code size has increased to:

code_region_size:          9,142,272

The blocks are fairly large objects still, so we might be able to do more to reduce their size, but if we could somehow reduce the number of blocks (or reduce code size?), that would probably have the biggest impact.

Code size may be larger on ARM64 than on x86 (to be verified).