pytorch

1
'use strict';
2

3
import * as d3 from "https://cdn.skypack.dev/d3@5";
4
import {axisLeft} from "https://cdn.skypack.dev/d3-axis@1";
5
import {scaleLinear} from "https://cdn.skypack.dev/d3-scale@1";
6
import {zoom, zoomIdentity} from "https://cdn.skypack.dev/d3-zoom@1";
7
import {brushX} from "https://cdn.skypack.dev/d3-brush@1";
8

9
const schemeTableau10 = [
10
  '#4e79a7',
11
  '#f28e2c',
12
  '#e15759',
13
  '#76b7b2',
14
  '#59a14f',
15
  '#edc949',
16
  '#af7aa1',
17
  '#ff9da7',
18
  '#9c755f',
19
  '#bab0ab',
20
];
21

22
function version_space() {
23
  const version = {};
24
  return (addr, increment) => {
25
    if (!(addr in version)) {
26
      version[addr] = 0;
27
    }
28
    const r = version[addr];
29
    if (increment) {
30
      version[addr]++;
31
    }
32
    return r;
33
  };
34
}
35

36
function Segment(addr, size, stream, frames, version) {
37
  return {addr, size, stream, version, frames};
38
}
39

40
function Block(addr, size, requested_size, frames, free_requested, version) {
41
  return {addr, size, requested_size, frames, free_requested, version};
42
}
43

44
function EventSelector(outer, events, stack_info, memory_view) {
45
  const events_div = outer
46
    .append('div')
47
    .attr(
48
      'style',
49
      'grid-column: 1; grid-row: 1; overflow: auto; font-family: monospace',
50
    );
51

52
  const events_selection = events_div
53
    .selectAll('pre')
54
    .data(events)
55
    .enter()
56
    .append('pre')
57
    .text(e => formatEvent(e))
58
    .attr('style', '');
59

60
  let selected_event_idx = null;
61

62
  const es = {
63
    select(idx) {
64
      if (selected_event_idx !== null) {
65
        const selected_event = d3.select(
66
          events_div.node().children[selected_event_idx],
67
        );
68
        selected_event.attr('style', '');
69
      }
70
      if (idx !== null) {
71
        const div = d3.select(events_div.node().children[idx]);
72
        div.attr('style', `background-color: ${schemeTableau10[5]}`);
73
        const [reserved, allocated] = memory_view.draw(idx);
74
        const enter = () => eventStack(div.datum(), allocated, reserved);
75
        stack_info.highlight(enter);
76
        div.node().scrollIntoViewIfNeeded(false);
77
      } else {
78
        memory_view.draw(0);
79
      }
80
      selected_event_idx = idx;
81
    },
82
  };
83
  d3.select('body').on('keydown', _e => {
84
    const key = d3.event.key;
85
    const actions = {ArrowDown: 1, ArrowUp: -1};
86
    if (selected_event_idx !== null && key in actions) {
87
      const new_idx = selected_event_idx + actions[key];
88
      es.select(Math.max(0, Math.min(new_idx, events.length - 1)));
89
      d3.event.preventDefault();
90
    }
91
  });
92

93
  stack_info.register(
94
    events_selection,
95
    t => eventStack(t.datum()),
96
    _t => {},
97
    d => es.select(d.datum().idx),
98
  );
99

100
  return es;
101
}
102

103
function formatSize(num) {
104
  const orig = num;
105
  // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size
106
  const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'];
107
  for (const unit of units) {
108
    if (Math.abs(num) < 1024.0) {
109
      return `${num.toFixed(1)}${unit}B (${orig} bytes)`;
110
    }
111
    num /= 1024.0;
112
  }
113
  return `${num.toFixed(1)}YiB`;
114
}
115
function formatAddr(event) {
116
  const prefix = event.action.startsWith('segment') ? 's' : 'b';
117
  return `${prefix}${event.addr.toString(16)}_${event.version}`;
118
}
119
function formatEvent(event) {
120
  const stream =
121
    event.stream === null ? '' : `\n              (stream ${event.stream})`;
122
  switch (event.action) {
123
    case 'oom':
124
      return `OOM (requested ${formatSize(event.size)}, CUDA has ${formatSize(
125
        event.device_free,
126
      )} memory free)${stream}`;
127
    case 'snapshot':
128
      return 'snapshot';
129
    default:
130
      return `${event.action.padEnd(14)} ${formatAddr(event).padEnd(
131
        18,
132
      )} ${formatSize(event.size)}${stream}`;
133
  }
134
}
135

136
function eventStack(e, allocated, reserved) {
137
  let event = formatEvent(e);
138
  if (reserved !== undefined) {
139
    event = `(${formatSize(allocated)} allocated / ${formatSize(
140
      reserved,
141
    )} reserved)\n${event}`;
142
  }
143
  return event + '\n' + format_frames(e.frames);
144
}
145

146
function hashCode(num) {
147
  const numStr = num.toString();
148
  let hash = 0;
149
  for (let i = 0; i < numStr.length; i++) {
150
    const charCode = numStr.charCodeAt(i);
151
    hash = (hash << 5) - hash + charCode;
152
    hash = hash & hash; // Convert to 32-bit integer
153
  }
154
  return hash;
155
}
156

157
function addStroke(d) {
158
  d.attr('stroke', 'red')
159
    .attr('stroke-width', '2')
160
    .attr('vector-effect', 'non-scaling-stroke');
161
}
162

163
function removeStroke(d) {
164
  d.attr('stroke', '');
165
}
166

167
function calculate_fragmentation(blocks, sorted_segments) {
168
  const sorted_blocks = Object.values(blocks).sort((a, b) => a.addr - b.addr);
169
  let block_i = 0;
170
  let total_size = 0;
171
  let sum_squared_free = 0;
172
  for (const seg of sorted_segments) {
173
    let addr = seg.addr;
174
    total_size += seg.size;
175
    while (
176
      block_i < sorted_blocks.length &&
177
      sorted_blocks[block_i].addr < seg.addr + seg.size
178
    ) {
179
      const block = sorted_blocks[block_i];
180
      if (block.addr > addr) {
181
        sum_squared_free += (block.addr - addr) ** 2;
182
      }
183
      addr = block.addr + block.size;
184
      block_i += 1;
185
    }
186
    if (addr < seg.addr + seg.size) {
187
      sum_squared_free += (seg.addr + seg.size - addr) ** 2;
188
    }
189
  }
190
  console.log(sum_squared_free / total_size ** 2);
191
}
192

