revert memorylevelannotationpass

Author: runwangdl

Deeploy/MemoryLevelExtension/OptimizationPasses/MemoryLevelAnnotationPasses.py

@@ -4,12 +4,11 @@
 
 from typing import List, Tuple
 
-import numpy as np
 import onnx_graphsurgeon as gs
 
 from Deeploy.CommonExtensions.OptimizationPasses.PassClasses import SequentialPass
-from Deeploy.DeeployTypes import ConstantBuffer, NetworkContext, VariableBuffer
-from Deeploy.MemoryLevelExtension.MemoryLevels import MemoryHierarchy, MemoryLevel
+from Deeploy.DeeployTypes import NetworkContext, VariableBuffer
+from Deeploy.MemoryLevelExtension.MemoryLevels import MemoryHierarchy
 
 
 class AnnotateDefaultMemoryLevel(SequentialPass):
@@ -19,15 +18,9 @@ def __init__(self, memoryHierarchy: MemoryHierarchy):
         self.memoryHierarchy = memoryHierarchy
 
     def apply(self, ctxt: NetworkContext, graph: gs.Graph) -> Tuple[NetworkContext, gs.Graph]:
-        defaultLevel = self.memoryHierarchy.getDefaultMemoryLevel().name
         for _buffer in {**ctxt.localObjects, **ctxt.globalObjects}.values():
             if not hasattr(_buffer, "_memoryLevel"):
-                if isinstance(_buffer, ConstantBuffer) and defaultLevel == "L3":
-                    # ConstantBuffer (weights, biases, BN params) are small and
-                    # excluded from the tiling DMA schedule.  Keep them in L2.
-                    _buffer._memoryLevel = "L2"
-                else:
-                    _buffer._memoryLevel = defaultLevel
+                _buffer._memoryLevel = self.memoryHierarchy.getDefaultMemoryLevel().name
         return ctxt, graph
 
 
@@ -53,37 +46,4 @@ def globalBuffers(tensors: List[gs.Tensor]) -> List[VariableBuffer]:
         for _buffer in buffers:
             _buffer._memoryLevel = self.ioLevel
 
-        return ctxt, graph
-
-
-class AnnotateNeurekaWeightMemoryLevel(SequentialPass):
-
-    def __init__(self, neurekaEngineName: str, weightMemoryLevel: MemoryLevel):
-        self._weightMemoryLevel = weightMemoryLevel
-        self.neurekaEngineName = neurekaEngineName
-        super().__init__()
-
-    def apply(self, ctxt: NetworkContext, graph: gs.Graph) -> Tuple[NetworkContext, gs.Graph]:
-
-        def _neurekaWeightBufferSize(buffer: ConstantBuffer) -> int:
-            return int(np.prod(buffer.shape))  # Weights are encoded as bytes so no need to check for typeWidth
-
-        weightMemoryOccupation = 0
-
-        # Current weight memory occupation
-        for buffer in {**ctxt.globalObjects, **ctxt.localObjects}.values():
-            if hasattr(buffer, "_memoryLevel") and buffer._memoryLevel == self._weightMemoryLevel.name:
-                weightMemoryOccupation += _neurekaWeightBufferSize(buffer)
-
-        neurekaNodes = [node for node in graph.nodes if node.attrs["engine"] == self.neurekaEngineName]
-        for node in neurekaNodes:
-            if node.op in ["Conv", "RequantizedConv"]:
-
-                if not (ctxt.is_local(node.inputs[1].name) or ctxt.is_global(node.inputs[1].name)):
-                    continue
-
-                buffer = ctxt.lookup(node.inputs[1].name)
-                if weightMemoryOccupation + _neurekaWeightBufferSize(buffer) < self._weightMemoryLevel.size:
-                    buffer._memoryLevel = self._weightMemoryLevel.name
-                    weightMemoryOccupation += _neurekaWeightBufferSize(buffer)
-        return ctxt, graph
+        return ctxt, graph