private findPeakMemoryUsage()

in frontend/app/components/memory_viewer/memory_usage/memory_usage.ts [243:344]
99 lines of code
19 McCabe index (conditional complexity)

  private findPeakMemoryUsage(
      trace: proto.HeapSimulatorTrace|null, color: number) {
    const heapSizes: number[] = [];
    const unpaddedHeapSizes: number[] = [];
    let logicalBuffers: number[] = [];
    let peakLogicalBuffers: number[] = [];
    let heapSizeBytes = 0;
    let unpaddedHeapSizeBytes = 0;
    let peakHeapSizeBytes = 0;
    let unpaddedPeakHeapSizeBytes = 0;
    let peakHeapSizePosition = 0;

    if (trace) {
      for (const event of trace.events || []) {
        heapSizes.push(utils.bytesToMiB(heapSizeBytes));
        unpaddedHeapSizes.push(utils.bytesToMiB(unpaddedHeapSizeBytes));
        const eventId = utils.toNumber(event.bufferId);
        const buffer = this.idToBuffer[eventId];
        this.unSeenLogicalBuffers.delete(eventId);
        const alloc = this.idToBufferAllocation[eventId];
        if (alloc) {
          this.seenBufferAllocations.add(alloc.index);
        }
        let shape: Shape|null = null;
        if (buffer.instructionName && buffer.instructionName !== '') {
          const hlo = this.nameToHlo[buffer.instructionName];
          if (hlo && hlo.shape) {
            shape = hlo.shape.resolveShapeIndex(buffer.shapeIndex);
          }
        }
        switch (event.kind) {
          // Default to 'ALLOC' when event.kind is undefined. This is because
          // by default proto3 primitive fields with default values will be
          // omitted in JSON output.
          case undefined:
          case 'ALLOC':
          case 'SHARE_WITH':
            logicalBuffers.push(eventId);
            heapSizeBytes += buffer.size;
            if (shape) {
              unpaddedHeapSizeBytes += shape.unpaddedHeapSizeBytes();
            }
            this.logicalBufferSpans[eventId] = [heapSizes.length - 1, -1];
            if (heapSizeBytes > peakHeapSizeBytes) {
              peakHeapSizeBytes = heapSizeBytes;
              unpaddedPeakHeapSizeBytes = unpaddedHeapSizeBytes;
              peakHeapSizePosition = heapSizes.length - 1;
              peakLogicalBuffers = logicalBuffers.slice();
            }
            break;
          case 'FREE':
            logicalBuffers = logicalBuffers.filter(item => {
              return item !== eventId;
            });
            heapSizeBytes -= buffer.size;
            if (shape) {
              unpaddedHeapSizeBytes -= shape.unpaddedHeapSizeBytes();
            }
            if (!this.logicalBufferSpans[eventId]) {
              // The logical buffer is not allocated in this module.
              this.logicalBufferSpans[eventId] = [0, heapSizes.length - 1];
              console.warn(
                  event, ' is freed but has seen no allocation event.');
            } else {
              this.logicalBufferSpans[eventId][1] = heapSizes.length - 1;
            }
            if (heapSizeBytes < 0) {
              console.error('heap_size_bytes < 0');
            }
            break;
          default:
            console.log('ERROR: unknown heap event kind: ', event);
            break;
        }
      }
      heapSizes.push(utils.bytesToMiB(heapSizeBytes));
      unpaddedHeapSizes.push(utils.bytesToMiB(unpaddedHeapSizeBytes));
    }

    this.peakLogicalBuffers = peakLogicalBuffers;
    this.peakHeapSizePosition = peakHeapSizePosition;

    if (this.includeNotSimulated) {
      const indefiniteMemoryUsageBytes =
          this.findIndefiniteMemoryUsage(this.unSeenLogicalBuffers, color);
      this.peakHeapSizeBytes =
          peakHeapSizeBytes + indefiniteMemoryUsageBytes.padded;
      this.unpaddedPeakHeapSizeBytes =
          unpaddedPeakHeapSizeBytes + indefiniteMemoryUsageBytes.unpadded;
      const addendPadded = utils.bytesToMiB(indefiniteMemoryUsageBytes.padded);
      this.heapSizes = heapSizes.map(item => item + addendPadded);
      const addendUnpadded =
          utils.bytesToMiB(indefiniteMemoryUsageBytes.unpadded);
      this.unpaddedHeapSizes =
          unpaddedHeapSizes.map(item => item + addendUnpadded);
    } else {
      this.peakHeapSizeBytes = peakHeapSizeBytes;
      this.unpaddedPeakHeapSizeBytes = unpaddedPeakHeapSizeBytes;
      this.heapSizes = heapSizes;
      this.unpaddedHeapSizes = unpaddedHeapSizes;
    }
  }