193
function MemoryView(outer, stack_info, snapshot, device) {
194
  const svg = outer
195
    .append('svg')
196
    .attr('style', 'grid-column: 2; grid-row: 1; width: 100%; height: 100%;')
197
    .attr('viewBox', '0 0 200 100')
198
    .attr('preserveAspectRatio', 'xMinYMin meet');
199
  const g = svg.append('g');
200
  const seg_zoom = zoom();
201
  seg_zoom.on('zoom', () => {
202
    g.attr('transform', d3.event.transform);
203
  });
204
  svg.call(seg_zoom);
205

206
  const sorted_segments = [];
207
  const block_map = {};
208
  for (const seg of snapshot.segments) {
209
    if (seg.device !== device) {
210
      continue;
211
    }
212
    sorted_segments.push(
213
      Segment(
214
        seg.address,
215
        seg.total_size,
216
        seg.stream,
217
        seg.frames || [],
218
        seg.version,
219
      ),
220
    );
221
    for (const b of seg.blocks) {
222
      if (b.state !== 'active_pending_free' && b.state !== 'active_allocated') {
223
        continue;
224
      }
225
      block_map[b.addr] = Block(
226
        b.addr,
227
        b.size,
228
        b.requested_size,
229
        b.frames,
230
        b.state === 'active_pending_free',
231
        b.version,
232
      );
233
    }
234
  }
235
  sorted_segments.sort((x, y) => x.addr - y.addr);
236

237
  function simulate_memory(idx) {
238
    // create a copy of segments because we edit size properties below
239
    const l_segments = sorted_segments.map(x => {
240
      return {...x};
241
    });
242
    const l_block_map = {...block_map};
243

244
    function map_segment(merge, seg) {
245
      let idx = l_segments.findIndex(e => e.addr > seg.addr);
246
      if (!merge) {
247
        l_segments.splice(idx, 0, seg);
248
        return;
249
      }
250
      if (idx === -1) {
251
        idx = l_segments.length;
252
      }
253
      l_segments.splice(idx, 0, seg);
254
      if (idx + 1 < l_segments.length) {
255
        const next = l_segments[idx + 1];
256
        if (seg.addr + seg.size === next.addr && seg.stream === next.stream) {
257
          seg.size += next.size;
258
          l_segments.splice(idx + 1, 1);
259
        }
260
      }
261
      if (idx > 0) {
262
        const prev = l_segments[idx - 1];
263
        if (prev.addr + prev.size === seg.addr && prev.stream === seg.stream) {
264
          prev.size += seg.size;
265
          l_segments.splice(idx, 1);
266
        }
267
      }
268
    }
269
    function unmap_segment(merge, seg) {
270
      if (!merge) {
271
        l_segments.splice(
272
          l_segments.findIndex(x => x.addr === seg.addr),
273
          1,
274
        );
275
        return;
276
      }
277
      const seg_end = seg.addr + seg.size;
278
      const idx = l_segments.findIndex(
279
        e => e.addr <= seg.addr && seg_end <= e.addr + e.size,
280
      );
281
      const existing = l_segments[idx];
282
      const existing_end = existing.addr + existing.size;
283
      if (existing.addr === seg.addr) {
284
        existing.addr += seg.size;
285
        existing.size -= seg.size;
286
        if (existing.size === 0) {
287
          l_segments.splice(idx, 1);
288
        }
289
      } else if (existing_end === seg_end) {
290
        existing.size -= seg.size;
291
      } else {
292
        existing.size = seg.addr - existing.addr;
293
        seg.addr = seg_end;
294
        seg.size = existing_end - seg_end;
295
        l_segments.splice(idx + 1, 0, seg);
296
      }
297
    }
298
    const events = snapshot.device_traces[device];
299
    for (let i = events.length - 1; i > idx; i--) {
300
      const event = events[i];
301
      switch (event.action) {
302
        case 'free':
303
          l_block_map[event.addr] = Block(
304
            event.addr,
305
            event.size,
306
            event.size,
307
            event.frames,
308
            false,
309
            event.version,
310
          );
311
          break;
312
        case 'free_requested':
313
          l_block_map[event.addr].free_requested = false;
314
          break;
315
        case 'free_completed':
316
          l_block_map[event.addr] = Block(
317
            event.addr,
318
            event.size,
319
            event.size,
320
            event.frames,
321
            true,
322
            event.version,
323
          );
324
          break;
325
        case 'alloc':
326
          delete l_block_map[event.addr];
327
          break;
328
        case 'segment_free':
329
        case 'segment_unmap':
330
          map_segment(
331
            event.action === 'segment_unmap',
332
            Segment(
333
              event.addr,
334
              event.size,
335
              event.stream,
336
              event.frames,
337
              event.version,
338
            ),
339
          );
340
          break;
341
        case 'segment_alloc':
342
        case 'segment_map':
343
          unmap_segment(
344
            event.action === 'segment_map',
345
            Segment(
346
              event.addr,
347
              event.size,
348
              event.stream,
349
              event.frames,
350
              event.version,
351
            ),
352
          );
353
          break;
354
        case 'oom':
355
          break;
356
        default:
357
          break;
358
      }
359
    }
360
    const new_blocks = Object.values(l_block_map);
361
    return [l_segments, new_blocks];
362
  }
363

364
  return {
365
    draw(idx) {
366
      const [segments_unsorted, blocks] = simulate_memory(idx);
367
      g.selectAll('g').remove();
368

369
      const segment_d = g.append('g');
370
      const block_g = g.append('g');
371
      const block_r = g.append('g');
372

373
      segment_d.selectAll('rect').remove();
374
      block_g.selectAll('rect').remove();
375
      block_r.selectAll('rect').remove();
376
      const segments = [...segments_unsorted].sort((x, y) =>
377
        x.size === y.size ? x.addr - y.addr : x.size - y.size,
378
      );
379

380
      const segments_by_addr = [...segments].sort((x, y) => x.addr - y.addr);
381

382
      const max_size = segments.length === 0 ? 0 : segments.at(-1).size;
383

384
      const xScale = scaleLinear().domain([0, max_size]).range([0, 200]);
385
      const padding = xScale.invert(1);
386

387
      let cur_row = 0;
388
      let cur_row_size = 0;
389
      for (const seg of segments) {
390
        seg.occupied = 0;
391
        seg.internal_free = 0;
392
        if (cur_row_size + seg.size > max_size) {
393
          cur_row_size = 0;
394
          cur_row += 1;
395
        }
396
        seg.offset = cur_row_size;
397
        seg.row = cur_row;
398
        cur_row_size += seg.size + padding;
399
      }
400

401
      const num_rows = cur_row + 1;
402

403
      const yScale = scaleLinear().domain([0, num_rows]).range([0, 100]);
404

405
      const segments_selection = segment_d
406
        .selectAll('rect')
407
        .data(segments)
408
        .enter()
409
        .append('rect')
410
        .attr('x', x => xScale(x.offset))
411
        .attr('y', x => yScale(x.row))
412
        .attr('width', x => xScale(x.size))
413
        .attr('height', yScale(4 / 5))
414
        .attr('stroke', 'black')
415
        .attr('stroke-width', '1')
416
        .attr('vector-effect', 'non-scaling-stroke')
417
        .attr('fill', 'white');
418

419
      stack_info.register(
420
        segments_selection,
421
        d => {
422
          addStroke(d);
423
          const t = d.datum();
424
          const free = t.size - t.occupied;
425
          let internal = '';
426
          if (t.internal_free > 0) {
427
            internal = ` (${(t.internal_free / free) * 100}% internal)`;
428
          }
429
          return (
430
            `s${t.addr.toString(16)}_${t.version}: segment ${formatSize(
431
              t.size,
432
            )} allocated, ` +
433
            `${formatSize(free)} free${internal} (stream ${
434
              t.stream
435
            })\n${format_frames(t.frames)}`
436
          );
437
        },
438
        d => {
439
          d.attr('stroke', 'black')
440
            .attr('stroke-width', '1')
441
            .attr('vector-effect', 'non-scaling-stroke');
442
        },
443
      );
444

445
      function find_segment(addr) {
446
        let left = 0;
447
        let right = segments_by_addr.length - 1;
448
        while (left <= right) {
449
          const mid = Math.floor((left + right) / 2);
450
          if (addr < segments_by_addr[mid].addr) {
451
            right = mid - 1;
452
          } else if (
453
            addr >=
454
            segments_by_addr[mid].addr + segments_by_addr[mid].size
455
          ) {
456
            left = mid + 1;
457
          } else {
458
            return segments_by_addr[mid];
459
          }
460
        }
461
        return null;
462
      }
463

464
      for (const b of blocks) {
465
        b.segment = find_segment(b.addr);
466
        b.segment.occupied += b.requested_size;
467
        b.segment.internal_free += b.size - b.requested_size;
468
      }
469

470
      const block_selection = block_g
471
        .selectAll('rect')
472
        .data(blocks)
473
        .enter()
474
        .append('rect')
475
        .attr('x', x => xScale(x.segment.offset + (x.addr - x.segment.addr)))
476
        .attr('y', x => yScale(x.segment.row))
477
        .attr('width', x => xScale(x.requested_size))
478
        .attr('height', yScale(4 / 5))
479
        .attr('fill', (x, _i) =>
480
          x.free_requested
481
            ? 'red'
482
            : schemeTableau10[
483
                Math.abs(hashCode(x.addr)) % schemeTableau10.length
484
              ],
485
        );
486

487
      stack_info.register(
488
        block_selection,
489
        d => {
490
          addStroke(d);
491
          const t = d.datum();
492
          let requested = '';
493
          if (t.free_requested) {
494
            requested = ' (block freed but waiting due to record_stream)';
495
          }
496
          return (
497
            `b${t.addr.toString(16)}_${t.version} ` +
498
            `${formatSize(t.requested_size)} allocation${requested} (stream ${
499
              t.segment.stream
500
            })\n` +
501
            format_frames(t.frames)
502
          );
503
        },
504
        removeStroke,
505
      );
506

507
      const free_selection = block_r
508
        .selectAll('rect')
509
        .data(blocks)
510
        .enter()
511
        .append('rect')
512
        .attr('x', x =>
513
          xScale(
514
            x.segment.offset + (x.addr - x.segment.addr) + x.requested_size,
515
          ),
516
        )
517
        .attr('y', x => yScale(x.segment.row))
518
        .attr('width', x => xScale(x.size - x.requested_size))
519
        .attr('height', yScale(4 / 5))
520
        .attr('fill', (_x, _i) => 'red');
521

522
      stack_info.register(
523
        free_selection,
524
        d => {
525
          addStroke(d);
526
          const t = d.datum();
527
          return (
528
            `Free space lost due to rounding ${formatSize(
529
              t.size - t.requested_size,
530
            )}` +
531
            ` (stream ${t.segment.stream})\n` +
532
            format_frames(t.frames)
533
          );
534
        },
535
        removeStroke,
536
      );
537

538
      const reserved = segments.reduce((x, y) => x + y.size, 0);
539
      const allocated = blocks.reduce((x, y) => x + y.requested_size, 0);
540
      return [reserved, allocated];
541
    },
542
  };
543
}
544

