Add trainingscheduler to keep input forever alive

Author: runwangdl

DeeployTest/testMVPTraining.py

@@ -14,7 +14,7 @@
 from testUtils.codeGenerate import generateTrainingTestNetwork
 from testUtils.platformMapping import mapDeployer, mapPlatform, setupMemoryPlatform
 from testUtils.testRunner import TestGeneratorArgumentParser
-from testUtils.tilingUtils import SBTiler
+from testUtils.tilingUtils import TrainingSBTiler
 from testUtils.typeMapping import inferTypeAndOffset
 
 from Deeploy.AbstractDataTypes import PointerClass
@@ -207,11 +207,11 @@ def generateTiledTrainingNetwork(args) -> None:
         AnnotateDefaultMemoryLevel(memoryHierarchy),
     ])
 
-    # 9. Wrap with tiler (SBTiler only — DBTiler conflicts with InPlaceAccumulatorV2 alias).
+    # 9. Wrap with tiler (TrainingSBTiler: SB strategy + extended input lifetimes for backward pass).
     unique_params = f"{args.dumpdir}_L1{args.l1}_L2{args.l2}_{args.defaultMemLevel}"
     testIdentifier = hashlib.md5(unique_params.encode()).hexdigest()[:16]
 
-    deployer = TilerDeployerWrapper(deployer, SBTiler, testName=testIdentifier, workDir=args.dumpdir)
+    deployer = TilerDeployerWrapper(deployer, TrainingSBTiler, testName=testIdentifier, workDir=args.dumpdir)
     deployer.tiler.visualizeMemoryAlloc = args.plotMemAlloc
     deployer.tiler.memoryAllocStrategy = args.memAllocStrategy
     deployer.tiler.searchStrategy = args.searchStrategy

DeeployTest/testUtils/tilingUtils.py

@@ -2,11 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import List, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from ortools.constraint_solver.pywrapcp import IntVar
 
 from Deeploy.DeeployTypes import NetworkContext, SubGraph, TransientBuffer
+from Deeploy.TilingExtension.MemoryConstraints import PatternMemoryConstraints
+from Deeploy.TilingExtension.MemoryScheduler import MemoryScheduler
 from Deeploy.TilingExtension.TilerExtension import Tiler
 from Deeploy.TilingExtension.TilerModel import TilerModel
 
@@ -43,3 +45,28 @@ class SBTiler(Tiler):
     def multiBufferStrategy(self, tilerModel: TilerModel, ctxt: NetworkContext, pattern: SubGraph, path: List[str],
                             hop: str, tensorName: str) -> Union[int, IntVar]:
         return 1
+
+
+class TrainingMemoryScheduler(MemoryScheduler):
+    """MemoryScheduler variant for training networks.
+
+    Extends input tensor lifetimes to the end of the full tiling schedule so
+    that forward-pass inputs remain live during the backward pass.
+    """
+
+    def _calculateLifetimes(
+            self, ctxt: NetworkContext, patternMemoryConstraint: PatternMemoryConstraints,
+            memoryLevel: str) -> Tuple[Dict[str, Tuple[int, int]], Dict]:
+        tensorLifetimeMap, tensorMap = super()._calculateLifetimes(ctxt, patternMemoryConstraint, memoryLevel)
+
+        maxStepIdx = len(patternMemoryConstraint.nodeConstraints)
+        for tensorName, lifetime in tensorLifetimeMap.items():
+            buffer = ctxt.lookup(tensorName)
+            if buffer.is_input:
+                tensorLifetimeMap[tensorName] = (0, maxStepIdx)
+
+        return tensorLifetimeMap, tensorMap
+
+
+class TrainingSBTiler(SBTiler):
+    memorySchedulerClass = TrainingMemoryScheduler