Refactor task graph configuration for clarity and maintainability

Author: mikepapadim

src/main/java/com/example/tornadovm/TornadoVMLayerPlanner.java

@@ -81,7 +81,6 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
                 .persistOnDevice(state.wrapX);
         taskGraphs.add(activationUpdate.snapshot());
 
-
         TaskGraph unifiedLayer = null;
         for (int layerIndex =0; layerIndex < config.numberOfLayers; layerIndex++) {
             unifiedLayer = new TaskGraph("layer_" + layerIndex);
@@ -135,7 +134,6 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
         }
 
         TaskGraph lastUnifiedLayer = unifiedLayer;
-
         TaskGraph logits = new TaskGraph("logits")
                 .consumeFromDevice(lastUnifiedLayer.getTaskGraphName(),
                         state.wrapX
@@ -186,18 +184,21 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
     private TaskGraph configureQuantizedMatrixVectorFinalWeight(TaskGraph logits) {
         switch (weights.weightType) {
             case Q8_0:
-                logits.task("projection", TransformerComputeKernels::matmulTornadoQ8Optimized, context, weights.wclsByteArray, state.wrapX, state.wrapLogits, config.dim);
+                logits.task("projection", TransformerComputeKernels::matmulTornadoQ8Optimized,  //
+                        context, weights.wclsByteArray, state.wrapX,  //
+                        state.wrapLogits, config.dim); //
                 break;
             case Q4_0:
-                logits.task("projection", TransformerComputeKernels::matmulTornadoQ4Optimized, context, weights.wclsByteArray, state.wrapX, state.wrapLogits, config.dim);
+                logits.task("projection", TransformerComputeKernels::matmulTornadoQ4Optimized, //
+                        context, weights.wclsByteArray, state.wrapX,  //
+                        state.wrapLogits, config.dim); //
                 break;
             default:
                 throw new UnsupportedOperationException("Unsupported weight quantization type: " + weights.weightType + ". Only Q8_0 and Q4_0 are supported.");
         }
         return logits;
     }
 
-    // @formatter:off
     /**
      * Configures data transfer operations for a specific layer in the neural network task graph.
      *
@@ -218,29 +219,21 @@ private TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int layerI
         // First layer: Transfer initial data to device (one-time transfer)
         if (layerIndex == 0) {
             // Transfer all attention-related data: query, key, value matrices and their caches
-            unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
-                    context, state.wrapXb, state.wrapXb2,
-                    state.wrapQ, state.wrapK, state.wrapV,
-                    state.wrapKeyCache, state.wrapValueCache,
-                    state.wrapAtt, state.wrapHb);
+            unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.positionHolder, state.temp, state.tempFFN); //
+            unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, //
+                    context, state.wrapXb, state.wrapXb2, //
+                    state.wrapQ, state.wrapK, state.wrapV, //
+                    state.wrapKeyCache, state.wrapValueCache, //
+                    state.wrapAtt, state.wrapHb); //
         } else {
             // Subsequent layers: Consume data already on device from previous layer
-            unifiedLayer.consumeFromDevice(context, state.wrapXb, state.wrapXb2,
-                    state.wrapQ, state.wrapK, state.wrapV,
-                    state.wrapKeyCache, state.wrapValueCache,
-                    state.wrapAtt, state.wrapHb
+            unifiedLayer.consumeFromDevice(context, state.wrapXb, state.wrapXb2, //
+                    state.wrapQ, state.wrapK, state.wrapV, //
+                    state.wrapKeyCache, state.wrapValueCache, //
+                    state.wrapAtt, state.wrapHb, //
+                    state.positionHolder //
             );
         }
-
-        // First layer: Transfer position and temp data (transferred every execution)
-        if ((layerIndex) == 0) {
-            // Transfer data that changes with each execution (position, temp buffers)
-            unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.positionHolder, state.temp, state.tempFFN);
-        } else {
-            // Subsequent layers: Only consume position data from device
-            unifiedLayer.consumeFromDevice(state.positionHolder);
-        }
-        // @formatter:on
         return unifiedLayer;
     }