545
function StackInfo(outer) {
546
  const stack_trace = outer
547
    .append('pre')
548
    .attr('style', 'grid-column: 1 / 3; grid-row: 2; overflow: auto');
549
  let selected = {
550
    enter: () => {
551
      stack_trace.text('');
552
    },
553
    leave: () => {},
554
  };
555
  return {
556
    register(dom, enter, leave = _e => {}, select = _e => {}) {
557
      dom
558
        .on('mouseover', _e => {
559
          selected.leave();
560
          stack_trace.text(enter(d3.select(d3.event.target)));
561
        })
562
        .on('mousedown', _e => {
563
          const obj = d3.select(d3.event.target);
564
          selected = {
565
            enter: () => stack_trace.text(enter(obj)),
566
            leave: () => leave(obj),
567
          };
568
          select(obj);
569
        })
570
        .on('mouseleave', _e => {
571
          leave(d3.select(d3.event.target));
572
          selected.enter();
573
        });
574
    },
575
    highlight(enter, leave = () => {}) {
576
      selected = {enter: () => stack_trace.text(enter()), leave};
577
      selected.enter();
578
    },
579
  };
580
}
581

582
function create_segment_view(dst, snapshot, device) {
583
  const outer = dst
584
    .append('div')
585
    .attr(
586
      'style',
587
      'display: grid; grid-template-columns: 1fr 2fr; grid-template-rows: 2fr 1fr; height: 100%; gap: 10px',
588
    );
589

590
  const events = snapshot.device_traces[device];
591
  const stack_info = StackInfo(outer);
592
  const memory_view = MemoryView(outer, stack_info, snapshot, device);
593
  const event_selector = EventSelector(outer, events, stack_info, memory_view);
594

595
  window.requestAnimationFrame(function () {
596
    event_selector.select(events.length > 0 ? events.length - 1 : null);
597
  });
598
}
599

600
function annotate_snapshot(snapshot) {
601
  snapshot.segment_version = version_space();
602
  snapshot.block_version = version_space();
603
  snapshot.categories = [];
604
  const empty_list = [];
605
  let next_stream = 1;
606
  const stream_names = {0: 0};
607
  function stream_name(s) {
608
    if (!(s in stream_names)) {
609
      stream_names[s] = next_stream++;
610
    }
611
    return stream_names[s];
612
  }
613
  const new_traces = [];
614
  for (const device_trace of snapshot.device_traces) {
615
    const new_trace = [];
616
    new_traces.push(new_trace);
617
    for (const t of device_trace) {
618
      if (!('frames' in t)) {
619
        t.frames = empty_list;
620
      }
621
      // set unique version for each time an address is used
622
      // so that ctrl-f can be used to search for the beginning
623
      // and end of allocations and segments
624
      t.stream = stream_name(t.stream);
625
      switch (t.action) {
626
        case 'free_completed':
627
          t.version = snapshot.block_version(t.addr, true);
628
          if (new_trace.length > 0) {
629
            // elide free_requested/free_completed into a single event
630
            const prev = new_trace.at(-1);
631
            if (prev.action === 'free_requested' && prev.addr === t.addr) {
632
              prev.action = 'free';
633
              continue;
634
            }
635
          }
636
          break;
637
        case 'free_requested':
638
        case 'alloc':
639
          t.version = snapshot.block_version(t.addr, false);
640
          break;
641
        case 'segment_free':
642
        case 'segment_unmap':
643
          t.version = snapshot.segment_version(t.addr, true);
644
          break;
645
        case 'segment_alloc':
646
        case 'segment_map':
647
          t.version = snapshot.segment_version(t.addr, false);
648
          break;
649
        default:
650
          break;
651
      }
652
      if ('category' in t && !snapshot.categories.includes(t.category)) {
653
        snapshot.categories.push(t.category);
654
      }
655
      t.idx = new_trace.length;
656
      new_trace.push(t);
657
    }
658
  }
659
  snapshot.device_traces = new_traces;
660
  // if every event was on the default stream, we elide stream printing
661
  if (next_stream == 1) {
662
    for (const device_trace of snapshot.device_traces) {
663
      for (const t of device_trace) {
664
        t.stream = null;
665
      }
666
    }
667
  }
668

669
  for (const seg of snapshot.segments) {
670
    seg.stream = stream_name(seg.stream);
671
    seg.version = snapshot.segment_version(seg.address, false);
672
    let addr = seg.address;
673
    for (const b of seg.blocks) {
674
      b.addr = addr;
675
      if (!('frames' in b)) {
676
        // legacy format where 'requested_size' may be missing
677
        // and frames might be in history rather than directly on block
678
        if ('history' in b) {
679
          b.frames = b.history[0].frames || empty_list;
680
          b.requested_size = b.requested_size || b.history[0].real_size;
681
        } else {
682
          b.frames = empty_list;
683
          b.requested_size = b.requested_size || b.size;
684
        }
685
      }
686
      b.version = snapshot.block_version(b.addr, false);
687
      addr += b.size;
688
    }
689
  }
690

691
  if (
692
    snapshot.categories.length > 0 &&
693
    !snapshot.categories.includes('unknown')
694
  ) {
695
    snapshot.categores.push('unknown');
696
  }
697
}
698

