Add allocation step for DMA future

Author: Xeratec

Deeploy/Targets/PULPOpen/DMA/L3Dma.py

@@ -21,6 +21,8 @@ class L3DmaFuture(Future):
 
     _deinitTemplate = NodeTemplate("")
 
+    _allocTemplate = NodeTemplate("")
+
     _waitTemplate = NodeTemplate("""
     % if comment:
     // ${comment}

Deeploy/Targets/PULPOpen/DMA/MchanDma.py

@@ -20,6 +20,13 @@ class MchanChannelFuture(Future):
 
     _deinitTemplate = NodeTemplate("")
 
+    _allocTemplate = NodeTemplate("""
+     % if comment:
+    // ${comment}
+    % endif
+    ${name} = mchan_channel_alloc();
+    """)
+
     _waitTemplate = NodeTemplate("""
     % if comment:
     // ${comment}
@@ -37,15 +44,13 @@ class MchanDma(AsyncDma):
         % if comment:
         // ${comment}
         % endif
-        ${future} = mchan_channel_alloc();
         mchan_transfer_1d(${cmd}, ${loc}, ${ext});
         """),
         2:
             NodeTemplate("""
         % if comment:
         // ${comment}
         % endif
-        ${future} = mchan_channel_alloc();
         mchan_transfer_2d_ext_strided(${cmd}, ${loc}, ${ext}, ${size_1d}, ${stride_2d});
         """),
     }

Deeploy/Targets/Snitch/DMA/SnitchDma.py

@@ -11,6 +11,7 @@
 class SnitchBarrierFuture(Future):
     _initTemplate = NodeTemplate("")
     _deinitTemplate = NodeTemplate("")
+    _allocTemplate = NodeTemplate("")
     _waitTemplate = NodeTemplate("""
     % if comment:
     // ${comment}
@@ -30,6 +31,8 @@ class SnitchFuture(Future):
 
     _deinitTemplate = NodeTemplate("")
 
+    _allocTemplate = NodeTemplate("")
+
     _waitTemplate = NodeTemplate("""
     % if comment:
     // ${comment}

Deeploy/TilingExtension/AsyncDma.py

@@ -16,6 +16,7 @@
 class Future:
 
     _initTemplate: NodeTemplate
+    _allocTemplate: NodeTemplate
     _deinitTemplate: NodeTemplate
     _waitTemplate: NodeTemplate
 
@@ -28,6 +29,9 @@ def _operatorRepresentation(self, comment: str = "") -> OperatorRepresentation:
     def init(self, comment: str = "") -> CodeSnippet:
         return CodeSnippet(self._initTemplate, self._operatorRepresentation(comment))
 
+    def alloc(self, comment: str = "") -> CodeSnippet:
+        return CodeSnippet(self._allocTemplate, self._operatorRepresentation(comment))
+
     def deinit(self, comment: str = "") -> CodeSnippet:
         return CodeSnippet(self._deinitTemplate, self._operatorRepresentation(comment))
 
@@ -105,12 +109,12 @@ def transfer(self,
                  strideLoc: Tuple[int, ...],
                  direction: DmaDirection,
                  future: Future,
-                 comment: str = "") -> List[CodeSnippet]:
+                 comment: str = "") -> Tuple[List[CodeSnippet], List[CodeSnippet], List[CodeSnippet]]:
         self.checkTransfer(ctxt, externalBuffer, localBuffer, shape, strideExt, strideLoc, direction)
         opRepr = self.transferOpRepr(externalBuffer, localBuffer, shape, strideExt, strideLoc, direction, future,
                                      comment)
         template = self._transferTemplates[len(shape)]
-        return [CodeSnippet(template, opRepr)]
+        return [future.alloc(comment)], [CodeSnippet(template, opRepr)], []
 
     def setup(self) -> List[CodeSnippet]:
         return []
@@ -122,6 +126,7 @@ def teardown(self) -> List[CodeSnippet]:
 class EmptyFuture(Future):
 
     _initTemplate = NodeTemplate("")
+    _allocTemplate = NodeTemplate("")
     _deinitTemplate = NodeTemplate("")
     _waitTemplate = NodeTemplate("")
 
@@ -158,23 +163,25 @@ def transfer(self,
                  strideLoc: Tuple[int, ...],
                  direction: DmaDirection,
                  future: Future,
-                 comment: str = "") -> List[CodeSnippet]:
+                 comment: str = "") -> Tuple[List[CodeSnippet], List[CodeSnippet], List[CodeSnippet]]:
         tmpFuture = self.dma.getFuture(future.name.removesuffix("_future"))
         callStack = []
-        callStack.append(tmpFuture.init())
-        callStack.extend(
-            self.dma.transfer(ctxt,
-                              externalBuffer,
-                              localBuffer,
-                              shape,
-                              strideExt,
-                              strideLoc,
-                              direction,
-                              tmpFuture,
-                              comment = comment))
-        callStack.append(tmpFuture.wait())
-        callStack.append(tmpFuture.deinit())
-        return callStack
+        callStack.append(tmpFuture.init(comment))
+        callStack.append(tmpFuture.alloc(comment))
+        _, dma_code, _ = self.dma.transfer(ctxt,
+                                           externalBuffer,
+                                           localBuffer,
+                                           shape,
+                                           strideExt,
+                                           strideLoc,
+                                           direction,
+                                           tmpFuture,
+                                           comment = comment)
+        callStack.extend(dma_code)
+        callStack.append(tmpFuture.wait(comment))
+        callStack.append(tmpFuture.deinit(comment))
+
+        return [], callStack, []
 
     def setup(self) -> List[CodeSnippet]:
         return self.dma.setup()
@@ -239,7 +246,7 @@ def transfer(self,
                  direction: DmaDirection,
                  future: Future,
                  strideExtPad: int = 0,
-                 comment: str = "") -> List[CodeSnippet]:
+                 comment: str = "") -> Tuple[List[CodeSnippet], List[CodeSnippet], List[CodeSnippet]]:
         transferRank = len(shape)
         kernelRank = self.nearestSupportedTransferRank(transferRank)
 
@@ -275,18 +282,19 @@ def transfer(self,
                     "offset": "ext_offset"
                 }))
 
-            callStack.extend(
-                self.dma.transfer(ctxt,
-                                  externalBufferOffseted,
-                                  localBufferOffseted,
-                                  shape[-kernelRank:],
-                                  strideExt[-kernelRank:],
-                                  strideLoc[-kernelRank:],
-                                  direction,
-                                  future,
-                                  comment = comment))
+            alloc_code, dma_code, deinit_code = self.dma.transfer(ctxt,
+                                                                  externalBufferOffseted,
+                                                                  localBufferOffseted,
+                                                                  shape[-kernelRank:],
+                                                                  strideExt[-kernelRank:],
+                                                                  strideLoc[-kernelRank:],
+                                                                  direction,
+                                                                  future,
+                                                                  comment = comment)
+
+            callStack.extend(dma_code)
             callStack.append(CodeSnippet(self.NestedForLoopCloseTemplate(nestedLoopDepth), {}))
-            return callStack
+            return alloc_code, callStack, deinit_code
         elif kernelRank == transferRank:
             return self.dma.transfer(ctxt,
                                      externalBuffer,

Deeploy/TilingExtension/CodeTransformationPasses/DoubleBufferingTilingCodeGeneration.py

@@ -131,6 +131,7 @@ def _tilingLoop(self, ctxt: NetworkContext, executionBlock: ExecutionBlock,
                                                              initialFuture,
                                                              math.prod(externalBufferShape),
                                                              comment = "DMA INITIAL")
+            initialDmaTransferCalls = [item for tup in initialDmaTransferCalls for item in tup]
             setupStatements.extend(initialDmaTransferCalls)
 
             referenceUpdate = self._generateExternalReferenceUpdate(ctxt, tensorName, rectangles, "TILING_I+1",

Deeploy/TilingExtension/CodeTransformationPasses/TilingCodeGeneration.py

@@ -117,6 +117,8 @@ def _generateDmaTransferCalls(self,
                                               math.prod(externalBuffer.shape,),
                                               comment = comment)
 
+        initSnippets = [item for tup in initSnippets for item in tup]
+
         templates = [snippet.template for snippet in initSnippets]
         opReprUpdates = [[] for _ in range(len(initSnippets))]
 
@@ -131,6 +133,8 @@ def _generateDmaTransferCalls(self,
                                               future,
                                               math.prod(externalBuffer.shape),
                                               comment = comment)
+
+            snippets = [item for tup in snippets for item in tup]
             for i, snippet in enumerate(snippets):
                 opReprUpdates[i].append(snippet.operatorRepresentation)