mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-06 17:24:59 +08:00
f9022ba93b
Summary: Enhanced the PyTorch CUDA memory visualizer to display user_metadata alongside stack frames when inspecting allocations. The user_metadata field is now shown in all views (Allocator State History, Active Memory Timeline, etc.) with consistent formatting. The implementation handles both string and object metadata types, displaying strings directly and objects as key-value pairs.
Test Plan:
1. Generate a memory snapshot with user_metadata
2. Open the memory visualizer in a browser
3. Load the snapshot file
4. Verify user_metadata appears
5. Test with both string metadata ("testing") and object metadata ({"key": "value"})
6. Verify formatting shows "User Metadata:\n <value>" for strings
{F1982860439}
Differential Revision: D85095152
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165939
Approved by: https://github.com/yushangdi
1827 lines
50 KiB
JavaScript
1827 lines
50 KiB
JavaScript
'use strict';
|
|
|
|
import * as d3 from "https://cdn.skypack.dev/d3@5";
|
|
import {axisLeft} from "https://cdn.skypack.dev/d3-axis@1";
|
|
import {scaleLinear} from "https://cdn.skypack.dev/d3-scale@1";
|
|
import {zoom, zoomIdentity} from "https://cdn.skypack.dev/d3-zoom@1";
|
|
import {brushX} from "https://cdn.skypack.dev/d3-brush@1";
|
|
|
|
const schemeTableau10 = [
|
|
'#4e79a7',
|
|
'#f28e2c',
|
|
'#e15759',
|
|
'#76b7b2',
|
|
'#59a14f',
|
|
'#edc949',
|
|
'#af7aa1',
|
|
'#ff9da7',
|
|
'#9c755f',
|
|
'#bab0ab',
|
|
];
|
|
|
|
function version_space() {
|
|
const version = {};
|
|
return (addr, increment) => {
|
|
if (!(addr in version)) {
|
|
version[addr] = 0;
|
|
}
|
|
const r = version[addr];
|
|
if (increment) {
|
|
version[addr]++;
|
|
}
|
|
return r;
|
|
};
|
|
}
|
|
|
|
function Segment(addr, size, stream, frames, version, user_metadata) {
|
|
return {addr, size, stream, version, frames, user_metadata};
|
|
}
|
|
|
|
function Block(addr, size, requested_size, frames, free_requested, version, user_metadata) {
|
|
return {addr, size, requested_size, frames, free_requested, version, user_metadata};
|
|
}
|
|
|
|
function EventSelector(outer, events, stack_info, memory_view) {
|
|
const events_div = outer
|
|
.append('div')
|
|
.attr(
|
|
'style',
|
|
'grid-column: 1; grid-row: 1; overflow: auto; font-family: monospace',
|
|
);
|
|
|
|
const events_selection = events_div
|
|
.selectAll('pre')
|
|
.data(events)
|
|
.enter()
|
|
.append('pre')
|
|
.text(e => formatEvent(e))
|
|
.attr('style', '');
|
|
|
|
let selected_event_idx = null;
|
|
|
|
const es = {
|
|
select(idx) {
|
|
if (selected_event_idx !== null) {
|
|
const selected_event = d3.select(
|
|
events_div.node().children[selected_event_idx],
|
|
);
|
|
selected_event.attr('style', '');
|
|
}
|
|
if (idx !== null) {
|
|
const div = d3.select(events_div.node().children[idx]);
|
|
div.attr('style', `background-color: ${schemeTableau10[5]}`);
|
|
const [reserved, allocated] = memory_view.draw(idx);
|
|
const enter = () => eventStack(div.datum(), allocated, reserved);
|
|
stack_info.highlight(enter);
|
|
div.node().scrollIntoViewIfNeeded(false);
|
|
} else {
|
|
memory_view.draw(0);
|
|
}
|
|
selected_event_idx = idx;
|
|
},
|
|
};
|
|
d3.select('body').on('keydown', _e => {
|
|
const key = d3.event.key;
|
|
const actions = {ArrowDown: 1, ArrowUp: -1};
|
|
if (selected_event_idx !== null && key in actions) {
|
|
const new_idx = selected_event_idx + actions[key];
|
|
es.select(Math.max(0, Math.min(new_idx, events.length - 1)));
|
|
d3.event.preventDefault();
|
|
}
|
|
});
|
|
|
|
stack_info.register(
|
|
events_selection,
|
|
t => eventStack(t.datum()),
|
|
_t => {},
|
|
d => es.select(d.datum().idx),
|
|
);
|
|
|
|
return es;
|
|
}
|
|
|
|
function formatSize(num) {
|
|
const orig = num;
|
|
// https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
|
|
const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'];
|
|
for (const unit of units) {
|
|
if (Math.abs(num) < 1024.0) {
|
|
return `${num.toFixed(1)}${unit}B (${orig} bytes)`;
|
|
}
|
|
num /= 1024.0;
|
|
}
|
|
return `${num.toFixed(1)}YiB`;
|
|
}
|
|
function formatAddr(event) {
|
|
const prefix = event.action.startsWith('segment') ? 's\'' : 'b\'';
|
|
return `${prefix}${event.addr.toString(16)}_${event.version}`;
|
|
}
|
|
function formatEvent(event) {
|
|
const stream =
|
|
event.stream === null ? '' : `\n (stream ${event.stream})`;
|
|
switch (event.action) {
|
|
case 'oom':
|
|
return `OOM (requested ${formatSize(event.size)}, CUDA has ${formatSize(
|
|
event.device_free,
|
|
)} memory free)${stream}`;
|
|
case 'snapshot':
|
|
return 'snapshot';
|
|
default:
|
|
return `${event.action.padEnd(14)} ${formatAddr(event).padEnd(
|
|
18,
|
|
)} ${formatSize(event.size)}${stream}`;
|
|
}
|
|
}
|
|
|
|
function eventStack(e, allocated, reserved) {
|
|
let event = formatEvent(e);
|
|
if (reserved !== undefined) {
|
|
event = `(${formatSize(allocated)} allocated / ${formatSize(
|
|
reserved,
|
|
)} reserved)\n${event}`;
|
|
}
|
|
const user_metadata_str = format_user_metadata(e.user_metadata);
|
|
const frames_str = format_frames(e.frames);
|
|
return event + '\n' + (user_metadata_str ? user_metadata_str + '\n' : '') + frames_str;
|
|
}
|
|
|
|
function hashCode(num) {
|
|
const numStr = num.toString();
|
|
let hash = 0;
|
|
for (let i = 0; i < numStr.length; i++) {
|
|
const charCode = numStr.charCodeAt(i);
|
|
hash = (hash << 5) - hash + charCode;
|
|
hash = hash & hash; // Convert to 32-bit integer
|
|
}
|
|
return hash;
|
|
}
|
|
|
|
function addStroke(d) {
|
|
d.attr('stroke', 'red')
|
|
.attr('stroke-width', '2')
|
|
.attr('vector-effect', 'non-scaling-stroke');
|
|
}
|
|
|
|
function removeStroke(d) {
|
|
d.attr('stroke', '');
|
|
}
|
|
|
|
function calculate_fragmentation(blocks, sorted_segments) {
|
|
const sorted_blocks = Object.values(blocks).sort((a, b) => a.addr - b.addr);
|
|
let block_i = 0;
|
|
let total_size = 0;
|
|
let sum_squared_free = 0;
|
|
for (const seg of sorted_segments) {
|
|
let addr = seg.addr;
|
|
total_size += seg.size;
|
|
while (
|
|
block_i < sorted_blocks.length &&
|
|
sorted_blocks[block_i].addr < seg.addr + seg.size
|
|
) {
|
|
const block = sorted_blocks[block_i];
|
|
if (block.addr > addr) {
|
|
sum_squared_free += (block.addr - addr) ** 2;
|
|
}
|
|
addr = block.addr + block.size;
|
|
block_i += 1;
|
|
}
|
|
if (addr < seg.addr + seg.size) {
|
|
sum_squared_free += (seg.addr + seg.size - addr) ** 2;
|
|
}
|
|
}
|
|
console.log(sum_squared_free / total_size ** 2);
|
|
}
|
|
|
|
function MemoryView(outer, stack_info, snapshot, device) {
|
|
const svg = outer
|
|
.append('svg')
|
|
.attr('style', 'grid-column: 2; grid-row: 1; width: 100%; height: 100%;')
|
|
.attr('viewBox', '0 0 200 100')
|
|
.attr('preserveAspectRatio', 'xMinYMin meet');
|
|
const g = svg.append('g');
|
|
const seg_zoom = zoom();
|
|
seg_zoom.on('zoom', () => {
|
|
g.attr('transform', d3.event.transform);
|
|
});
|
|
svg.call(seg_zoom);
|
|
|
|
const sorted_segments = [];
|
|
const block_map = {};
|
|
for (const seg of snapshot.segments) {
|
|
if (seg.device !== device) {
|
|
continue;
|
|
}
|
|
sorted_segments.push(
|
|
Segment(
|
|
seg.address,
|
|
seg.total_size,
|
|
seg.stream,
|
|
seg.frames || [],
|
|
seg.version,
|
|
seg.user_metadata,
|
|
),
|
|
);
|
|
for (const b of seg.blocks) {
|
|
if (b.state !== 'active_pending_free' && b.state !== 'active_allocated') {
|
|
continue;
|
|
}
|
|
block_map[b.addr] = Block(
|
|
b.addr,
|
|
b.size,
|
|
b.requested_size,
|
|
b.frames,
|
|
b.state === 'active_pending_free',
|
|
b.version,
|
|
b.user_metadata,
|
|
);
|
|
}
|
|
}
|
|
sorted_segments.sort((x, y) => x.addr - y.addr);
|
|
|
|
function simulate_memory(idx) {
|
|
// create a copy of segments because we edit size properties below
|
|
const l_segments = sorted_segments.map(x => {
|
|
return {...x};
|
|
});
|
|
const l_block_map = {...block_map};
|
|
|
|
function map_segment(merge, seg) {
|
|
let idx = l_segments.findIndex(e => e.addr > seg.addr);
|
|
if (!merge) {
|
|
l_segments.splice(idx, 0, seg);
|
|
return;
|
|
}
|
|
if (idx === -1) {
|
|
idx = l_segments.length;
|
|
}
|
|
l_segments.splice(idx, 0, seg);
|
|
if (idx + 1 < l_segments.length) {
|
|
const next = l_segments[idx + 1];
|
|
if (seg.addr + seg.size === next.addr && seg.stream === next.stream) {
|
|
seg.size += next.size;
|
|
l_segments.splice(idx + 1, 1);
|
|
}
|
|
}
|
|
if (idx > 0) {
|
|
const prev = l_segments[idx - 1];
|
|
if (prev.addr + prev.size === seg.addr && prev.stream === seg.stream) {
|
|
prev.size += seg.size;
|
|
l_segments.splice(idx, 1);
|
|
}
|
|
}
|
|
}
|
|
function unmap_segment(merge, seg) {
|
|
if (!merge) {
|
|
l_segments.splice(
|
|
l_segments.findIndex(x => x.addr === seg.addr),
|
|
1,
|
|
);
|
|
return;
|
|
}
|
|
const seg_end = seg.addr + seg.size;
|
|
const idx = l_segments.findIndex(
|
|
e => e.addr <= seg.addr && seg_end <= e.addr + e.size,
|
|
);
|
|
const existing = l_segments[idx];
|
|
const existing_end = existing.addr + existing.size;
|
|
if (existing.addr === seg.addr) {
|
|
existing.addr += seg.size;
|
|
existing.size -= seg.size;
|
|
if (existing.size === 0) {
|
|
l_segments.splice(idx, 1);
|
|
}
|
|
} else if (existing_end === seg_end) {
|
|
existing.size -= seg.size;
|
|
} else {
|
|
existing.size = seg.addr - existing.addr;
|
|
seg.addr = seg_end;
|
|
seg.size = existing_end - seg_end;
|
|
l_segments.splice(idx + 1, 0, seg);
|
|
}
|
|
}
|
|
const events = snapshot.device_traces[device];
|
|
for (let i = events.length - 1; i > idx; i--) {
|
|
const event = events[i];
|
|
switch (event.action) {
|
|
case 'free':
|
|
l_block_map[event.addr] = Block(
|
|
event.addr,
|
|
event.size,
|
|
event.size,
|
|
event.frames,
|
|
false,
|
|
event.version,
|
|
event.user_metadata,
|
|
);
|
|
break;
|
|
case 'free_requested':
|
|
l_block_map[event.addr].free_requested = false;
|
|
break;
|
|
case 'free_completed':
|
|
l_block_map[event.addr] = Block(
|
|
event.addr,
|
|
event.size,
|
|
event.size,
|
|
event.frames,
|
|
true,
|
|
event.version,
|
|
event.user_metadata,
|
|
);
|
|
break;
|
|
case 'alloc':
|
|
delete l_block_map[event.addr];
|
|
break;
|
|
case 'segment_free':
|
|
case 'segment_unmap':
|
|
map_segment(
|
|
event.action === 'segment_unmap',
|
|
Segment(
|
|
event.addr,
|
|
event.size,
|
|
event.stream,
|
|
event.frames,
|
|
event.version,
|
|
event.user_metadata,
|
|
),
|
|
);
|
|
break;
|
|
case 'segment_alloc':
|
|
case 'segment_map':
|
|
unmap_segment(
|
|
event.action === 'segment_map',
|
|
Segment(
|
|
event.addr,
|
|
event.size,
|
|
event.stream,
|
|
event.frames,
|
|
event.version,
|
|
event.user_metadata,
|
|
),
|
|
);
|
|
break;
|
|
case 'oom':
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
const new_blocks = Object.values(l_block_map);
|
|
return [l_segments, new_blocks];
|
|
}
|
|
|
|
return {
|
|
draw(idx) {
|
|
const [segments_unsorted, blocks] = simulate_memory(idx);
|
|
g.selectAll('g').remove();
|
|
|
|
const segment_d = g.append('g');
|
|
const block_g = g.append('g');
|
|
const block_r = g.append('g');
|
|
|
|
segment_d.selectAll('rect').remove();
|
|
block_g.selectAll('rect').remove();
|
|
block_r.selectAll('rect').remove();
|
|
const segments = [...segments_unsorted].sort((x, y) =>
|
|
x.size === y.size ? x.addr - y.addr : x.size - y.size,
|
|
);
|
|
|
|
const segments_by_addr = [...segments].sort((x, y) => x.addr - y.addr);
|
|
|
|
const max_size = segments.length === 0 ? 0 : segments.at(-1).size;
|
|
|
|
const xScale = scaleLinear().domain([0, max_size]).range([0, 200]);
|
|
const padding = xScale.invert(1);
|
|
|
|
let cur_row = 0;
|
|
let cur_row_size = 0;
|
|
for (const seg of segments) {
|
|
seg.occupied = 0;
|
|
seg.internal_free = 0;
|
|
if (cur_row_size + seg.size > max_size) {
|
|
cur_row_size = 0;
|
|
cur_row += 1;
|
|
}
|
|
seg.offset = cur_row_size;
|
|
seg.row = cur_row;
|
|
cur_row_size += seg.size + padding;
|
|
}
|
|
|
|
const num_rows = cur_row + 1;
|
|
|
|
const yScale = scaleLinear().domain([0, num_rows]).range([0, 100]);
|
|
|
|
const segments_selection = segment_d
|
|
.selectAll('rect')
|
|
.data(segments)
|
|
.enter()
|
|
.append('rect')
|
|
.attr('x', x => xScale(x.offset))
|
|
.attr('y', x => yScale(x.row))
|
|
.attr('width', x => xScale(x.size))
|
|
.attr('height', yScale(4 / 5))
|
|
.attr('stroke', 'black')
|
|
.attr('stroke-width', '1')
|
|
.attr('vector-effect', 'non-scaling-stroke')
|
|
.attr('fill', 'white');
|
|
|
|
stack_info.register(
|
|
segments_selection,
|
|
d => {
|
|
addStroke(d);
|
|
const t = d.datum();
|
|
const free = t.size - t.occupied;
|
|
let internal = '';
|
|
if (t.internal_free > 0) {
|
|
internal = ` (${(t.internal_free / free) * 100}% internal)`;
|
|
}
|
|
const user_metadata_str = format_user_metadata(t.user_metadata);
|
|
const frames_str = format_frames(t.frames);
|
|
return (
|
|
`s${t.addr.toString(16)}_${t.version}: segment ${formatSize(
|
|
t.size,
|
|
)} allocated, ` +
|
|
`${formatSize(free)} free${internal} (stream ${
|
|
t.stream
|
|
})\n` +
|
|
(user_metadata_str ? user_metadata_str + '\n' : '') +
|
|
frames_str
|
|
);
|
|
},
|
|
d => {
|
|
d.attr('stroke', 'black')
|
|
.attr('stroke-width', '1')
|
|
.attr('vector-effect', 'non-scaling-stroke');
|
|
},
|
|
);
|
|
|
|
function find_segment(addr) {
|
|
let left = 0;
|
|
let right = segments_by_addr.length - 1;
|
|
while (left <= right) {
|
|
const mid = Math.floor((left + right) / 2);
|
|
if (addr < segments_by_addr[mid].addr) {
|
|
right = mid - 1;
|
|
} else if (
|
|
addr >=
|
|
segments_by_addr[mid].addr + segments_by_addr[mid].size
|
|
) {
|
|
left = mid + 1;
|
|
} else {
|
|
return segments_by_addr[mid];
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
for (const b of blocks) {
|
|
b.segment = find_segment(b.addr);
|
|
b.segment.occupied += b.requested_size;
|
|
b.segment.internal_free += b.size - b.requested_size;
|
|
}
|
|
|
|
const block_selection = block_g
|
|
.selectAll('rect')
|
|
.data(blocks)
|
|
.enter()
|
|
.append('rect')
|
|
.attr('x', x => xScale(x.segment.offset + (x.addr - x.segment.addr)))
|
|
.attr('y', x => yScale(x.segment.row))
|
|
.attr('width', x => xScale(x.requested_size))
|
|
.attr('height', yScale(4 / 5))
|
|
.attr('fill', (x, _i) =>
|
|
x.free_requested
|
|
? 'red'
|
|
: schemeTableau10[
|
|
Math.abs(hashCode(x.addr)) % schemeTableau10.length
|
|
],
|
|
);
|
|
|
|
stack_info.register(
|
|
block_selection,
|
|
d => {
|
|
addStroke(d);
|
|
const t = d.datum();
|
|
let requested = '';
|
|
if (t.free_requested) {
|
|
requested = ' (block freed but waiting due to record_stream)';
|
|
}
|
|
const user_metadata_str = format_user_metadata(t.user_metadata);
|
|
const frames_str = format_frames(t.frames);
|
|
return (
|
|
`b${t.addr.toString(16)}_${t.version} ` +
|
|
`${formatSize(t.requested_size)} allocation${requested} (stream ${
|
|
t.segment.stream
|
|
})\n` +
|
|
(user_metadata_str ? user_metadata_str + '\n' : '') +
|
|
frames_str
|
|
);
|
|
},
|
|
removeStroke,
|
|
);
|
|
|
|
const free_selection = block_r
|
|
.selectAll('rect')
|
|
.data(blocks)
|
|
.enter()
|
|
.append('rect')
|
|
.attr('x', x =>
|
|
xScale(
|
|
x.segment.offset + (x.addr - x.segment.addr) + x.requested_size,
|
|
),
|
|
)
|
|
.attr('y', x => yScale(x.segment.row))
|
|
.attr('width', x => xScale(x.size - x.requested_size))
|
|
.attr('height', yScale(4 / 5))
|
|
.attr('fill', (_x, _i) => 'red');
|
|
|
|
stack_info.register(
|
|
free_selection,
|
|
d => {
|
|
addStroke(d);
|
|
const t = d.datum();
|
|
const user_metadata_str = format_user_metadata(t.user_metadata);
|
|
const frames_str = format_frames(t.frames);
|
|
return (
|
|
`Free space lost due to rounding ${formatSize(
|
|
t.size - t.requested_size,
|
|
)}` +
|
|
` (stream ${t.segment.stream})\n` +
|
|
(user_metadata_str ? user_metadata_str + '\n' : '') +
|
|
frames_str
|
|
);
|
|
},
|
|
removeStroke,
|
|
);
|
|
|
|
const reserved = segments.reduce((x, y) => x + y.size, 0);
|
|
const allocated = blocks.reduce((x, y) => x + y.requested_size, 0);
|
|
return [reserved, allocated];
|
|
},
|
|
};
|
|
}
|
|
|
|
function StackInfo(outer) {
|
|
const stack_trace = outer
|
|
.append('pre')
|
|
.attr('style', 'grid-column: 1 / 3; grid-row: 2; overflow: auto');
|
|
let selected = {
|
|
enter: () => {
|
|
stack_trace.text('');
|
|
},
|
|
leave: () => {},
|
|
};
|
|
return {
|
|
register(dom, enter, leave = _e => {}, select = _e => {}) {
|
|
dom
|
|
.on('mouseover', _e => {
|
|
selected.leave();
|
|
stack_trace.text(enter(d3.select(d3.event.target)));
|
|
})
|
|
.on('mousedown', _e => {
|
|
const obj = d3.select(d3.event.target);
|
|
selected = {
|
|
enter: () => stack_trace.text(enter(obj)),
|
|
leave: () => leave(obj),
|
|
};
|
|
select(obj);
|
|
})
|
|
.on('mouseleave', _e => {
|
|
leave(d3.select(d3.event.target));
|
|
selected.enter();
|
|
});
|
|
},
|
|
highlight(enter, leave = () => {}) {
|
|
selected = {enter: () => stack_trace.text(enter()), leave};
|
|
selected.enter();
|
|
},
|
|
};
|
|
}
|
|
|
|
function create_segment_view(dst, snapshot, device) {
|
|
const outer = dst
|
|
.append('div')
|
|
.attr(
|
|
'style',
|
|
'display: grid; grid-template-columns: 1fr 2fr; grid-template-rows: 2fr 1fr; height: 100%; gap: 10px',
|
|
);
|
|
|
|
const events = snapshot.device_traces[device];
|
|
const stack_info = StackInfo(outer);
|
|
const memory_view = MemoryView(outer, stack_info, snapshot, device);
|
|
const event_selector = EventSelector(outer, events, stack_info, memory_view);
|
|
|
|
window.requestAnimationFrame(function () {
|
|
event_selector.select(events.length > 0 ? events.length - 1 : null);
|
|
});
|
|
}
|
|
|
|
function annotate_snapshot(snapshot) {
|
|
snapshot.segment_version = version_space();
|
|
snapshot.block_version = version_space();
|
|
snapshot.categories = [];
|
|
const empty_list = [];
|
|
let next_stream = 1;
|
|
const stream_names = {0: 0};
|
|
function stream_name(s) {
|
|
if (!(s in stream_names)) {
|
|
stream_names[s] = next_stream++;
|
|
}
|
|
return stream_names[s];
|
|
}
|
|
const new_traces = [];
|
|
for (const device_trace of snapshot.device_traces) {
|
|
const new_trace = [];
|
|
new_traces.push(new_trace);
|
|
for (const t of device_trace) {
|
|
if (!('frames' in t)) {
|
|
t.frames = empty_list;
|
|
}
|
|
// set unique version for each time an address is used
|
|
// so that ctrl-f can be used to search for the beginning
|
|
// and end of allocations and segments
|
|
t.stream = stream_name(t.stream);
|
|
switch (t.action) {
|
|
case 'free_completed':
|
|
t.version = snapshot.block_version(t.addr, true);
|
|
if (new_trace.length > 0) {
|
|
// elide free_requested/free_completed into a single event
|
|
const prev = new_trace.at(-1);
|
|
if (prev.action === 'free_requested' && prev.addr === t.addr) {
|
|
prev.action = 'free';
|
|
continue;
|
|
}
|
|
}
|
|
break;
|
|
case 'free_requested':
|
|
case 'alloc':
|
|
t.version = snapshot.block_version(t.addr, false);
|
|
break;
|
|
case 'segment_free':
|
|
case 'segment_unmap':
|
|
t.version = snapshot.segment_version(t.addr, true);
|
|
break;
|
|
case 'segment_alloc':
|
|
case 'segment_map':
|
|
t.version = snapshot.segment_version(t.addr, false);
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
if ('category' in t && !snapshot.categories.includes(t.category)) {
|
|
snapshot.categories.push(t.category);
|
|
}
|
|
t.idx = new_trace.length;
|
|
new_trace.push(t);
|
|
}
|
|
}
|
|
snapshot.device_traces = new_traces;
|
|
// if every event was on the default stream, we elide stream printing
|
|
if (next_stream == 1) {
|
|
for (const device_trace of snapshot.device_traces) {
|
|
for (const t of device_trace) {
|
|
t.stream = null;
|
|
}
|
|
}
|
|
}
|
|
|
|
for (const seg of snapshot.segments) {
|
|
seg.stream = stream_name(seg.stream);
|
|
seg.version = snapshot.segment_version(seg.address, false);
|
|
let addr = seg.address;
|
|
for (const b of seg.blocks) {
|
|
b.addr = addr;
|
|
if (!('frames' in b)) {
|
|
// legacy format where 'requested_size' may be missing
|
|
// and frames might be in history rather than directly on block
|
|
if ('history' in b) {
|
|
b.frames = b.history[0].frames || empty_list;
|
|
b.requested_size = b.requested_size || b.history[0].real_size;
|
|
} else {
|
|
b.frames = empty_list;
|
|
b.requested_size = b.requested_size || b.size;
|
|
}
|
|
}
|
|
b.version = snapshot.block_version(b.addr, false);
|
|
addr += b.size;
|
|
}
|
|
}
|
|
|
|
if (
|
|
snapshot.categories.length > 0 &&
|
|
!snapshot.categories.includes('unknown')
|
|
) {
|
|
snapshot.categores.push('unknown');
|
|
}
|
|
}
|
|
|
|
function elideRepeats(frames) {
|
|
const result = [];
|
|
const length = frames.length;
|
|
for (let i = 0; i < length; ) {
|
|
let j = i + 1;
|
|
const f = frames[i];
|
|
while (j < length && f === frames[j]) {
|
|
j++;
|
|
}
|
|
switch (j - i) {
|
|
case 1:
|
|
result.push(f);
|
|
break;
|
|
case 2:
|
|
result.push(f, f);
|
|
break;
|
|
default:
|
|
result.push(f, `<repeats ${j - i - 1} times>`);
|
|
break;
|
|
}
|
|
i = j;
|
|
}
|
|
return result;
|
|
}
|
|
function frameFilter({name, filename}) {
|
|
const omitFunctions = [
|
|
'unwind::unwind',
|
|
'CapturedTraceback::gather',
|
|
'gather_with_cpp',
|
|
'_start',
|
|
'__libc_start_main',
|
|
'PyEval_',
|
|
'PyObject_',
|
|
'PyFunction_',
|
|
];
|
|
|
|
const omitFilenames = [
|
|
'core/boxing',
|
|
'/Register',
|
|
'/Redispatch',
|
|
'pythonrun.c',
|
|
'Modules/main.c',
|
|
'Objects/call.c',
|
|
'Objects/methodobject.c',
|
|
'pycore_ceval.h',
|
|
'ceval.c',
|
|
'cpython/abstract.h',
|
|
];
|
|
|
|
for (const of of omitFunctions) {
|
|
if (name.includes(of)) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
for (const of of omitFilenames) {
|
|
if (filename.includes(of)) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
function format_user_metadata(user_metadata) {
|
|
if (!user_metadata) {
|
|
return '';
|
|
}
|
|
// Handle string metadata
|
|
if (typeof user_metadata === 'string') {
|
|
return `User Metadata:\n ${user_metadata}`;
|
|
}
|
|
// Handle object metadata
|
|
if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) {
|
|
return '';
|
|
}
|
|
const metadata_lines = Object.entries(user_metadata)
|
|
.map(([key, value]) => ` ${key}: ${value}`);
|
|
return 'User Metadata:\n' + metadata_lines.join('\n');
|
|
}
|
|
|
|
function format_frames(frames) {
|
|
if (frames.length === 0) {
|
|
return (
|
|
`This block has no frames. Potential causes:\n` +
|
|
`1) This block was allocated before _record_memory_history was enabled.\n` +
|
|
`2) The context or stacks passed to _record_memory_history does not include this block. Consider changing context to 'state', 'alloc', or 'all', or changing stacks to 'all'.\n` +
|
|
`3) This event occurred during backward, which has no python frames, and memory history did not include C++ frames. Use stacks='all' to record both C++ and python frames.`
|
|
);
|
|
}
|
|
const frame_strings = frames
|
|
.filter(frameFilter)
|
|
.map(f => `${f.filename}:${f.line}:${f.name}`);
|
|
return elideRepeats(frame_strings).join('\n');
|
|
}
|
|
|
|
function process_alloc_data(snapshot, device, plot_segments, max_entries) {
|
|
const elements = [];
|
|
const initially_allocated = [];
|
|
const actions = [];
|
|
const addr_to_alloc = {};
|
|
|
|
const alloc = plot_segments ? 'segment_alloc' : 'alloc';
|
|
const [free, free_completed] = plot_segments
|
|
? ['segment_free', 'segment_free']
|
|
: ['free', 'free_completed'];
|
|
for (const e of snapshot.device_traces[device]) {
|
|
switch (e.action) {
|
|
case alloc:
|
|
elements.push(e);
|
|
addr_to_alloc[e.addr] = elements.length - 1;
|
|
actions.push(elements.length - 1);
|
|
break;
|
|
case free:
|
|
case free_completed:
|
|
if (e.addr in addr_to_alloc) {
|
|
actions.push(addr_to_alloc[e.addr]);
|
|
delete addr_to_alloc[e.addr];
|
|
} else {
|
|
elements.push(e);
|
|
initially_allocated.push(elements.length - 1);
|
|
actions.push(elements.length - 1);
|
|
}
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
for (const seg of snapshot.segments) {
|
|
if (seg.device !== device) {
|
|
continue;
|
|
}
|
|
if (plot_segments) {
|
|
if (!(seg.address in addr_to_alloc)) {
|
|
const element = {
|
|
action: 'alloc',
|
|
addr: seg.address,
|
|
size: seg.total_size,
|
|
frames: [],
|
|
stream: seg.stream,
|
|
version: seg.version,
|
|
};
|
|
elements.push(element);
|
|
initially_allocated.push(elements.length - 1);
|
|
}
|
|
} else {
|
|
for (const b of seg.blocks) {
|
|
if (b.state === 'active_allocated' && !(b.addr in addr_to_alloc)) {
|
|
const element = {
|
|
action: 'alloc',
|
|
addr: b.addr,
|
|
size: b.requested_size,
|
|
frames: b.frames,
|
|
stream: seg.stream,
|
|
version: b.version,
|
|
};
|
|
elements.push(element);
|
|
initially_allocated.push(elements.length - 1);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
initially_allocated.reverse();
|
|
// if there are no actions, the graph will be blank,
|
|
// but if there are existing allocations we do not want to hide them
|
|
// by having just one allocate action it will show a flat graph with all segments
|
|
if (actions.length === 0 && initially_allocated.length > 0) {
|
|
actions.push(initially_allocated.pop());
|
|
}
|
|
|
|
const current = [];
|
|
const current_data = [];
|
|
const data = [];
|
|
let max_size = 0;
|
|
|
|
let total_mem = 0;
|
|
let total_summarized_mem = 0;
|
|
let timestep = 0;
|
|
|
|
const max_at_time = [];
|
|
|
|
const summarized_mem = {
|
|
elem: 'summarized',
|
|
timesteps: [],
|
|
offsets: [total_mem],
|
|
size: [],
|
|
color: 0,
|
|
};
|
|
const summarized_elems = {};
|
|
|
|
function advance(n) {
|
|
summarized_mem.timesteps.push(timestep);
|
|
summarized_mem.offsets.push(total_mem);
|
|
summarized_mem.size.push(total_summarized_mem);
|
|
timestep += n;
|
|
for (let i = 0; i < n; i++) {
|
|
max_at_time.push(total_mem + total_summarized_mem);
|
|
}
|
|
}
|
|
|
|
const sizes = elements
|
|
.map((x, i) => [x.size, i])
|
|
.sort(([x, _xi], [y, _yi]) => y - x);
|
|
|
|
const draw_elem = {};
|
|
for (const [_s, e] of sizes.slice(0, max_entries)) {
|
|
draw_elem[e] = true;
|
|
}
|
|
|
|
function add_allocation(elem) {
|
|
const element_obj = elements[elem];
|
|
const size = element_obj.size;
|
|
current.push(elem);
|
|
let color = elem;
|
|
if (snapshot.categories.length > 0) {
|
|
color = snapshot.categories.indexOf(element_obj.category || 'unknown');
|
|
}
|
|
const e = {
|
|
elem,
|
|
timesteps: [timestep],
|
|
offsets: [total_mem],
|
|
size,
|
|
color,
|
|
};
|
|
current_data.push(e);
|
|
data.push(e);
|
|
total_mem += size;
|
|
element_obj.max_allocated_mem = total_mem + total_summarized_mem;
|
|
}
|
|
|
|
for (const elem of initially_allocated) {
|
|
if (elem in draw_elem) {
|
|
add_allocation(elem);
|
|
} else {
|
|
total_summarized_mem += elements[elem].size;
|
|
summarized_elems[elem] = true;
|
|
}
|
|
}
|
|
|
|
for (const elem of actions) {
|
|
const size = elements[elem].size;
|
|
if (!(elem in draw_elem)) {
|
|
if (elem in summarized_elems) {
|
|
advance(1);
|
|
total_summarized_mem -= size;
|
|
summarized_elems[elem] = null;
|
|
} else {
|
|
total_summarized_mem += size;
|
|
summarized_elems[elem] = true;
|
|
advance(1);
|
|
}
|
|
continue;
|
|
}
|
|
const idx = current.findLastIndex(x => x === elem);
|
|
// first time we see an action we add it
|
|
// second time we remove it
|
|
if (idx === -1) {
|
|
add_allocation(elem);
|
|
advance(1);
|
|
} else {
|
|
advance(1);
|
|
const removed = current_data[idx];
|
|
removed.timesteps.push(timestep);
|
|
removed.offsets.push(removed.offsets.at(-1));
|
|
current.splice(idx, 1);
|
|
current_data.splice(idx, 1);
|
|
|
|
if (idx < current.length) {
|
|
for (let j = idx; j < current.length; j++) {
|
|
const e = current_data[j];
|
|
e.timesteps.push(timestep);
|
|
e.offsets.push(e.offsets.at(-1));
|
|
e.timesteps.push(timestep + 3);
|
|
e.offsets.push(e.offsets.at(-1) - size);
|
|
}
|
|
advance(3);
|
|
}
|
|
total_mem -= size;
|
|
}
|
|
max_size = Math.max(total_mem + total_summarized_mem, max_size);
|
|
}
|
|
|
|
for (const elem of current_data) {
|
|
elem.timesteps.push(timestep);
|
|
elem.offsets.push(elem.offsets.at(-1));
|
|
}
|
|
data.push(summarized_mem);
|
|
|
|
return {
|
|
max_size,
|
|
allocations_over_time: data,
|
|
max_at_time,
|
|
summarized_mem,
|
|
elements_length: elements.length,
|
|
context_for_id: id => {
|
|
const elem = elements[id];
|
|
let text = `Addr: ${formatAddr(elem)}`;
|
|
text = `${text}, Size: ${formatSize(elem.size)} allocation`;
|
|
text = `${text}, Total memory used after allocation: ${formatSize(
|
|
elem.max_allocated_mem,
|
|
)}`;
|
|
const context = elem?.compile_context ?? 'None';
|
|
text = `${text}, Compile context: ${context}`;
|
|
if (elem.stream !== null) {
|
|
text = `${text}, stream ${elem.stream}`;
|
|
}
|
|
if (elem.timestamp !== null) {
|
|
var d = new Date(elem.time_us / 1000);
|
|
text = `${text}, timestamp ${d}`;
|
|
}
|
|
if (!elem.action.includes('alloc')) {
|
|
text = `${text}\nalloc not recorded, stack trace for free:`;
|
|
}
|
|
const user_metadata_str = format_user_metadata(elem.user_metadata);
|
|
if (user_metadata_str) {
|
|
text = `${text}\n${user_metadata_str}`;
|
|
}
|
|
text = `${text}\n${format_frames(elem.frames)}`;
|
|
return text;
|
|
},
|
|
};
|
|
}
|
|
|
|
function MemoryPlot(
|
|
svg,
|
|
data,
|
|
left_pad,
|
|
width,
|
|
height,
|
|
colors = schemeTableau10,
|
|
) {
|
|
function format_points(d) {
|
|
const size = d.size;
|
|
const xs = d.timesteps.map(t => xscale(t));
|
|
const bottom = d.offsets.map(t => yscale(t));
|
|
const m = Array.isArray(size)
|
|
? (t, i) => yscale(t + size[i])
|
|
: t => yscale(t + size);
|
|
const top = d.offsets.map(m);
|
|
const p0 = xs.map((x, i) => `${x},${bottom[i]}`);
|
|
const p1 = xs.map((x, i) => `${x},${top[i]}`).reverse();
|
|
return `${p0.join(' ')} ${p1.join(' ')}`;
|
|
}
|
|
|
|
const max_timestep = data.max_at_time.length;
|
|
const max_size = data.max_size;
|
|
|
|
const plot_width = width - left_pad;
|
|
const plot_height = height;
|
|
|
|
const yscale = scaleLinear().domain([0, max_size]).range([plot_height, 0]);
|
|
const yaxis = axisLeft(yscale).tickFormat(d3.format('.3s'));
|
|
const xscale = scaleLinear().domain([0, max_timestep]).range([0, plot_width]);
|
|
const plot_coordinate_space = svg
|
|
.append('g')
|
|
.attr('transform', `translate(${left_pad}, ${0})`);
|
|
const plot_outer = plot_coordinate_space.append('g');
|
|
|
|
function view_rect(a) {
|
|
return a
|
|
.append('rect')
|
|
.attr('x', 0)
|
|
.attr('y', 0)
|
|
.attr('width', plot_width)
|
|
.attr('height', plot_height)
|
|
.attr('fill', 'white');
|
|
}
|
|
|
|
view_rect(plot_outer);
|
|
|
|
const cp = svg.append('clipPath').attr('id', 'clip');
|
|
view_rect(cp);
|
|
plot_outer.attr('clip-path', 'url(#clip)');
|
|
|
|
const zoom_group = plot_outer.append('g');
|
|
const scrub_group = zoom_group.append('g');
|
|
|
|
const plot = scrub_group
|
|
.selectAll('polygon')
|
|
.data(data.allocations_over_time)
|
|
.enter()
|
|
.append('polygon')
|
|
.attr('points', format_points)
|
|
.attr('fill', d => colors[d.color % colors.length]);
|
|
|
|
const axis = plot_coordinate_space.append('g').call(yaxis);
|
|
|
|
function handleZoom() {
|
|
const t = d3.event.transform;
|
|
zoom_group.attr('transform', t);
|
|
axis.call(yaxis.scale(d3.event.transform.rescaleY(yscale)));
|
|
}
|
|
|
|
const thezoom = zoom().on('zoom', handleZoom);
|
|
plot_outer.call(thezoom);
|
|
|
|
return {
|
|
select_window: (stepbegin, stepend, max) => {
|
|
const begin = xscale(stepbegin);
|
|
const size = xscale(stepend) - xscale(stepbegin);
|
|
const scale = plot_width / size;
|
|
const translate = -begin;
|
|
const yscale = max_size / max;
|
|
scrub_group.attr(
|
|
'transform',
|
|
`scale(${scale / yscale}, 1) translate(${translate}, 0)`,
|
|
);
|
|
plot_outer.call(
|
|
thezoom.transform,
|
|
zoomIdentity
|
|
.scale(yscale)
|
|
.translate(0, -(plot_height - plot_height / yscale)),
|
|
);
|
|
},
|
|
set_delegate: delegate => {
|
|
plot
|
|
.on('mouseover', function (_e, _d) {
|
|
delegate.set_selected(d3.select(this));
|
|
})
|
|
.on('mousedown', function (_e, _d) {
|
|
delegate.default_selected = d3.select(this);
|
|
})
|
|
.on('mouseleave', function (_e, _d) {
|
|
delegate.set_selected(delegate.default_selected);
|
|
});
|
|
},
|
|
};
|
|
}
|
|
|
|
function ContextViewer(text, data) {
|
|
let current_selected = null;
|
|
|
|
return {
|
|
default_selected: null,
|
|
set_selected: d => {
|
|
if (current_selected !== null) {
|
|
current_selected.attr('stroke', null).attr('stroke-width', null);
|
|
}
|
|
if (d === null) {
|
|
text.text('');
|
|
} else {
|
|
const dd = d.datum();
|
|
if (dd.elem === 'summarized') {
|
|
text.html(
|
|
'Small tensors that were not plotted to cutdown on render time.\n' +
|
|
'Use detail slider to see smaller allocations.',
|
|
);
|
|
} else {
|
|
text.text(`${dd.elem} ${data.context_for_id(dd.elem)}`);
|
|
}
|
|
d.attr('stroke', 'black')
|
|
.attr('stroke-width', 1)
|
|
.attr('vector-effect', 'non-scaling-stroke');
|
|
}
|
|
current_selected = d;
|
|
},
|
|
};
|
|
}
|
|
|
|
function MiniMap(mini_svg, plot, data, left_pad, width, height = 70) {
|
|
const max_at_time = data.max_at_time;
|
|
const plot_width = width - left_pad;
|
|
const yscale = scaleLinear().domain([0, data.max_size]).range([height, 0]);
|
|
const minixscale = scaleLinear()
|
|
.domain([0, max_at_time.length])
|
|
.range([left_pad, width]);
|
|
|
|
const mini_points = [
|
|
[max_at_time.length, 0],
|
|
[0, 0],
|
|
];
|
|
|
|
for (const [i, m] of max_at_time.entries()) {
|
|
const [_lastx, lasty] = mini_points[mini_points.length - 1];
|
|
if (m !== lasty) {
|
|
mini_points.push([i, lasty]);
|
|
mini_points.push([i, m]);
|
|
} else if (i === max_at_time.length - 1) {
|
|
mini_points.push([i, m]);
|
|
}
|
|
}
|
|
|
|
let points = mini_points.map(([t, o]) => `${minixscale(t)}, ${yscale(o)}`);
|
|
points = points.join(' ');
|
|
mini_svg
|
|
.append('polygon')
|
|
.attr('points', points)
|
|
.attr('fill', schemeTableau10[0]);
|
|
|
|
const xscale = scaleLinear()
|
|
.domain([0, max_at_time.length])
|
|
.range([0, plot_width]);
|
|
|
|
const brush = brushX();
|
|
brush.extent([
|
|
[left_pad, 0],
|
|
[width, height],
|
|
]);
|
|
brush.on('brush', function () {
|
|
const [begin, end] = d3.event.selection.map(x => x - left_pad);
|
|
|
|
const stepbegin = Math.floor(xscale.invert(begin));
|
|
const stepend = Math.floor(xscale.invert(end));
|
|
let max = 0;
|
|
for (let i = stepbegin; i < stepend; i++) {
|
|
max = Math.max(max, max_at_time[i]);
|
|
}
|
|
plot.select_window(stepbegin, stepend, max);
|
|
});
|
|
mini_svg.call(brush);
|
|
return {};
|
|
}
|
|
|
|
function Legend(plot_svg, categories) {
|
|
const xstart = 100;
|
|
const ystart = 5;
|
|
plot_svg
|
|
.append('g')
|
|
.selectAll('rect')
|
|
.data(categories)
|
|
.enter()
|
|
.append('rect')
|
|
.attr('x', (c, i) => xstart)
|
|
.attr('y', (c, i) => ystart + i * 15)
|
|
.attr('width', 10)
|
|
.attr('height', 10)
|
|
.attr('fill', (c, i) => schemeTableau10[i % schemeTableau10.length]);
|
|
plot_svg
|
|
.append('g')
|
|
.selectAll('text')
|
|
.data(categories)
|
|
.enter()
|
|
.append('text')
|
|
.attr('x', (c, i) => xstart + 20)
|
|
.attr('y', (c, i) => ystart + i * 15 + 8)
|
|
.attr('font-family', 'helvetica')
|
|
.attr('font-size', 10)
|
|
.text(c => c);
|
|
return {};
|
|
}
|
|
|
|
function create_trace_view(
|
|
dst,
|
|
snapshot,
|
|
device,
|
|
plot_segments = false,
|
|
max_entries = 15000,
|
|
) {
|
|
const left_pad = 70;
|
|
const data = process_alloc_data(snapshot, device, plot_segments, max_entries);
|
|
dst.selectAll('svg').remove();
|
|
dst.selectAll('div').remove();
|
|
|
|
max_entries = Math.min(max_entries, data.elements_length);
|
|
const d = dst.append('div');
|
|
d.append('input')
|
|
.attr('type', 'range')
|
|
.attr('min', 0)
|
|
.attr('max', data.elements_length)
|
|
.attr('value', max_entries)
|
|
.on('change', function () {
|
|
create_trace_view(dst, snapshot, device, plot_segments, this.value);
|
|
});
|
|
d.append('label').text(
|
|
`Detail: ${max_entries} of ${data.elements_length} entries`,
|
|
);
|
|
|
|
const grid_container = dst
|
|
.append('div')
|
|
.attr(
|
|
'style',
|
|
'display: grid; grid-template-columns: 1fr; grid-template-rows: 10fr 1fr 8fr; height: 100%; gap: 10px',
|
|
);
|
|
|
|
const plot_svg = grid_container
|
|
.append('svg')
|
|
.attr('display', 'block')
|
|
.attr('viewBox', '0 0 1024 576')
|
|
.attr('preserveAspectRatio', 'none')
|
|
.attr('style', 'grid-column: 1; grid-row: 1; width: 100%; height: 100%;');
|
|
|
|
const plot = MemoryPlot(plot_svg, data, left_pad, 1024, 576);
|
|
|
|
if (snapshot.categories.length !== 0) {
|
|
Legend(plot_svg.append('g'), snapshot.categories);
|
|
}
|
|
|
|
const mini_svg = grid_container
|
|
.append('svg')
|
|
.attr('display', 'block')
|
|
.attr('viewBox', '0 0 1024 60')
|
|
.attr('preserveAspectRatio', 'none')
|
|
.attr('style', 'grid-column: 1; grid-row: 2; width: 100%; height: 100%;');
|
|
|
|
MiniMap(mini_svg, plot, data, left_pad, 1024);
|
|
const context_div = grid_container
|
|
.append('div')
|
|
.attr(
|
|
'style',
|
|
'grid-column: 1; grid-row: 3; width: 100%; height: 100%; overflow: auto;',
|
|
);
|
|
const delegate = ContextViewer(context_div.append('pre').text('none'), data);
|
|
plot.set_delegate(delegate);
|
|
}
|
|
|
|
function create_settings_view(dst, snapshot, device) {
|
|
dst.selectAll('svg').remove();
|
|
dst.selectAll('div').remove();
|
|
const settings_div = dst.append('div');
|
|
settings_div.append('p').text('CUDA Caching Allocator Settings:');
|
|
|
|
// Check if allocator_settings exists in snapshot
|
|
if ('allocator_settings' in snapshot) {
|
|
settings_div
|
|
.append('pre')
|
|
.text(JSON.stringify(snapshot.allocator_settings, null, 2));
|
|
} else {
|
|
settings_div.append('p').text('No allocator settings found.');
|
|
}
|
|
}
|
|
|
|
function unpickle(buffer) {
|
|
try {
|
|
const decoder = new TextDecoder();
|
|
const jsonString = decoder.decode(new Uint8Array(buffer));
|
|
const data = JSON.parse(jsonString);
|
|
|
|
return data;
|
|
} catch (e) {
|
|
console.log('Failed to decode the data as JSON, fall back to pickle', e);
|
|
}
|
|
return unpickleData(buffer);
|
|
}
|
|
|
|
function unpickleData(buffer) {
|
|
const bytebuffer = new Uint8Array(buffer);
|
|
const decoder = new TextDecoder();
|
|
|
|
const stack = [];
|
|
const marks = [];
|
|
const memo = [];
|
|
let offset = 0;
|
|
let memo_id = 0;
|
|
|
|
const APPENDS = 'e'.charCodeAt(0);
|
|
const BINGET = 'h'.charCodeAt(0);
|
|
const BININT = 'J'.charCodeAt(0);
|
|
const BININT1 = 'K'.charCodeAt(0);
|
|
const BININT2 = 'M'.charCodeAt(0);
|
|
const EMPTY_DICT = '}'.charCodeAt(0);
|
|
const EMPTY_LIST = ']'.charCodeAt(0);
|
|
const FRAME = 0x95;
|
|
const LONG1 = 0x8a;
|
|
const LONG_BINGET = 'j'.charCodeAt(0);
|
|
const MARK = '('.charCodeAt(0);
|
|
const MEMOIZE = 0x94;
|
|
const PROTO = 0x80;
|
|
const SETITEMS = 'u'.charCodeAt(0);
|
|
const SHORT_BINUNICODE = 0x8c;
|
|
const STOP = '.'.charCodeAt(0);
|
|
const TUPLE2 = 0x86;
|
|
const APPEND = 'a'.charCodeAt(0);
|
|
const NEWFALSE = 0x89;
|
|
const BINPUT = 'q'.charCodeAt(0);
|
|
const BINUNICODE = 'X'.charCodeAt(0);
|
|
const EMPTY_TUPLE = ')'.charCodeAt(0);
|
|
const NEWTRUE = 0x88;
|
|
const NONE = 'N'.charCodeAt(0);
|
|
const BINFLOAT = 'G'.charCodeAt(0);
|
|
const TUPLE = 't'.charCodeAt(0);
|
|
const TUPLE1 = 0x85;
|
|
const TUPLE3 = 0x87;
|
|
// untested
|
|
const LONG_BINPUT = 'r'.charCodeAt(0);
|
|
const LIST = 'l'.charCodeAt(0);
|
|
const DICT = 'd'.charCodeAt(0);
|
|
const SETITEM = 's'.charCodeAt(0);
|
|
|
|
const scratch_buffer = new ArrayBuffer(8);
|
|
const scratch_bytes = new Uint8Array(scratch_buffer);
|
|
const big = new BigInt64Array(scratch_buffer);
|
|
const float64 = new Float64Array(scratch_buffer);
|
|
|
|
function read_uint4() {
|
|
const n =
|
|
bytebuffer[offset] +
|
|
bytebuffer[offset + 1] * 256 +
|
|
bytebuffer[offset + 2] * 65536 +
|
|
bytebuffer[offset + 3] * 16777216;
|
|
offset += 4;
|
|
return n;
|
|
}
|
|
function setitems(d, mark) {
|
|
for (let i = mark; i < stack.length; i += 2) {
|
|
d[stack[i]] = stack[i + 1];
|
|
}
|
|
stack.splice(mark, Infinity);
|
|
}
|
|
|
|
while (true) {
|
|
const opcode = bytebuffer[offset++];
|
|
switch (opcode) {
|
|
case PROTO:
|
|
{
|
|
const version = bytebuffer[offset++];
|
|
if (version < 2 || version > 4) {
|
|
throw new Error(`Unhandled version ${version}`);
|
|
}
|
|
}
|
|
break;
|
|
case APPEND:
|
|
{
|
|
const v = stack.pop();
|
|
stack.at(-1).push(v);
|
|
}
|
|
break;
|
|
case APPENDS:
|
|
{
|
|
const mark = marks.pop();
|
|
const arr = stack[mark - 1];
|
|
arr.push(...stack.splice(mark, Infinity));
|
|
}
|
|
break;
|
|
case LIST:
|
|
case TUPLE:
|
|
{
|
|
const mark = marks.pop();
|
|
stack.push([...stack.splice(mark, Infinity)]);
|
|
}
|
|
break;
|
|
case NEWFALSE:
|
|
stack.push(false);
|
|
break;
|
|
case NEWTRUE:
|
|
stack.push(true);
|
|
break;
|
|
case NONE:
|
|
stack.push(null);
|
|
break;
|
|
case BINGET:
|
|
stack.push(memo[bytebuffer[offset++]]);
|
|
break;
|
|
case BININT:
|
|
{
|
|
let i32 = read_uint4();
|
|
if (i32 > 0x7fffffff) {
|
|
i32 -= 0x100000000;
|
|
}
|
|
stack.push(i32);
|
|
}
|
|
break;
|
|
case BININT1:
|
|
stack.push(bytebuffer[offset++]);
|
|
break;
|
|
case BININT2:
|
|
{
|
|
const v = bytebuffer[offset] + bytebuffer[offset + 1] * 256;
|
|
stack.push(v);
|
|
offset += 2;
|
|
}
|
|
break;
|
|
case EMPTY_DICT:
|
|
stack.push({});
|
|
break;
|
|
case EMPTY_LIST:
|
|
stack.push([]);
|
|
break;
|
|
case FRAME:
|
|
offset += 8;
|
|
break;
|
|
case LONG1:
|
|
{
|
|
const s = bytebuffer[offset++];
|
|
if (s <= 8) {
|
|
for (let i = 0; i < s; i++) {
|
|
scratch_bytes[i] = bytebuffer[offset++];
|
|
}
|
|
const fill = scratch_bytes[s - 1] >= 128 ? 0xff : 0x0;
|
|
for (let i = s; i < 8; i++) {
|
|
scratch_bytes[i] = fill;
|
|
}
|
|
stack.push(Number(big[0]));
|
|
} else { // BigInt
|
|
let scratch_bytes_unbounded = [];
|
|
for (let i = 0; i < s; i++) {
|
|
scratch_bytes_unbounded.push(bytebuffer[offset++]);
|
|
}
|
|
|
|
// BigInt can only convert from unsigned hex, thus we need to
|
|
// convert from twos-complement if negative
|
|
const negative = scratch_bytes_unbounded[s - 1] >= 128;
|
|
if (negative) {
|
|
// implements scratch_bytes_unbounded = ~scratch_bytes_unbounded + 1
|
|
// byte-by-byte.
|
|
let carry = 1;
|
|
for (let i = 0; i < s; i++) {
|
|
const twos_complement = (0xff ^ scratch_bytes_unbounded[i]) + carry;
|
|
carry = twos_complement > 0xff ? 1 : 0;
|
|
scratch_bytes_unbounded[i] = 0xff & twos_complement;
|
|
}
|
|
}
|
|
|
|
const hex_str = Array.from(scratch_bytes_unbounded.reverse(), byte => {
|
|
return byte.toString(16).padStart(2, '0');
|
|
}).join('');
|
|
|
|
const big_int = negative ? -BigInt(`0x${hex_str}`) : BigInt(`0x${hex_str}`);
|
|
stack.push(big_int);
|
|
}
|
|
}
|
|
break;
|
|
case LONG_BINGET:
|
|
{
|
|
const idx = read_uint4();
|
|
stack.push(memo[idx]);
|
|
}
|
|
break;
|
|
case MARK:
|
|
marks.push(stack.length);
|
|
break;
|
|
case MEMOIZE:
|
|
memo[memo_id++] = stack.at(-1);
|
|
break;
|
|
case BINPUT:
|
|
memo[bytebuffer[offset++]] = stack.at(-1);
|
|
break;
|
|
case LONG_BINPUT:
|
|
memo[read_uint4()] = stack.at(-1);
|
|
break;
|
|
case SETITEMS:
|
|
{
|
|
const mark = marks.pop();
|
|
const d = stack[mark - 1];
|
|
setitems(d, mark);
|
|
}
|
|
break;
|
|
case SETITEM: {
|
|
const v = stack.pop();
|
|
const k = stack.pop();
|
|
stack.at(-1)[k] = v;
|
|
break;
|
|
}
|
|
case DICT:
|
|
{
|
|
const mark = marks.pop();
|
|
const d = {};
|
|
setitems(d, mark);
|
|
stack.push(d);
|
|
}
|
|
break;
|
|
case SHORT_BINUNICODE:
|
|
{
|
|
const n = bytebuffer[offset++];
|
|
stack.push(decoder.decode(new Uint8Array(buffer, offset, n)));
|
|
offset += n;
|
|
}
|
|
break;
|
|
case BINUNICODE:
|
|
{
|
|
const n = read_uint4();
|
|
stack.push(decoder.decode(new Uint8Array(buffer, offset, n)));
|
|
offset += n;
|
|
}
|
|
break;
|
|
case STOP:
|
|
return stack.pop();
|
|
case EMPTY_TUPLE:
|
|
stack.push([]);
|
|
break;
|
|
case TUPLE1:
|
|
stack.push([stack.pop()]);
|
|
break;
|
|
case TUPLE2:
|
|
stack.push(stack.splice(-2, Infinity));
|
|
break;
|
|
case TUPLE3:
|
|
stack.push(stack.splice(-3, Infinity));
|
|
break;
|
|
case BINFLOAT:
|
|
for (let i = 7; i >= 0; i--) {
|
|
// stored in big-endian order
|
|
scratch_bytes[i] = bytebuffer[offset++];
|
|
}
|
|
stack.push(float64[0]);
|
|
break;
|
|
default:
|
|
throw new Error(`UNKNOWN OPCODE: ${opcode}`);
|
|
}
|
|
}
|
|
}
|
|
|
|
function decode_base64(input) {
|
|
function decode_char(i, shift) {
|
|
const nChr = input.charCodeAt(i);
|
|
const r =
|
|
nChr > 64 && nChr < 91
|
|
? nChr - 65
|
|
: nChr > 96 && nChr < 123
|
|
? nChr - 71
|
|
: nChr > 47 && nChr < 58
|
|
? nChr + 4
|
|
: nChr === 43
|
|
? 62
|
|
: nChr === 47
|
|
? 63
|
|
: 0;
|
|
return r << shift;
|
|
}
|
|
const output = new Uint8Array((input.length / 4) * 3);
|
|
for (let i = 0, j = 0; i < input.length; i += 4, j += 3) {
|
|
const u24 =
|
|
decode_char(i, 18) +
|
|
decode_char(i + 1, 12) +
|
|
decode_char(i + 2, 6) +
|
|
decode_char(i + 3);
|
|
output[j] = u24 >> 16;
|
|
output[j + 1] = (u24 >> 8) & 0xff;
|
|
output[j + 2] = u24 & 0xff;
|
|
}
|
|
return output.buffer;
|
|
}
|
|
|
|
const kinds = {
|
|
'Active Memory Timeline': create_trace_view,
|
|
'Allocator State History': create_segment_view,
|
|
'Active Cached Segment Timeline': (dst, snapshot, device) =>
|
|
create_trace_view(dst, snapshot, device, true),
|
|
'Allocator Settings': create_settings_view,
|
|
};
|
|
|
|
const snapshot_cache = {};
|
|
const snapshot_to_loader = {};
|
|
const snapshot_to_url = {};
|
|
const selection_to_div = {};
|
|
|
|
const style = `
|
|
pre {
|
|
margin: 0px;
|
|
}
|
|
html, body {
|
|
height: 100%;
|
|
overflow: clip;
|
|
}`;
|
|
|
|
const head = d3.select('head');
|
|
head.append('style').text(style);
|
|
const body = d3.select('body');
|
|
const snapshot_select = body.append('select');
|
|
const view = body.append('select');
|
|
for (const x in kinds) {
|
|
view.append('option').text(x);
|
|
}
|
|
const gpu = body.append('select');
|
|
|
|
function unpickle_and_annotate(data) {
|
|
data = unpickle(data);
|
|
console.log(data);
|
|
annotate_snapshot(data);
|
|
return data;
|
|
}
|
|
|
|
function snapshot_change(f) {
|
|
const view_value = view.node().value;
|
|
let no_starting_gpu = gpu.node().value == '';
|
|
let device = Number(gpu.node().value);
|
|
const snapshot = snapshot_cache[f];
|
|
gpu.selectAll('option').remove();
|
|
const has_segments = {};
|
|
for (const s of snapshot.segments) {
|
|
has_segments[s.device] = true;
|
|
}
|
|
let device_valid = false;
|
|
let maxTraceLength = -1;
|
|
let defaultDevice = null;
|
|
for (const [i, trace] of snapshot.device_traces.entries()) {
|
|
if (trace.length > 0 || i in has_segments) {
|
|
gpu.append('option').text(i);
|
|
if (trace.length > maxTraceLength) {
|
|
maxTraceLength = trace.length;
|
|
defaultDevice = i;
|
|
}
|
|
if (i === device) {
|
|
device_valid = true;
|
|
gpu.node().selectedIndex = gpu.node().children.length - 1;
|
|
}
|
|
}
|
|
}
|
|
if (!device_valid) {
|
|
device = Number(gpu.node().value);
|
|
}
|
|
|
|
if (no_starting_gpu) {
|
|
device = defaultDevice;
|
|
gpu.node().value = device;
|
|
}
|
|
|
|
const key = [f, view_value, device];
|
|
if (!(key in selection_to_div)) {
|
|
selection_to_div[key] = d3.select('body').append('div');
|
|
kinds[view_value](selection_to_div[key], snapshot, device);
|
|
}
|
|
const selected_div = selection_to_div[key];
|
|
|
|
selected_div.attr('style', 'display: float; height: 100%');
|
|
}
|
|
|
|
function selected_change() {
|
|
for (const d of Object.values(selection_to_div)) {
|
|
d.attr('style', 'display: none; height: 100%');
|
|
}
|
|
const f = snapshot_select.node().value;
|
|
if (f === '') {
|
|
return;
|
|
}
|
|
if (!(f in snapshot_cache)) {
|
|
snapshot_to_loader[f](f);
|
|
} else {
|
|
snapshot_change(f);
|
|
}
|
|
}
|
|
|
|
snapshot_select.on('change', selected_change);
|
|
view.on('change', selected_change);
|
|
gpu.on('change', selected_change);
|
|
|
|
body.on('dragover', e => {
|
|
event.preventDefault();
|
|
});
|
|
|
|
body.on('drop', () => {
|
|
console.log(event.dataTransfer.files);
|
|
Array.from(event.dataTransfer.files).forEach(file => {
|
|
add_snapshot(file.name, unique_name => {
|
|
const reader = new FileReader();
|
|
reader.onload = e => {
|
|
finished_loading(unique_name, e.target.result);
|
|
};
|
|
reader.readAsArrayBuffer(file);
|
|
});
|
|
});
|
|
event.preventDefault();
|
|
snapshot_select.node().selectedIndex =
|
|
snapshot_select.node().options.length - 1;
|
|
selected_change();
|
|
});
|
|
|
|
selection_to_div[''] = body
|
|
.append('div')
|
|
.text(
|
|
'Drag and drop or select a file to load a local snapshot. No data from the snapshot is uploaded.',
|
|
);
|
|
|
|
const fileInput = body.append('input')
|
|
.attr('type', 'file')
|
|
.attr('multiple', true) // allow several snapshots at once
|
|
.style('margin-left', '8px')
|
|
.on('change', function () {
|
|
Array.from(this.files).forEach(file => {
|
|
add_snapshot(file.name, unique_name => {
|
|
const reader = new FileReader();
|
|
reader.onload = e =>
|
|
finished_loading(unique_name, e.target.result);
|
|
reader.readAsArrayBuffer(file);
|
|
});
|
|
});
|
|
this.value = null; // reset so the same file can be picked again
|
|
snapshot_select.node().selectedIndex =
|
|
snapshot_select.node().options.length - 1;
|
|
selected_change(); // refresh the UI
|
|
});
|
|
|
|
let next_unique_n = 1;
|
|
function add_snapshot(name, loader) {
|
|
if (name in snapshot_to_loader) {
|
|
name = `${name} (${next_unique_n++})`;
|
|
}
|
|
snapshot_select.append('option').text(name);
|
|
snapshot_to_loader[name] = loader;
|
|
}
|
|
|
|
function finished_loading(name, data) {
|
|
snapshot_cache[name] = unpickle_and_annotate(data);
|
|
snapshot_change(name);
|
|
}
|
|
|
|
export function add_remote_files(files) {
|
|
files.forEach(f =>
|
|
add_snapshot(f.name, unique_name => {
|
|
console.log('fetching', f.url);
|
|
fetch(f.url)
|
|
.then(x => x.arrayBuffer())
|
|
.then(data => finished_loading(unique_name, data));
|
|
}),
|
|
);
|
|
if (files.length > 0) {
|
|
selected_change();
|
|
}
|
|
}
|
|
|
|
export function add_local_files(files, view_value) {
|
|
view.node().value = view_value;
|
|
files.forEach(f =>
|
|
add_snapshot(f.name, unique_name => {
|
|
finished_loading(unique_name, decode_base64(f.base64));
|
|
}),
|
|
);
|
|
if (files.length > 0) {
|
|
selected_change();
|
|
}
|
|
}
|