699
function elideRepeats(frames) {
700
  const result = [];
701
  const length = frames.length;
702
  for (let i = 0; i < length; ) {
703
    let j = i + 1;
704
    const f = frames[i];
705
    while (j < length && f === frames[j]) {
706
      j++;
707
    }
708
    switch (j - i) {
709
      case 1:
710
        result.push(f);
711
        break;
712
      case 2:
713
        result.push(f, f);
714
        break;
715
      default:
716
        result.push(f, `<repeats ${j - i - 1} times>`);
717
        break;
718
    }
719
    i = j;
720
  }
721
  return result;
722
}
723
function frameFilter({name, filename}) {
724
  const omitFunctions = [
725
    'unwind::unwind',
726
    'CapturedTraceback::gather',
727
    'gather_with_cpp',
728
    '_start',
729
    '__libc_start_main',
730
    'PyEval_',
731
    'PyObject_',
732
    'PyFunction_',
733
  ];
734

735
  const omitFilenames = [
736
    'core/boxing',
737
    '/Register',
738
    '/Redispatch',
739
    'pythonrun.c',
740
    'Modules/main.c',
741
    'Objects/call.c',
742
    'Objects/methodobject.c',
743
    'pycore_ceval.h',
744
    'ceval.c',
745
    'cpython/abstract.h',
746
  ];
747

748
  for (const of of omitFunctions) {
749
    if (name.includes(of)) {
750
      return false;
751
    }
752
  }
753

754
  for (const of of omitFilenames) {
755
    if (filename.includes(of)) {
756
      return false;
757
    }
758
  }
759

760
  return true;
761
}
762

763
function format_frames(frames) {
764
  if (frames.length === 0) {
765
    return `<block was allocated before _record_history was enabled>`;
766
  }
767
  const frame_strings = frames
768
    .filter(frameFilter)
769
    .map(f => `${f.filename}:${f.line}:${f.name}`);
770
  return elideRepeats(frame_strings).join('\n');
771
}
772

773
function process_alloc_data(snapshot, device, plot_segments, max_entries) {
774
  const elements = [];
775
  const initially_allocated = [];
776
  const actions = [];
777
  const addr_to_alloc = {};
778

779
  const alloc = plot_segments ? 'segment_alloc' : 'alloc';
780
  const [free, free_completed] = plot_segments
781
    ? ['segment_free', 'segment_free']
782
    : ['free', 'free_completed'];
783
  for (const e of snapshot.device_traces[device]) {
784
    switch (e.action) {
785
      case alloc:
786
        elements.push(e);
787
        addr_to_alloc[e.addr] = elements.length - 1;
788
        actions.push(elements.length - 1);
789
        break;
790
      case free:
791
      case free_completed:
792
        if (e.addr in addr_to_alloc) {
793
          actions.push(addr_to_alloc[e.addr]);
794
          delete addr_to_alloc[e.addr];
795
        } else {
796
          elements.push(e);
797
          initially_allocated.push(elements.length - 1);
798
          actions.push(elements.length - 1);
799
        }
800
        break;
801
      default:
802
        break;
803
    }
804
  }
805
  for (const seg of snapshot.segments) {
806
    if (seg.device !== device) {
807
      continue;
808
    }
809
    if (plot_segments) {
810
      if (!(seg.address in addr_to_alloc)) {
811
        const element = {
812
          action: 'alloc',
813
          addr: seg.address,
814
          size: seg.total_size,
815
          frames: [],
816
          stream: seg.stream,
817
          version: seg.version,
818
        };
819
        elements.push(element);
820
        initially_allocated.push(elements.length - 1);
821
      }
822
    } else {
823
      for (const b of seg.blocks) {
824
        if (b.state === 'active_allocated' && !(b.addr in addr_to_alloc)) {
825
          const element = {
826
            action: 'alloc',
827
            addr: b.addr,
828
            size: b.requested_size,
829
            frames: b.frames,
830
            stream: seg.stream,
831
            version: b.version,
832
          };
833
          elements.push(element);
834
          initially_allocated.push(elements.length - 1);
835
        }
836
      }
837
    }
838
  }
839
  initially_allocated.reverse();
840
  // if there are no actions, the graph will be blank,
841
  // but if there are existing allocations we do not want to hide them
842
  // by having just one allocate action it will show a flat graph with all segments
843
  if (actions.length === 0 && initially_allocated.length > 0) {
844
    actions.push(initially_allocated.pop());
845
  }
846

847
  const current = [];
848
  const current_data = [];
849
  const data = [];
850
  let max_size = 0;
851

852
  let total_mem = 0;
853
  let total_summarized_mem = 0;
854
  let timestep = 0;
855

856
  const max_at_time = [];
857

858
  const summarized_mem = {
859
    elem: 'summarized',
860
    timesteps: [],
861
    offsets: [total_mem],
862
    size: [],
863
    color: 0,
864
  };
865
  const summarized_elems = {};
866

867
  function advance(n) {
868
    summarized_mem.timesteps.push(timestep);
869
    summarized_mem.offsets.push(total_mem);
870
    summarized_mem.size.push(total_summarized_mem);
871
    timestep += n;
872
    for (let i = 0; i < n; i++) {
873
      max_at_time.push(total_mem + total_summarized_mem);
874
    }
875
  }
876

877
  const sizes = elements
878
    .map((x, i) => [x.size, i])
879
    .sort(([x, _xi], [y, _yi]) => y - x);
880

881
  const draw_elem = {};
882
  for (const [_s, e] of sizes.slice(0, max_entries)) {
883
    draw_elem[e] = true;
884
  }
885

886
  function add_allocation(elem) {
887
    const element_obj = elements[elem];
888
    const size = element_obj.size;
889
    current.push(elem);
890
    let color = elem;
891
    if (snapshot.categories.length > 0) {
892
      color = snapshot.categories.indexOf(element_obj.category || 'unknown');
893
    }
894
    const e = {
895
      elem,
896
      timesteps: [timestep],
897
      offsets: [total_mem],
898
      size,
899
      color,
900
    };
901
    current_data.push(e);
902
    data.push(e);
903
    total_mem += size;
904
    element_obj.max_allocated_mem = total_mem + total_summarized_mem;
905
  }
906

907
  for (const elem of initially_allocated) {
908
    if (elem in draw_elem) {
909
      add_allocation(elem);
910
    } else {
911
      total_summarized_mem += elements[elem].size;
912
      summarized_elems[elem] = true;
913
    }
914
  }
915

916
  for (const elem of actions) {
917
    const size = elements[elem].size;
918
    if (!(elem in draw_elem)) {
919
      if (elem in summarized_elems) {
920
        advance(1);
921
        total_summarized_mem -= size;
922
        summarized_elems[elem] = null;
923
      } else {
924
        total_summarized_mem += size;
925
        summarized_elems[elem] = true;
926
        advance(1);
927
      }
928
      continue;
929
    }
930
    const idx = current.findLastIndex(x => x === elem);
931
    // first time we see an action we add it
932
    // second time we remove it
933
    if (idx === -1) {
934
      add_allocation(elem);
935
      advance(1);
936
    } else {
937
      advance(1);
938
      const removed = current_data[idx];
939
      removed.timesteps.push(timestep);
940
      removed.offsets.push(removed.offsets.at(-1));
941
      current.splice(idx, 1);
942
      current_data.splice(idx, 1);
943

944
      if (idx < current.length) {
945
        for (let j = idx; j < current.length; j++) {
946
          const e = current_data[j];
947
          e.timesteps.push(timestep);
948
          e.offsets.push(e.offsets.at(-1));
949
          e.timesteps.push(timestep + 3);
950
          e.offsets.push(e.offsets.at(-1) - size);
951
        }
952
        advance(3);
953
      }
954
      total_mem -= size;
955
    }
956
    max_size = Math.max(total_mem + total_summarized_mem, max_size);
957
  }
958

959
  for (const elem of current_data) {
960
    elem.timesteps.push(timestep);
961
    elem.offsets.push(elem.offsets.at(-1));
962
  }
963
  data.push(summarized_mem);
964

965
  return {
966
    max_size,
967
    allocations_over_time: data,
968
    max_at_time,
969
    summarized_mem,
970
    elements_length: elements.length,
971
    context_for_id: id => {
972
      const elem = elements[id];
973
      let text = `Addr: ${formatAddr(elem)}`;
974
      text = `${text}, Size: ${formatSize(elem.size)} allocation`;
975
      text = `${text}, Total memory used after allocation: ${formatSize(
976
        elem.max_allocated_mem,
977
      )}`;
978
      if (elem.stream !== null) {
979
        text = `${text}, stream ${elem.stream}`;
980
      }
981
      if (!elem.action.includes('alloc')) {
982
        text = `${text}\nalloc not recorded, stack trace for free:`;
983
      }
984
      text = `${text}\n${format_frames(elem.frames)}`;
985
      return text;
986
    },
987
  };
988
}
989

