Calculate dynamic and static buffer sizes for non tiling platforms

Author: Xeratec

Deeploy/CommonExtensions/CodeTransformationPasses/MemoryAllocation.py

@@ -142,13 +142,31 @@ def apply(self,
         for buffer in reversed(self.topologicallySortBuffers(outputs + transients)):
             assert buffer._live == False, f"Tried to allocate already live buffer {buffer.name}"
             buffer._live = True
+
+            memoryLevel = "None" if not hasattr(buffer, "_memoryLevel") else buffer._memoryLevel
+            if memoryLevel not in ctxt._dynamicSize:
+                ctxt._dynamicSize[memoryLevel] = int(buffer.sizeInBytes())
+            else:
+                ctxt._dynamicSize[memoryLevel] += int(buffer.sizeInBytes())
+
             executionBlock.addLeft(buffer.allocTemplate, buffer._bufferRepresentation())
 
+        for levels in ctxt._dynamicSize.keys():
+            if levels not in ctxt._maxDynamicSize:
+                ctxt._maxDynamicSize[levels] = max(0, ctxt._dynamicSize[levels])
+            else:
+                ctxt._maxDynamicSize[levels] = max(ctxt._maxDynamicSize.get(levels, 0), ctxt._dynamicSize[levels])
+
         for buffer in inputs + transients:
             assert buffer._live == True, f"Tried to deallocate already dead buffer {buffer.name}"
             buffer._live = False
             # Don't deallocate if it's an alias of a live buffer
             if not buffer.has_live_ancestors(ctxt = ctxt):
+                memoryLevel = "None" if not hasattr(buffer, "_memoryLevel") else buffer._memoryLevel
+                if memoryLevel not in ctxt._dynamicSize:
+                    ctxt._dynamicSize[memoryLevel] = 0
+                else:
+                    ctxt._dynamicSize[memoryLevel] -= int(buffer.sizeInBytes())
                 executionBlock.addRight(buffer.deallocTemplate, buffer._bufferRepresentation())
 
         return ctxt, executionBlock
@@ -177,10 +195,30 @@ def apply(self,
 
         for buffer in outputs + transients:
             assert buffer._live == False, f"Tried to allocate already live buffer {buffer.name}"
+
+            memoryLevel = "None" if not hasattr(buffer, "_memoryLevel") else buffer._memoryLevel
+            if memoryLevel not in ctxt._dynamicSize:
+                ctxt._dynamicSize[memoryLevel] = int(buffer.sizeInBytes())
+            else:
+                ctxt._dynamicSize[memoryLevel] += int(buffer.sizeInBytes())
+
             buffer._live = True
 
+        for levels in ctxt._dynamicSize.keys():
+            if levels not in ctxt._maxDynamicSize:
+                ctxt._maxDynamicSize[levels] = max(0, ctxt._dynamicSize[levels])
+            else:
+                ctxt._maxDynamicSize[levels] = max(ctxt._maxDynamicSize.get(levels, 0), ctxt._dynamicSize[levels])
+
         for buffer in inputs + transients:
             assert buffer._live == True, f"Tried to deallocate already dead buffer {buffer.name}"
+
+            memoryLevel = "None" if not hasattr(buffer, "_memoryLevel") else buffer._memoryLevel
+            if memoryLevel not in ctxt._dynamicSize:
+                ctxt._dynamicSize[memoryLevel] = 0
+            else:
+                ctxt._dynamicSize[memoryLevel] -= int(buffer.sizeInBytes())
+
             buffer._live = False
 
         return ctxt, executionBlock

Deeploy/CommonExtensions/NetworkDeployers/NetworkDeployerWrapper.py

@@ -96,3 +96,9 @@ def generateBufferAllocationCode(self) -> str:
     # MultiEngineDeployer augment
     def _mapNode(self, node: gs.Node) -> Union[ONNXLayer, Any]:
         return self._innerObject._mapNode(node)
+
+    def _printMemorySummary(self):
+        return self._innerObject._printMemorySummary()
+
+    def _printInputOutputSummary(self):
+        return self._innerObject._printInputOutputSummary()

Deeploy/CommonExtensions/NetworkDeployers/SignPropDeployer.py

@@ -28,7 +28,7 @@
 import onnx_graphsurgeon as gs
 
 from Deeploy.AbstractDataTypes import Pointer
-from Deeploy.DeeployTypes import CodeGenVerbosity, DeploymentPlatform, NetworkDeployer, TopologyOptimizer, _NoVerbosity
+from Deeploy.DeeployTypes import DeploymentPlatform, NetworkDeployer, TopologyOptimizer
 from Deeploy.Logging import DEFAULT_LOGGER as log
 
 
@@ -64,35 +64,15 @@ def _createIOBindings(self, ctxt, graph):
 
         return ctxt
 
-    def generateFunction(self, verbose: CodeGenVerbosity = _NoVerbosity) -> str:
-        """Helper function to prepare deployment and return generated function code
-
-        """
-
-        if not self.prepared:
-            self.prepare(verbose = verbose)
-
-        log.info("=" * 80)
-        log.info("Deeploy Code Generation")
-        log.info("=" * 80)
-
+    def _printInputOutputSummary(self):
         log.info('Input:')
-        for name in self.inputTypes.keys():
-            buf = self.ctxt.lookup(name)
+        for buf in self.inputs():
             log.info(
-                f" - '{name}': Type: {buf._type.referencedType.typeName}, nLevels: {buf.nLevels}, Signed: {buf._signed}, Offset: {self.inputOffsets[name]}"
+                f" - '{buf.name}': Type: {buf._type.referencedType.typeName}, nLevels: {buf.nLevels}, Signed: {buf._signed}, Offset: {self.inputOffsets[buf.name]}"
             )
 
         log.info('Output:')
         for buf in self.outputs():
             log.info(
                 f" - '{buf.name}': Type: {buf._type.referencedType.typeName}, nLevels: {buf.nLevels}, Signed: {buf._signed}"
             )
-
-        log.info("-" * 80)
-        num_ops = self.numberOfOps(verbose = True)
-        log.info("-" * 80)
-        log.info(f"Number of Ops.                : {num_ops}")
-        log.info(f"Model Parameters              : {self.getParameterSize()}")
-
-        return self.generateInferenceCode()

Deeploy/DeeployTypes.py

@@ -399,6 +399,17 @@ def has_live_ancestors(self, ctxt: NetworkContext) -> bool:
 
         return False
 
+    def sizeInBytes(self) -> int:
+        """Returns the size of this VariableBuffer in bytes
+
+        Returns
+        -------
+        int
+            Size of this VariableBuffer in bytes
+
+        """
+        return int(np.prod(self.shape) * (self._type.referencedType.typeWidth // 8))
+
 
 class TransientBuffer(VariableBuffer):
     """Class to represent memory space required by kernels that is not covered by input and output tensors, e.g. im2col buffers in convolutions
@@ -445,6 +456,9 @@ def __repr__(self) -> str:
     def fromVariableBuffer(cls, buffer: VariableBuffer):
         ret = cls(name = buffer.name, size = np.prod(buffer.shape) * buffer._type.typeWidth // 8)
 
+    def sizeInBytes(self) -> int:
+        return int(self.size)
+
 
 class ConstantBuffer(VariableBuffer):
     """Class to represent compile-time constant tensors (weights, biases, other parameters) within Deeploy.
@@ -587,6 +601,9 @@ def __init__(self,
         self.TransientBuffer = transientBuffer
         self.name = name
 
+        self._maxDynamicSize = {}  #: int: Maximum dynamic memory size occupied by live buffers at any point in time
+        self._dynamicSize = {}  #: int: Current dynamic memory size occupied by live buffers
+
     def dealiasBuffer(self, name: str) -> str:
         """Function to find the underlying aliased VariableBuffer
 
@@ -2965,8 +2982,7 @@ def generateIOBufferInitializationCode(self) -> str:
     def worstCaseBufferSize(self):
         """Return the worst-case buffer size occupied by the network implementaiton
         """
-        # WIESEP: There is no reasonable value for a worst case buffer size without tiling
-        raise NotImplementedError("Worst case buffer size is not known or not implemented!")
+        return self.ctxt._maxDynamicSize
 
     # Don't override this
     def generateBufferInitializationCode(self) -> str:
@@ -3116,54 +3132,6 @@ def generateEngineInitializationCode(self) -> str:
         """
         return ("\n").join([engine.initCode for engine in self.Platform.engines])
 
-    # Don't override this - Returns parameter size in bytes
-    def getParameterSize(self) -> int:
-        """Return the BYTE size of all static network parameters (weights, biases, parameters,...)
-
-        Returns
-        -------
-        int
-            Size of all network parameters
-
-        Raises
-        ------
-        RuntimeError
-            Raises a RuntimeError if network is not parsed and bound
-
-
-        """
-        if not self.parsed or not self.bound:
-            raise RuntimeError('You need to parse and bind the network before getting RAM Size!')
-
-        size = 0
-        for _buffer in self.ctxt.globalObjects.values():
-            # We do not count structs for now, since they are not properly modeled
-            if isinstance(_buffer, ConstantBuffer) and _buffer._deploy:
-                size += int((np.prod(_buffer.shape) * _buffer._type.typeWidth // 8))
-
-        return size
-
-    # Don't override this - Returns worst case layer and buffering size in bytes
-    def getTotalSize(self) -> int:
-        """Returns total size of the network, consisting of all parameters and intermediate buffer size
-
-        Returns
-        -------
-        int
-            Total network size
-
-        Raises
-        ------
-        RuntimeError
-            Raises a RuntimeError if network is not parsed and bound
-
-
-        """
-        if not self.parsed or not self.bound:
-            raise RuntimeError('You need to parse and bind the network before getting RAM Size!')
-
-        return self.getParameterSize() + self.worstCaseBufferSize
-
     def numberOfOps(self, verbose: bool) -> int:
         """Returns the total number of operations per network inference
 
@@ -3584,6 +3552,41 @@ def prepare(self, verbose: CodeGenVerbosity = _NoVerbosity):
         self.backEnd(verbose = verbose)
         self.prepared = True
 
+    def _printInputOutputSummary(self):
+        log.info("Input:")
+        for buf in self.inputs():
+            log.info(f" - '{buf.name}': Type: {buf._type.referencedType.typeName}")
+
+        log.info('Output:')
+        for buf in self.outputs():
+            log.info(f" - '{buf.name}': Type: {buf._type.referencedType.typeName}")
+
+    def _printMemorySummary(self):
+        log.info("")
+        log.info("Memory Usage Report:")
+        log.info(f"Level                 Total (bytes)   (Static + Dynamic)    ")
+        log.info("-" * 80)
+
+        _worstCaseBufferSize = self.worstCaseBufferSize
+        if len(_worstCaseBufferSize) == 0:
+            _worstCaseBufferSize = {"None": 0}
+
+        for level, dynamicSize in _worstCaseBufferSize.items():
+            staticSize = 0
+            for _buffer in self.ctxt.globalObjects.values():
+                # We do not count structs for now, since they are not properly modeled
+                if isinstance(_buffer, ConstantBuffer) or (isinstance(_buffer, VariableBuffer) and _buffer._deploy):
+                    # SCHEREMO: We only
+                    if (hasattr(_buffer, "_memoryLevel") and _buffer._memoryLevel == level) or level == "None":
+                        staticSize += int((np.prod(_buffer.shape) * _buffer._type.referencedType.typeWidth // 8))
+                    else:
+                        log.warning(f"Buffer {_buffer.name} does not have a valid memory level")
+
+            total = staticSize + dynamicSize
+
+            log.info(f"{level:<22}     {total:8,d}   "
+                     f"({staticSize:6,d} + {dynamicSize:7,d})  ")
+
     def generateFunction(self, verbose: CodeGenVerbosity = _NoVerbosity) -> str:
         """Helper function to prepare deployment and return generated function code
 
@@ -3596,19 +3599,13 @@ def generateFunction(self, verbose: CodeGenVerbosity = _NoVerbosity) -> str:
         log.info("Deeploy Code Generation")
         log.info("=" * 80)
 
-        log.info('Input:')
-        for name in self.inputTypes.keys():
-            buf = self.ctxt.lookup(name)
-            log.info(f" - '{name}': Type: {buf._type.referencedType.typeName}")
-
-        log.info('Output:')
-        for buf in self.outputs():
-            log.info(f" - '{buf.name}': Type: {buf._type.referencedType.typeName}")
+        self._printInputOutputSummary()
 
         num_ops = self.numberOfOps(verbose = True)
         log.info("-" * 80)
 
         log.info(f"Number of Ops.                : {num_ops}")
-        log.info(f"Model Parameters              : {self.getParameterSize()}")
+
+        self._printMemorySummary()
 
         return self.generateInferenceCode()

Deeploy/MemoryLevelExtension/NetworkDeployers/MemoryLevelDeployer.py

@@ -28,6 +28,7 @@
 from types import MappingProxyType
 from typing import Any, Callable, Dict, List, Tuple, Type, Union
 
+import numpy as np
 import onnx_graphsurgeon as gs
 
 from Deeploy.AbstractDataTypes import Pointer
@@ -149,6 +150,25 @@ def codeTransform(self, verbose: CodeGenVerbosity = _NoVerbosity):
         self.ctxt, self.graph = self.memoryLevelAnnotationOptimizer.optimize(self.ctxt, self.graph)
         super().codeTransform(verbose)
 
+    def _printMemorySummary(self):
+        log.info("Memory Usage Report:")
+        log.info(f"{'Level':<22} {'Capacity (bytes)':>16}   {'Total':>8}   {'(Static + Dynamic)':<21} {'Usage':<6}")
+        log.info("-" * 80)
+
+        for level, dynamicSize in self.worstCaseBufferSize.items():
+            staticSize = 0
+            for _buffer in self.ctxt.globalObjects.values():
+                # We do not count structs for now, since they are not properly modeled
+                if isinstance(_buffer, ConstantBuffer) and _buffer._deploy and _buffer._memoryLevel == level:
+                    staticSize += int((np.prod(_buffer.shape) * _buffer._type.referencedType.typeWidth // 8))
+
+            capacity = self.Platform.memoryHierarchy.memoryLevels[level].size
+            total = staticSize + dynamicSize
+
+            log.info(f"{level:<22} {capacity:16,}   {total:8,d}   "
+                     f"({staticSize:6,d} + {dynamicSize:7,d})  "
+                     f"({total / capacity * 100:5.1f}%)")
+
 
 class MemoryLevelAwareSignPropDeployer(SignPropDeployer):
 
@@ -206,6 +226,25 @@ def codeTransform(self, verbose: CodeGenVerbosity = _NoVerbosity):
         self.ctxt, self.graph = self.memoryLevelAnnotationOptimizer.optimize(self.ctxt, self.graph)
         super().codeTransform(verbose)
 
+    def _printMemorySummary(self):
+        log.info("Memory Usage Report:")
+        log.info(f"{'Level':<22} {'Capacity (bytes)':>16}   {'Total':>8}   {'(Static + Dynamic)':<21} {'Usage':<6}")
+        log.info("-" * 80)
+
+        for level, dynamicSize in self.worstCaseBufferSize.items():
+            staticSize = 0
+            for _buffer in self.ctxt.globalObjects.values():
+                # We do not count structs for now, since they are not properly modeled
+                if isinstance(_buffer, ConstantBuffer) and _buffer._deploy and _buffer._memoryLevel == level:
+                    staticSize += int((np.prod(_buffer.shape) * _buffer._type.referencedType.typeWidth // 8))
+
+            capacity = self.Platform.memoryHierarchy.memoryLevels[level].size
+            total = staticSize + dynamicSize
+
+            log.info(f"{level:<22} {capacity:16,}   {total:8,d}   "
+                     f"({staticSize:6,d} + {dynamicSize:7,d})  "
+                     f"({total / capacity * 100:5.1f}%)")
+
 
 class MemoryDeployerWrapper(NetworkDeployerWrapper):
 
@@ -253,3 +292,22 @@ def bind(self):
     def codeTransform(self, verbose: CodeGenVerbosity = _NoVerbosity):
         self.ctxt, self.graph = self.memoryLevelAnnotationOptimizer.optimize(self.ctxt, self.graph)
         super().codeTransform(verbose)
+
+    def _printMemorySummary(self):
+        log.info("Memory Usage Report:")
+        log.info(f"{'Level':<22} {'Capacity (bytes)':>16}   {'Total':>8}   {'(Static + Dynamic)':<21} {'Usage':<6}")
+        log.info("-" * 80)
+
+        for level, dynamicSize in self.worstCaseBufferSize.items():
+            staticSize = 0
+            for _buffer in self.ctxt.globalObjects.values():
+                # We do not count structs for now, since they are not properly modeled
+                if isinstance(_buffer, ConstantBuffer) and _buffer._deploy and _buffer._memoryLevel == level:
+                    staticSize += int((np.prod(_buffer.shape) * _buffer._type.referencedType.typeWidth // 8))
+
+            capacity = self.Platform.memoryHierarchy.memoryLevels[level].size
+            total = staticSize + dynamicSize
+
+            log.info(f"{level:<22} {capacity:16,}   {total:8,d}   "
+                     f"({staticSize:6,d} + {dynamicSize:7,d})  "
+                     f"({total / capacity * 100:5.1f}%)")