Fix memory analysis crash when there is multiple free with same addre…

…ss (#560)

tb_plugin/torch_tb_profiler/profiler/memory_parser.py

@@ -305,7 +305,6 @@ def get_preprocessed_records(self):
         memory_records = sorted(self.memory_records, key=lambda r: r.ts)
 
         alloc = {}  # allocation events may or may not have paired free event
-        free = {}  # free events that does not have paired alloc event
         prev_ts = float('-inf')  # ensure ordered memory records is ordered
         for i, r in enumerate(memory_records):
             if r.addr is None:
@@ -326,10 +325,4 @@ def get_preprocessed_records(self):
                     r.op_name = alloc_r.op_name
                     r.parent_op_name = alloc_r.parent_op_name
                     del alloc[addr]
-                else:
-                    assert addr not in free
-                    free[addr] = i
-
-        if free:
-            logger.debug(f'{len(free)} memory records do not have associated operator.')
         return memory_records

tb_plugin/torch_tb_profiler/run.py

@@ -4,13 +4,15 @@
 from collections import defaultdict
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
-from . import consts
+from . import consts, utils
 from .profiler.diffrun import compare_op_tree, diff_summary
 from .profiler.memory_parser import MemoryMetrics, MemoryRecord, MemorySnapshot
 from .profiler.module_op import Stats
 from .profiler.node import OperatorNode
 from .utils import Canonicalizer, DisplayRounder
 
+logger = utils.get_logger()
+
 
 class Run(object):
     """ A profiler run. For visualization purpose only.
@@ -341,7 +343,7 @@ def get_op_name_or_ctx(record: MemoryRecord):
                 # profile json data prior to pytorch 1.10 do not have addr
                 # we should ignore them
                 continue
-            assert prev_ts < r.ts
+            assert prev_ts <= r.ts
             prev_ts = r.ts
             addr = r.addr
             size = r.bytes
@@ -362,7 +364,8 @@ def get_op_name_or_ctx(record: MemoryRecord):
                     ])
                     del alloc[addr]
                 else:
-                    assert addr not in free
+                    if addr in free:
+                        logger.warning(f'Address {addr} is freed multiple times')
                     free[addr] = i
 
         for i in alloc.values():