990
function MemoryPlot(
991
  svg,
992
  data,
993
  left_pad,
994
  width,
995
  height,
996
  colors = schemeTableau10,
997
) {
998
  function format_points(d) {
999
    const size = d.size;
1000
    const xs = d.timesteps.map(t => xscale(t));
1001
    const bottom = d.offsets.map(t => yscale(t));
1002
    const m = Array.isArray(size)
1003
      ? (t, i) => yscale(t + size[i])
1004
      : t => yscale(t + size);
1005
    const top = d.offsets.map(m);
1006
    const p0 = xs.map((x, i) => `${x},${bottom[i]}`);
1007
    const p1 = xs.map((x, i) => `${x},${top[i]}`).reverse();
1008
    return `${p0.join(' ')} ${p1.join(' ')}`;
1009
  }
1010

1011
  const max_timestep = data.max_at_time.length;
1012
  const max_size = data.max_size;
1013

1014
  const plot_width = width - left_pad;
1015
  const plot_height = height;
1016

1017
  const yscale = scaleLinear().domain([0, max_size]).range([plot_height, 0]);
1018
  const yaxis = axisLeft(yscale).tickFormat(d3.format('.3s'));
1019
  const xscale = scaleLinear().domain([0, max_timestep]).range([0, plot_width]);
1020
  const plot_coordinate_space = svg
1021
    .append('g')
1022
    .attr('transform', `translate(${left_pad}, ${0})`);
1023
  const plot_outer = plot_coordinate_space.append('g');
1024

1025
  function view_rect(a) {
1026
    return a
1027
      .append('rect')
1028
      .attr('x', 0)
1029
      .attr('y', 0)
1030
      .attr('width', plot_width)
1031
      .attr('height', plot_height)
1032
      .attr('fill', 'white');
1033
  }
1034

1035
  view_rect(plot_outer);
1036

1037
  const cp = svg.append('clipPath').attr('id', 'clip');
1038
  view_rect(cp);
1039
  plot_outer.attr('clip-path', 'url(#clip)');
1040

1041
  const zoom_group = plot_outer.append('g');
1042
  const scrub_group = zoom_group.append('g');
1043

1044
  const plot = scrub_group
1045
    .selectAll('polygon')
1046
    .data(data.allocations_over_time)
1047
    .enter()
1048
    .append('polygon')
1049
    .attr('points', format_points)
1050
    .attr('fill', d => colors[d.color % colors.length]);
1051

1052
  const axis = plot_coordinate_space.append('g').call(yaxis);
1053

1054
  function handleZoom() {
1055
    const t = d3.event.transform;
1056
    zoom_group.attr('transform', t);
1057
    axis.call(yaxis.scale(d3.event.transform.rescaleY(yscale)));
1058
  }
1059

1060
  const thezoom = zoom().on('zoom', handleZoom);
1061
  plot_outer.call(thezoom);
1062

1063
  return {
1064
    select_window: (stepbegin, stepend, max) => {
1065
      const begin = xscale(stepbegin);
1066
      const size = xscale(stepend) - xscale(stepbegin);
1067
      const scale = plot_width / size;
1068
      const translate = -begin;
1069
      const yscale = max_size / max;
1070
      scrub_group.attr(
1071
        'transform',
1072
        `scale(${scale / yscale}, 1) translate(${translate}, 0)`,
1073
      );
1074
      plot_outer.call(
1075
        thezoom.transform,
1076
        zoomIdentity
1077
          .scale(yscale)
1078
          .translate(0, -(plot_height - plot_height / yscale)),
1079
      );
1080
    },
1081
    set_delegate: delegate => {
1082
      plot
1083
        .on('mouseover', function (_e, _d) {
1084
          delegate.set_selected(d3.select(this));
1085
        })
1086
        .on('mousedown', function (_e, _d) {
1087
          delegate.default_selected = d3.select(this);
1088
        })
1089
        .on('mouseleave', function (_e, _d) {
1090
          delegate.set_selected(delegate.default_selected);
1091
        });
1092
    },
1093
  };
1094
}
1095

1096
function ContextViewer(text, data) {
1097
  let current_selected = null;
1098

1099
  return {
1100
    default_selected: null,
1101
    set_selected: d => {
1102
      if (current_selected !== null) {
1103
        current_selected.attr('stroke', null).attr('stroke-width', null);
1104
      }
1105
      if (d === null) {
1106
        text.text('');
1107
      } else {
1108
        const dd = d.datum();
1109
        if (dd.elem === 'summarized') {
1110
          text.html(
1111
            'Small tensors that were not plotted to cutdown on render time.\n' +
1112
              'Use detail slider to see smaller allocations.',
1113
          );
1114
        } else {
1115
          text.text(`${dd.elem} ${data.context_for_id(dd.elem)}`);
1116
        }
1117
        d.attr('stroke', 'black')
1118
          .attr('stroke-width', 1)
1119
          .attr('vector-effect', 'non-scaling-stroke');
1120
      }
1121
      current_selected = d;
1122
    },
1123
  };
1124
}
1125

