[refactor] Remove unused fields and methods across FFN layer implementations

Author: orionpapadakis

src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractFFNLayers.java

@@ -11,12 +11,10 @@
 /**
  * Abstract base class for all FFN (Feed-Forward Network) layer implementations.
  *
- * Extends AbstractLayer and adds FFN-specific methods: - getFfnLayerTaskGraphs(): Returns task graphs for all transformer layers - getLastTaskGraphID(): Tracks the ID of the last task graph
+ * Each subclass builds N ImmutableTaskGraphs (one per FFN layer) via
+ * {@link #setupFFNLayerTaskGraphs}, covering RMSNorm, Attention, and FFN computations.
  *
- * All model-specific FFN layers extend this: - LlamaFP16FFNLayers, Qwen2FP16FFNLayers, Qwen3FP16FFNLayers, Phi3FP16FFNLayers - LlamaQ8_0FFNLayers, Qwen2Q8_0FFNLayers, Qwen3Q8_0FFNLayers,
- * Phi3Q8_0FFNLayers
- *
- * Used by FP16LayerPlanner and Q8_0LayerPlanner template methods for type-safe polymorphic access to any FFN layer implementation.
+ * Model-specific subclasses: Llama, Qwen2, Qwen3, Phi3, Granite — each in FP16 and Q8_0 variants.
  */
 public abstract class AbstractFFNLayers extends AbstractLayer {
 
@@ -55,7 +53,7 @@ protected AbstractFFNLayers(String taskGraphName, State state, Weights weights,
      *
      * @return List of immutable task graphs (one per transformer layer)
      */
-    public abstract List<ImmutableTaskGraph> getFfnLayerTaskGraphs();
+    public abstract List<ImmutableTaskGraph> getFFNLayerTaskGraphs();
 
     /**
      * Returns the TaskGraph ID of the last FFN layer.

src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractLayer.java

@@ -4,15 +4,11 @@
 import org.beehive.gpullama3.inference.weights.Weights;
 import org.beehive.gpullama3.model.Configuration;
 import uk.ac.manchester.tornado.api.GridScheduler;
-import uk.ac.manchester.tornado.api.ImmutableTaskGraph;
 import uk.ac.manchester.tornado.api.KernelContext;
 import uk.ac.manchester.tornado.api.TaskGraph;
 
-import java.util.ArrayList;
-import java.util.List;
-
 /**
- * Minimal base with common fields/utilities so subclasses compile cleanly. Adjust or remove fields if they already exist in your project.
+ * Abstract base class for Activations, FFN Layers, and Logits.
  */
 public abstract class AbstractLayer {
 
@@ -22,17 +18,10 @@ public abstract class AbstractLayer {
 
     protected final Weights weights;
     protected final Configuration config;
-    /** Often a small context/config buffer passed into kernels. Use your real type if available. */
+    protected final State state;
     protected final KernelContext context = new KernelContext();
-    /** Collected snapshots for scheduling / debugging. */
-    protected final List<ImmutableTaskGraph> taskGraphs = new ArrayList<>();
-    /** Optional: track the "main" task graph for the layer if one exists. */
-    protected TaskGraph taskGraph;
-    /** Shared runtime objects (exposed because kernels expect them). */
-    protected State state;
 
     protected AbstractLayer(String taskGraphName, State state, Weights weights, Configuration config) {
-        this.taskGraph = null;
         this.state = state;
         this.weights = weights;
         this.config = config;
@@ -48,13 +37,7 @@ protected static <T> T requireWeightsType(Object weights, Class<T> expectedType,
 
     public abstract GridScheduler updateGridScheduler(GridScheduler scheduler);
 
-    public abstract GridScheduler getGridScheduler();
-
-    public abstract TaskGraph getTaskGraph();
-
-    public abstract ImmutableTaskGraph getImmutableTaskGraph();
-
-    /** Allow subclasses to override if they need custom transfers. */
+    /** Allow subclasses to override if they need custom data transfers. */
     protected TaskGraph configureLayerDataTransfers(TaskGraph tg, int layerIndex) {
         return tg;
     }

src/main/java/org/beehive/gpullama3/tornadovm/layers/Activation.java

@@ -50,17 +50,10 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
         return scheduler;
     }
 
-    @Override
-    public GridScheduler getGridScheduler() {
-        return null;
-    }
-
-    @Override
     public TaskGraph getTaskGraph() {
         return activationUpdate;
     }
 
-    @Override
     public ImmutableTaskGraph getImmutableTaskGraph() {
         return activationUpdate.snapshot();
     }

src/main/java/org/beehive/gpullama3/tornadovm/layers/ActivationGranite.java

@@ -41,28 +41,3 @@ public ActivationGranite(String taskGraphHandle, State state, Weights weights, G
         }
         // @formatter:on
     }
-
-    @Override
-    public GridScheduler updateGridScheduler(GridScheduler scheduler) {
-        WorkerGrid worker = new WorkerGrid1D(config.dim());
-        worker.setLocalWork(128, 1, 1);
-        scheduler.addWorkerGrid("activationUpdate.updateX", worker);
-        return scheduler;
-    }
-
-    @Override
-    public GridScheduler getGridScheduler() {
-        return null;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return activationUpdate;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return activationUpdate.snapshot();
-    }
-
-}

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/GraniteFP16FFNLayers.java

@@ -21,8 +21,6 @@
 
 public class GraniteFP16FFNLayers extends AbstractFFNLayers {
 
-    TaskGraph ffnTaskGraphs;
-    GridScheduler scheduler;
     List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public GraniteFP16FFNLayers(String taskGraph, State state, Weights weights, GraniteConfiguration config, SchedulerType schedulerType) {
@@ -64,22 +62,7 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
         return tornadoForwardScheduler;
     }
 
-    @Override
-    public GridScheduler getGridScheduler() {
-        return scheduler;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return ffnTaskGraphs;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return null;
-    }
-
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java

@@ -20,9 +20,7 @@
 
 public class LlamaFP16FFNLayers extends AbstractFFNLayers {
 
-    TaskGraph ffnTaskGraphs;
-    GridScheduler scheduler;
-    List<ImmutableTaskGraph> ffnLayerTaskGraphs;
+    private List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public LlamaFP16FFNLayers(String taskGraph, State state, Weights weights, Configuration config, SchedulerType schedulerType) {
         super(taskGraph, state, weights, config, schedulerType);
@@ -64,21 +62,7 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
     }
 
     @Override
-    public GridScheduler getGridScheduler() {
-        return scheduler;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return ffnTaskGraphs;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return null;
-    }
-
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Phi3FP16FFNLayers.java

@@ -33,8 +33,6 @@ public class Phi3FP16FFNLayers extends AbstractFFNLayers {
     private final Phi3Configuration phi3Config;
     // Phi3-specific dimension for combined QKV buffer
     private final int opSize;
-    TaskGraph ffnLayerTaskGraph;
-    GridScheduler scheduler;
     List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public Phi3FP16FFNLayers(String taskGraphName, Phi3State state, Phi3TornadoWeights weights, Phi3Configuration config, SchedulerType schedulerType) {
@@ -88,22 +86,7 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) {
         return gridScheduler;
     }
 
-    @Override
-    public GridScheduler getGridScheduler() {
-        return scheduler;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return ffnLayerTaskGraph;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return null;
-    }
-
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen2FP16FFNLayers.java

@@ -34,8 +34,6 @@ public class Qwen2FP16FFNLayers extends AbstractFFNLayers {
     // Typed references to Qwen2-specific state and config
     private final Qwen2State qwen2State;
     private final Qwen2Configuration qwen2Config;
-    TaskGraph ffnLayerTaskGraph;
-    GridScheduler scheduler;
     List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public Qwen2FP16FFNLayers(String taskGraphName, Qwen2State state, Qwen2TornadoWeights weights, Qwen2Configuration config, SchedulerType schedulerType) {
@@ -116,22 +114,7 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
         return tornadoForwardScheduler;
     }
 
-    @Override
-    public GridScheduler getGridScheduler() {
-        return scheduler;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return ffnLayerTaskGraph;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return null;
-    }
-
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java

@@ -38,8 +38,6 @@ public class Qwen3FP16FFNLayers extends AbstractFFNLayers {
     private final int nEmbdHead;
     private final int nEmbdGqa;
     private final int gqa;
-    TaskGraph ffnLayerTaskGraph;
-    GridScheduler scheduler;
     List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public Qwen3FP16FFNLayers(String taskGraphName, Qwen3State state, Qwen3TornadoWeights weights, Qwen3Configuration config, SchedulerType schedulerType) {
@@ -102,22 +100,7 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) {
         return gridScheduler;
     }
 
-    @Override
-    public GridScheduler getGridScheduler() {
-        return scheduler;
-    }
-
-    @Override
-    public TaskGraph getTaskGraph() {
-        return ffnLayerTaskGraph;
-    }
-
-    @Override
-    public ImmutableTaskGraph getImmutableTaskGraph() {
-        return null;
-    }
-
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/GraniteQ8_0FFNLayers.java

@@ -1,11 +1,7 @@
 package org.beehive.gpullama3.tornadovm.layers.type.q8_0;
 
 import org.beehive.gpullama3.inference.state.GraniteState;
-import org.beehive.gpullama3.inference.state.LlamaState;
 import org.beehive.gpullama3.inference.weights.tornado.GraniteTornadoWeights;
-import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights;
-import org.beehive.gpullama3.model.Configuration;
-import org.beehive.gpullama3.model.granite.Granite;
 import org.beehive.gpullama3.model.granite.GraniteConfiguration;
 import org.beehive.gpullama3.tornadovm.kernels.GraniteKernels;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered;
@@ -23,7 +19,6 @@
 
 public class GraniteQ8_0FFNLayers extends AbstractFFNLayers {
 
-    GridScheduler scheduler;
     List<ImmutableTaskGraph> ffnLayerTaskGraphs;
 
     public GraniteQ8_0FFNLayers(String taskGraphName, GraniteState state, GraniteTornadoWeights weights, GraniteConfiguration config, SchedulerType schedulerType) {
@@ -314,7 +309,7 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
         return tornadoForwardScheduler;
     }
 
-    public List<ImmutableTaskGraph> getFfnLayerTaskGraphs() {
+    public List<ImmutableTaskGraph> getFFNLayerTaskGraphs() {
         return ffnLayerTaskGraphs;
     }