1126
function MiniMap(mini_svg, plot, data, left_pad, width, height = 70) {
1127
  const max_at_time = data.max_at_time;
1128
  const plot_width = width - left_pad;
1129
  const yscale = scaleLinear().domain([0, data.max_size]).range([height, 0]);
1130
  const minixscale = scaleLinear()
1131
    .domain([0, max_at_time.length])
1132
    .range([left_pad, width]);
1133

1134
  const mini_points = [
1135
    [max_at_time.length, 0],
1136
    [0, 0],
1137
  ];
1138

1139
  for (const [i, m] of max_at_time.entries()) {
1140
    const [_lastx, lasty] = mini_points[mini_points.length - 1];
1141
    if (m !== lasty) {
1142
      mini_points.push([i, lasty]);
1143
      mini_points.push([i, m]);
1144
    } else if (i === max_at_time.length - 1) {
1145
      mini_points.push([i, m]);
1146
    }
1147
  }
1148

1149
  let points = mini_points.map(([t, o]) => `${minixscale(t)}, ${yscale(o)}`);
1150
  points = points.join(' ');
1151
  mini_svg
1152
    .append('polygon')
1153
    .attr('points', points)
1154
    .attr('fill', schemeTableau10[0]);
1155

1156
  const xscale = scaleLinear()
1157
    .domain([0, max_at_time.length])
1158
    .range([0, plot_width]);
1159

1160
  const brush = brushX();
1161
  brush.extent([
1162
    [left_pad, 0],
1163
    [width, height],
1164
  ]);
1165
  brush.on('brush', function () {
1166
    const [begin, end] = d3.event.selection.map(x => x - left_pad);
1167

1168
    const stepbegin = Math.floor(xscale.invert(begin));
1169
    const stepend = Math.floor(xscale.invert(end));
1170
    let max = 0;
1171
    for (let i = stepbegin; i < stepend; i++) {
1172
      max = Math.max(max, max_at_time[i]);
1173
    }
1174
    plot.select_window(stepbegin, stepend, max);
1175
  });
1176
  mini_svg.call(brush);
1177
  return {};
1178
}
1179

1180
function Legend(plot_svg, categories) {
1181
  const xstart = 100;
1182
  const ystart = 5;
1183
  plot_svg
1184
    .append('g')
1185
    .selectAll('rect')
1186
    .data(categories)
1187
    .enter()
1188
    .append('rect')
1189
    .attr('x', (c, i) => xstart)
1190
    .attr('y', (c, i) => ystart + i * 15)
1191
    .attr('width', 10)
1192
    .attr('height', 10)
1193
    .attr('fill', (c, i) => schemeTableau10[i % schemeTableau10.length]);
1194
  plot_svg
1195
    .append('g')
1196
    .selectAll('text')
1197
    .data(categories)
1198
    .enter()
1199
    .append('text')
1200
    .attr('x', (c, i) => xstart + 20)
1201
    .attr('y', (c, i) => ystart + i * 15 + 8)
1202
    .attr('font-family', 'helvetica')
1203
    .attr('font-size', 10)
1204
    .text(c => c);
1205
  return {};
1206
}
1207

1208
function create_trace_view(
1209
  dst,
1210
  snapshot,
1211
  device,
1212
  plot_segments = false,
1213
  max_entries = 15000,
1214
) {
1215
  const left_pad = 70;
1216
  const data = process_alloc_data(snapshot, device, plot_segments, max_entries);
1217
  dst.selectAll('svg').remove();
1218
  dst.selectAll('div').remove();
1219

1220
  const d = dst.append('div');
1221
  d.append('input')
1222
    .attr('type', 'range')
1223
    .attr('min', 0)
1224
    .attr('max', data.elements_length)
1225
    .attr('value', max_entries)
1226
    .on('change', function () {
1227
      create_trace_view(dst, snapshot, device, plot_segments, this.value);
1228
    });
1229
  d.append('label').text('Detail');
1230

1231
  const grid_container = dst
1232
    .append('div')
1233
    .attr(
1234
      'style',
1235
      'display: grid; grid-template-columns: 1fr; grid-template-rows: 10fr 1fr 8fr; height: 100%; gap: 10px',
1236
    );
1237

1238
  const plot_svg = grid_container
1239
    .append('svg')
1240
    .attr('display', 'block')
1241
    .attr('viewBox', '0 0 1024 576')
1242
    .attr('preserveAspectRatio', 'none')
1243
    .attr('style', 'grid-column: 1; grid-row: 1; width: 100%; height: 100%;');
1244

1245
  const plot = MemoryPlot(plot_svg, data, left_pad, 1024, 576);
1246

1247
  if (snapshot.categories.length !== 0) {
1248
    Legend(plot_svg.append('g'), snapshot.categories);
1249
  }
1250

1251
  const mini_svg = grid_container
1252
    .append('svg')
1253
    .attr('display', 'block')
1254
    .attr('viewBox', '0 0 1024 60')
1255
    .attr('preserveAspectRatio', 'none')
1256
    .attr('style', 'grid-column: 1; grid-row: 2; width: 100%; height: 100%;');
1257

1258
  MiniMap(mini_svg, plot, data, left_pad, 1024);
1259
  const context_div = grid_container
1260
    .append('div')
1261
    .attr(
1262
      'style',
1263
      'grid-column: 1; grid-row: 3; width: 100%; height: 100%; overflow: auto;',
1264
    );
1265
  const delegate = ContextViewer(context_div.append('pre').text('none'), data);
1266
  plot.set_delegate(delegate);
1267
}
1268

1269
function unpickle(buffer) {
1270
  const bytebuffer = new Uint8Array(buffer);
1271
  const decoder = new TextDecoder();
1272

1273
  const stack = [];
1274
  const marks = [];
1275
  const memo = [];
1276
  let offset = 0;
1277
  let memo_id = 0;
1278

1279
  const APPENDS = 'e'.charCodeAt(0);
1280
  const BINGET = 'h'.charCodeAt(0);
1281
  const BININT = 'J'.charCodeAt(0);
1282
  const BININT1 = 'K'.charCodeAt(0);
1283
  const BININT2 = 'M'.charCodeAt(0);
1284
  const EMPTY_DICT = '}'.charCodeAt(0);
1285
  const EMPTY_LIST = ']'.charCodeAt(0);
1286
  const FRAME = 0x95;
1287
  const LONG1 = 0x8a;
1288
  const LONG_BINGET = 'j'.charCodeAt(0);
1289
  const MARK = '('.charCodeAt(0);
1290
  const MEMOIZE = 0x94;
1291
  const PROTO = 0x80;
1292
  const SETITEMS = 'u'.charCodeAt(0);
1293
  const SHORT_BINUNICODE = 0x8c;
1294
  const STOP = '.'.charCodeAt(0);
1295
  const TUPLE2 = 0x86;
1296
  const APPEND = 'a'.charCodeAt(0);
1297
  const NEWFALSE = 0x89;
1298
  const BINPUT = 'q'.charCodeAt(0);
1299
  const BINUNICODE = 'X'.charCodeAt(0);
1300
  const EMPTY_TUPLE = ')'.charCodeAt(0);
1301
  const NEWTRUE = 0x88;
1302
  const NONE = 'N'.charCodeAt(0);
1303
  const BINFLOAT = 'G'.charCodeAt(0);
1304
  const TUPLE = 't'.charCodeAt(0);
1305
  const TUPLE1 = 0x85;
1306
  const TUPLE3 = 0x87;
1307
  // untested
1308
  const LONG_BINPUT = 'r'.charCodeAt(0);
1309
  const LIST = 'l'.charCodeAt(0);
1310
  const DICT = 'd'.charCodeAt(0);
1311
  const SETITEM = 's'.charCodeAt(0);
1312

1313
  const scratch_buffer = new ArrayBuffer(8);
1314
  const scratch_bytes = new Uint8Array(scratch_buffer);
1315
  const big = new BigInt64Array(scratch_buffer);
1316
  const float64 = new Float64Array(scratch_buffer);
1317

1318
  function read_uint4() {
1319
    const n =
1320
      bytebuffer[offset] +
1321
      bytebuffer[offset + 1] * 256 +
1322
      bytebuffer[offset + 2] * 65536 +
1323
      bytebuffer[offset + 3] * 16777216;
1324
    offset += 4;
1325
    return n;
1326
  }
1327
  function setitems(d, mark) {
1328
    for (let i = mark; i < stack.length; i += 2) {
1329
      d[stack[i]] = stack[i + 1];
1330
    }
1331
    stack.splice(mark, Infinity);
1332
  }
1333

1334
  while (true) {
1335
    const opcode = bytebuffer[offset++];
1336
    switch (opcode) {
1337
      case PROTO:
1338
        {
1339
          const version = bytebuffer[offset++];
1340
          if (version < 2 || version > 4) {
1341
            throw new Error(`Unhandled version ${version}`);
1342
          }
1343
        }
1344
        break;
1345
      case APPEND:
1346
        {
1347
          const v = stack.pop();
1348
          stack.at(-1).push(v);
1349
        }
1350
        break;
1351
      case APPENDS:
1352
        {
1353
          const mark = marks.pop();
1354
          const arr = stack[mark - 1];
1355
          arr.push(...stack.splice(mark, Infinity));
1356
        }
1357
        break;
1358
      case LIST:
1359
      case TUPLE:
1360
        {
1361
          const mark = marks.pop();
1362
          stack.push([...stack.splice(mark, Infinity)]);
1363
        }
1364
        break;
1365
      case NEWFALSE:
1366
        stack.push(false);
1367
        break;
1368
      case NEWTRUE:
1369
        stack.push(true);
1370
        break;
1371
      case NONE:
1372
        stack.push(null);
1373
        break;
1374
      case BINGET:
1375
        stack.push(memo[bytebuffer[offset++]]);
1376
        break;
1377
      case BININT:
1378
        {
1379
          let i32 = read_uint4();
1380
          if (i32 > 0x7fffffff) {
1381
            i32 -= 0x100000000;
1382
          }
1383
          stack.push(i32);
1384
        }
1385
        break;
1386
      case BININT1:
1387
        stack.push(bytebuffer[offset++]);
1388
        break;
1389
      case BININT2:
1390
        {
1391
          const v = bytebuffer[offset] + bytebuffer[offset + 1] * 256;
1392
          stack.push(v);
1393
          offset += 2;
1394
        }
1395
        break;
1396
      case EMPTY_DICT:
1397
        stack.push({});
1398
        break;
1399
      case EMPTY_LIST:
1400
        stack.push([]);
1401
        break;
1402
      case FRAME:
1403
        offset += 8;
1404
        break;
1405
      case LONG1:
1406
        {
1407
          const s = bytebuffer[offset++];
1408
          if (s <= 8) {
1409
            for (let i = 0; i < s; i++) {
1410
              scratch_bytes[i] = bytebuffer[offset++];
1411
            }
1412
            const fill = scratch_bytes[s - 1] >= 128 ? 0xff : 0x0;
1413
            for (let i = s; i < 8; i++) {
1414
              scratch_bytes[i] = fill;
1415
            }
1416
            stack.push(Number(big[0]));
1417
          } else { // BigInt
1418
            let scratch_bytes_unbounded = [];
1419
            for (let i = 0; i < s; i++) {
1420
              scratch_bytes_unbounded.push(bytebuffer[offset++]);
1421
            }
1422

1423
            // BigInt can only convert from unsigned hex, thus we need to
1424
            // convert from twos-complement if negative
1425
            const negative = scratch_bytes_unbounded[s - 1] >= 128;
1426
            if (negative) {
1427
              // implements scratch_bytes_unbounded = ~scratch_bytes_unbounded + 1
1428
              // byte-by-byte.
1429
              let carry = 1;
1430
              for (let i = 0; i < s; i++) {
1431
                const twos_complement = (0xff ^ scratch_bytes_unbounded[i]) + carry;
1432
                carry = twos_complement > 0xff ? 1 : 0;
1433
                scratch_bytes_unbounded[i] = 0xff & twos_complement;
1434
              }
1435
            }
1436

1437
            const hex_str = Array.from(scratch_bytes_unbounded.reverse(), byte => {
1438
              return byte.toString(16).padStart(2, '0');
1439
            }).join('');
1440

1441
            const big_int = negative ? -BigInt(`0x${hex_str}`) : BigInt(`0x${hex_str}`);
1442
            stack.push(big_int);
1443
          }
1444
        }
1445
        break;
1446
      case LONG_BINGET:
1447
        {
1448
          const idx = read_uint4();
1449
          stack.push(memo[idx]);
1450
        }
1451
        break;
1452
      case MARK:
1453
        marks.push(stack.length);
1454
        break;
1455
      case MEMOIZE:
1456
        memo[memo_id++] = stack.at(-1);
1457
        break;
1458
      case BINPUT:
1459
        memo[bytebuffer[offset++]] = stack.at(-1);
1460
        break;
1461
      case LONG_BINPUT:
1462
        memo[read_uint4()] = stack.at(-1);
1463
        break;
1464
      case SETITEMS:
1465
        {
1466
          const mark = marks.pop();
1467
          const d = stack[mark - 1];
1468
          setitems(d, mark);
1469
        }
1470
        break;
1471
      case SETITEM: {
1472
        const v = stack.pop();
1473
        const k = stack.pop();
1474
        stack.at(-1)[k] = v;
1475
        break;
1476
      }
1477
      case DICT:
1478
        {
1479
          const mark = marks.pop();
1480
          const d = {};
1481
          setitems(d, mark);
1482
          stack.push(d);
1483
        }
1484
        break;
1485
      case SHORT_BINUNICODE:
1486
        {
1487
          const n = bytebuffer[offset++];
1488
          stack.push(decoder.decode(new Uint8Array(buffer, offset, n)));
1489
          offset += n;
1490
        }
1491
        break;
1492
      case BINUNICODE:
1493
        {
1494
          const n = read_uint4();
1495
          stack.push(decoder.decode(new Uint8Array(buffer, offset, n)));
1496
          offset += n;
1497
        }
1498
        break;
1499
      case STOP:
1500
        return stack.pop();
1501
      case EMPTY_TUPLE:
1502
        stack.push([]);
1503
        break;
1504
      case TUPLE1:
1505
        stack.push([stack.pop()]);
1506
        break;
1507
      case TUPLE2:
1508
        stack.push(stack.splice(-2, Infinity));
1509
        break;
1510
      case TUPLE3:
1511
        stack.push(stack.splice(-3, Infinity));
1512
        break;
1513
      case BINFLOAT:
1514
        for (let i = 7; i >= 0; i--) {
1515
          // stored in big-endian order
1516
          scratch_bytes[i] = bytebuffer[offset++];
1517
        }
1518
        stack.push(float64[0]);
1519
        break;
1520
      default:
1521
        throw new Error(`UNKNOWN OPCODE: ${opcode}`);
1522
    }
1523
  }
1524
}
1525

1526
function decode_base64(input) {
1527
  function decode_char(i, shift) {
1528
    const nChr = input.charCodeAt(i);
1529
    const r =
1530
      nChr > 64 && nChr < 91
1531
        ? nChr - 65
1532
        : nChr > 96 && nChr < 123
1533
        ? nChr - 71
1534
        : nChr > 47 && nChr < 58
1535
        ? nChr + 4
1536
        : nChr === 43
1537
        ? 62
1538
        : nChr === 47
1539
        ? 63
1540
        : 0;
1541
    return r << shift;
1542
  }
1543
  const output = new Uint8Array((input.length / 4) * 3);
1544
  for (let i = 0, j = 0; i < input.length; i += 4, j += 3) {
1545
    const u24 =
1546
      decode_char(i, 18) +
1547
      decode_char(i + 1, 12) +
1548
      decode_char(i + 2, 6) +
1549
      decode_char(i + 3);
1550
    output[j] = u24 >> 16;
1551
    output[j + 1] = (u24 >> 8) & 0xff;
1552
    output[j + 2] = u24 & 0xff;
1553
  }
1554
  return output.buffer;
1555
}
1556

1557
const kinds = {
1558
  'Active Memory Timeline': create_trace_view,
1559
  'Allocator State History': create_segment_view,
1560
  'Active Cached Segment Timeline': (dst, snapshot, device) =>
1561
    create_trace_view(dst, snapshot, device, true),
1562
};
1563

1564
const snapshot_cache = {};
1565
const snapshot_to_loader = {};
1566
const snapshot_to_url = {};
1567
const selection_to_div = {};
1568

1569
const style = `
1570
pre {
1571
  margin: 0px;
1572
}
1573
html, body {
1574
  height: 100%;
1575
  overflow: clip;
1576
}`;
1577

1578
const head = d3.select('head');
1579
head.append('style').text(style);
1580
const body = d3.select('body');
1581
const snapshot_select = body.append('select');
1582
const view = body.append('select');
1583
for (const x in kinds) {
1584
  view.append('option').text(x);
1585
}
1586
const gpu = body.append('select');
1587

1588
function unpickle_and_annotate(data) {
1589
  data = unpickle(data);
1590
  console.log(data);
1591
  annotate_snapshot(data);
1592
  return data;
1593
}
1594

1595
function snapshot_change(f) {
1596
  const view_value = view.node().value;
1597
  let device = Number(gpu.node().value);
1598
  const snapshot = snapshot_cache[f];
1599
  gpu.selectAll('option').remove();
1600
  const has_segments = {};
1601
  for (const s of snapshot.segments) {
1602
    has_segments[s.device] = true;
1603
  }
1604
  let device_valid = false;
1605
  for (const [i, trace] of snapshot.device_traces.entries()) {
1606
    if (trace.length > 0 || i in has_segments) {
1607
      gpu.append('option').text(i);
1608
      if (i === device) {
1609
        device_valid = true;
1610
        gpu.node().selectedIndex = gpu.node().children.length - 1;
1611
      }
1612
    }
1613
  }
1614
  if (!device_valid) {
1615
    device = Number(gpu.node().value);
1616
  }
1617
  const key = [f, view_value, device];
1618
  if (!(key in selection_to_div)) {
1619
    selection_to_div[key] = d3.select('body').append('div');
1620
    kinds[view_value](selection_to_div[key], snapshot, device);
1621
  }
1622
  const selected_div = selection_to_div[key];
1623

1624
  selected_div.attr('style', 'display: float; height: 100%');
1625
}
1626

1627
function selected_change() {
1628
  for (const d of Object.values(selection_to_div)) {
1629
    d.attr('style', 'display: none; height: 100%');
1630
  }
1631
  const f = snapshot_select.node().value;
1632
  if (f === '') {
1633
    return;
1634
  }
1635
  if (!(f in snapshot_cache)) {
1636
    snapshot_to_loader[f](f);
1637
  } else {
1638
    snapshot_change(f);
1639
  }
1640
}
1641

1642
snapshot_select.on('change', selected_change);
1643
view.on('change', selected_change);
1644
gpu.on('change', selected_change);
1645

1646
body.on('dragover', e => {
1647
  event.preventDefault();
1648
});
1649

1650
body.on('drop', () => {
1651
  console.log(event.dataTransfer.files);
1652
  Array.from(event.dataTransfer.files).forEach(file => {
1653
    add_snapshot(file.name, unique_name => {
1654
      const reader = new FileReader();
1655
      reader.onload = e => {
1656
        finished_loading(unique_name, e.target.result);
1657
      };
1658
      reader.readAsArrayBuffer(file);
1659
    });
1660
  });
1661
  event.preventDefault();
1662
  snapshot_select.node().selectedIndex =
1663
    snapshot_select.node().options.length - 1;
1664
  selected_change();
1665
});
1666

1667
selection_to_div[''] = body
1668
  .append('div')
1669
  .text(
1670
    'Drag and drop a file to load a local snapshot. No data from the snapshot is uploaded.',
1671
  );
1672

1673
let next_unique_n = 1;
1674
function add_snapshot(name, loader) {
1675
  if (name in snapshot_to_loader) {
1676
    name = `${name} (${next_unique_n++})`;
1677
  }
1678
  snapshot_select.append('option').text(name);
1679
  snapshot_to_loader[name] = loader;
1680
}
1681

1682
function finished_loading(name, data) {
1683
  snapshot_cache[name] = unpickle_and_annotate(data);
1684
  snapshot_change(name);
1685
}
1686

1687
export function add_remote_files(files) {
1688
  files.forEach(f =>
1689
    add_snapshot(f.name, unique_name => {
1690
      console.log('fetching', f.url);
1691
      fetch(f.url)
1692
        .then(x => x.arrayBuffer())
1693
        .then(data => finished_loading(unique_name, data));
1694
    }),
1695
  );
1696
  if (files.length > 0) {
1697
    selected_change();
1698
  }
1699
}
1700

1701
export function add_local_files(files, view_value) {
1702
  view.node().value = view_value;
1703
  files.forEach(f =>
1704
    add_snapshot(f.name, unique_name => {
1705
      finished_loading(unique_name, decode_base64(f.base64));
1706
    }),
1707
  );
1708
  if (files.length > 0) {
1709
    selected_change();
1710
  }
1711
}
1712