[Feature] Select graph from a multi-graph execution plan to launch.

tornado-api/src/main/java/module-info.java

@@ -50,4 +50,6 @@
     opens uk.ac.manchester.tornado.api.types.tensors;
     opens uk.ac.manchester.tornado.api.types;
     opens uk.ac.manchester.tornado.api.runtime;
+    exports uk.ac.manchester.tornado.api.plan.types;
+    opens uk.ac.manchester.tornado.api.plan.types;
 }

tornado-api/src/main/java/uk/ac/manchester/tornado/api/ExecutionPlanType.java

@@ -22,6 +22,7 @@
 import uk.ac.manchester.tornado.api.plan.types.OffPrintKernel;
 import uk.ac.manchester.tornado.api.plan.types.OffProfiler;
 import uk.ac.manchester.tornado.api.plan.types.OffThreadInfo;
+import uk.ac.manchester.tornado.api.plan.types.WithAllGraphs;
 import uk.ac.manchester.tornado.api.plan.types.WithBatch;
 import uk.ac.manchester.tornado.api.plan.types.WithClearProfiles;
 import uk.ac.manchester.tornado.api.plan.types.WithCompilerFlags;
@@ -30,6 +31,7 @@
 import uk.ac.manchester.tornado.api.plan.types.WithDevice;
 import uk.ac.manchester.tornado.api.plan.types.WithDynamicReconfiguration;
 import uk.ac.manchester.tornado.api.plan.types.WithFreeDeviceMemory;
+import uk.ac.manchester.tornado.api.plan.types.WithGraph;
 import uk.ac.manchester.tornado.api.plan.types.WithGridScheduler;
 import uk.ac.manchester.tornado.api.plan.types.WithMemoryLimit;
 import uk.ac.manchester.tornado.api.plan.types.WithPrintKernel;
@@ -39,12 +41,11 @@
 import uk.ac.manchester.tornado.api.plan.types.WithWarmUp;
 
 public abstract sealed class ExecutionPlanType extends TornadoExecutionPlan //
-        permits OffConcurrentDevices, OffMemoryLimit, OffPrintKernel, //
-        OffProfiler, OffThreadInfo, WithWarmUp, WithBatch, WithClearProfiles,  //
-        WithCompilerFlags, WithConcurrentDevices, WithDefaultScheduler, //
-        WithDevice, WithDynamicReconfiguration, WithFreeDeviceMemory, //
-        WithGridScheduler, WithMemoryLimit, WithPrintKernel, WithProfiler, //
-        WithResetDevice, WithThreadInfo {
+        permits OffConcurrentDevices, OffMemoryLimit, OffPrintKernel, OffProfiler, //
+        OffThreadInfo, WithAllGraphs, WithBatch, WithClearProfiles, WithCompilerFlags,  //
+        WithConcurrentDevices, WithDefaultScheduler, WithDevice, WithDynamicReconfiguration, //
+        WithFreeDeviceMemory, WithGraph, WithGridScheduler, WithMemoryLimit, WithPrintKernel, //
+        WithProfiler, WithResetDevice, WithThreadInfo, WithWarmUp { //
 
     public ExecutionPlanType(TornadoExecutionPlan parentNode) {
 

tornado-api/src/main/java/uk/ac/manchester/tornado/api/TornadoExecutionPlan.java

@@ -17,9 +17,9 @@
  */
 package uk.ac.manchester.tornado.api;
 
-import java.util.Objects;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicLong;
 
 import uk.ac.manchester.tornado.api.common.TornadoDevice;
@@ -32,6 +32,7 @@
 import uk.ac.manchester.tornado.api.plan.types.OffPrintKernel;
 import uk.ac.manchester.tornado.api.plan.types.OffProfiler;
 import uk.ac.manchester.tornado.api.plan.types.OffThreadInfo;
+import uk.ac.manchester.tornado.api.plan.types.WithAllGraphs;
 import uk.ac.manchester.tornado.api.plan.types.WithBatch;
 import uk.ac.manchester.tornado.api.plan.types.WithClearProfiles;
 import uk.ac.manchester.tornado.api.plan.types.WithCompilerFlags;
@@ -40,6 +41,7 @@
 import uk.ac.manchester.tornado.api.plan.types.WithDevice;
 import uk.ac.manchester.tornado.api.plan.types.WithDynamicReconfiguration;
 import uk.ac.manchester.tornado.api.plan.types.WithFreeDeviceMemory;
+import uk.ac.manchester.tornado.api.plan.types.WithGraph;
 import uk.ac.manchester.tornado.api.plan.types.WithGridScheduler;
 import uk.ac.manchester.tornado.api.plan.types.WithMemoryLimit;
 import uk.ac.manchester.tornado.api.plan.types.WithPrintKernel;
@@ -165,6 +167,34 @@ public TornadoExecutionResult execute() {
         return executionResult;
     }
 
+    /**
+     * Select a graph from the {@link TornadoExecutionPlan} to execute.
+     * This method allows developers to select a specific graph from the
+     * execution plan to launch. Developers can choose which graph from
+     * the input list to use (passed in the constructor).
+     *
+     * 
+     * @since 1.0.9
+     * @param graphIndex
+     * @return {@link TornadoExecutionPlan}
+     */
+    public TornadoExecutionPlan withGraph(int graphIndex) {
+        tornadoExecutor.selectGraph(graphIndex);
+        return new WithGraph(this, graphIndex);
+    }
+
+    /**
+     * Select all graphs from the {@link TornadoExecutionPlan}. This method
+     * has an effect if the {@link #withGraph(int)} method was invoked.
+     *
+     * @since 1.0.9
+     * @return {@link TornadoExecutionPlan}
+     */
+    public TornadoExecutionPlan withAllGraphs() {
+        tornadoExecutor.selectAll();
+        return new WithAllGraphs(this);
+    }
+
     /**
      * It invokes the JIT compiler for all immutable tasks-graphs associated to an
      * executor.
@@ -522,4 +552,5 @@ public TornadoExecutionResult getPlanResult(int index) {
         }
         return planResults.get(index);
     }
+
 }

tornado-api/src/main/java/uk/ac/manchester/tornado/api/TornadoExecutor.java

@@ -17,16 +17,16 @@
  */
 package uk.ac.manchester.tornado.api;
 
-import uk.ac.manchester.tornado.api.common.TornadoDevice;
-import uk.ac.manchester.tornado.api.enums.TornadoVMBackendType;
-import uk.ac.manchester.tornado.api.exceptions.TornadoRuntimeException;
-import uk.ac.manchester.tornado.api.runtime.ExecutorFrame;
-
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 
+import uk.ac.manchester.tornado.api.common.TornadoDevice;
+import uk.ac.manchester.tornado.api.enums.TornadoVMBackendType;
+import uk.ac.manchester.tornado.api.exceptions.TornadoRuntimeException;
+import uk.ac.manchester.tornado.api.runtime.ExecutorFrame;
+
 /**
  * Executor Class to dispatch Tornado Task-Graphs. An executor plan
  * {@link TornadoExecutionPlan} contains an executor object, which in turn,
@@ -36,6 +36,7 @@
 class TornadoExecutor {
 
     private final List<ImmutableTaskGraph> immutableTaskGraphList;
+    private List<ImmutableTaskGraph> subgraphList;
 
     TornadoExecutor(ImmutableTaskGraph... immutableTaskGraphs) {
         immutableTaskGraphList = new ArrayList<>();
@@ -222,4 +223,22 @@ long getTotalDeviceMemoryUsage() {
     long getCurrentDeviceMemoryUsage() {
         return immutableTaskGraphList.stream().mapToLong(ImmutableTaskGraph::getCurrentDeviceMemoryUsage).sum();
     }
+
+    void selectGraph(int graphIndex) {
+        if (subgraphList == null) {
+            subgraphList = new ArrayList<>();
+            immutableTaskGraphList.forEach(g -> Collections.addAll(subgraphList, g));
+        }
+        immutableTaskGraphList.clear();
+        Collections.addAll(immutableTaskGraphList, subgraphList.get(graphIndex));
+    }
+
+    void selectAll() {
+        if (subgraphList == null) {
+            return;
+        }
+        immutableTaskGraphList.clear();
+        subgraphList.forEach(g -> Collections.addAll(immutableTaskGraphList, g));
+        subgraphList = null;
+    }
 }

tornado-api/src/main/java/uk/ac/manchester/tornado/api/plan/types/WithAllGraphs.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2024, APT Group, Department of Computer Science,
+ * The University of Manchester.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package uk.ac.manchester.tornado.api.plan.types;
+
+import uk.ac.manchester.tornado.api.ExecutionPlanType;
+import uk.ac.manchester.tornado.api.TornadoExecutionPlan;
+
+public final class WithAllGraphs extends ExecutionPlanType {
+
+    public WithAllGraphs(TornadoExecutionPlan parent) {
+        super(parent);
+    }
+
+    @Override
+    public String toString() {
+        return parentLink.toString() + "\n -> withAllGraphs ";
+    }
+
+}

tornado-api/src/main/java/uk/ac/manchester/tornado/api/plan/types/WithGraph.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2024, APT Group, Department of Computer Science,
+ * The University of Manchester.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package uk.ac.manchester.tornado.api.plan.types;
+
+import uk.ac.manchester.tornado.api.ExecutionPlanType;
+import uk.ac.manchester.tornado.api.TornadoExecutionPlan;
+
+public final class WithGraph extends ExecutionPlanType {
+
+    private final int graphIndexSelected;
+
+    public WithGraph(TornadoExecutionPlan parent, int graphIndexSelected) {
+        super(parent);
+        this.graphIndexSelected = graphIndexSelected;
+    }
+
+    @Override
+    public String toString() {
+        return parentLink.toString() + "\n -> withGraph(" + graphIndexSelected + ") ";
+    }
+}

tornado-unittests/src/main/java/uk/ac/manchester/tornado/unittests/executor/TestExecutor.java

@@ -22,6 +22,7 @@
 import static org.junit.Assert.fail;
 
 import java.util.Arrays;
+import java.util.Random;
 
 import org.junit.Test;
 
@@ -303,5 +304,104 @@ public void test05() throws TornadoExecutionPlanException {
         }
     }
 
+    /**
+     * Test Multi-Graphs in an execution plan. An execution plan can hold and launch
+     * multiple immutable task graphs. This tests shows how to execute individual immutable
+     * task graphs in any order.
+     * 
+     * @throws TornadoExecutionPlanException
+     */
+    @Test
+    public void test06() throws TornadoExecutionPlanException {
+        int numElements = 16;
+        IntArray a = new IntArray(numElements);
+        IntArray b = new IntArray(numElements);
+        IntArray c = new IntArray(numElements);
+
+        a.init(1);
+        b.init(2);
+
+        TaskGraph tg1 = new TaskGraph("s0") //
+                .transferToDevice(DataTransferMode.FIRST_EXECUTION, a, b) //
+                .task("t0", TestHello::add, a, b, c) //
+                .transferToHost(DataTransferMode.EVERY_EXECUTION, c);
+
+        TaskGraph tg2 = new TaskGraph("s1") //
+                .transferToDevice(DataTransferMode.FIRST_EXECUTION, a, b) //
+                .task("t1", TestHello::add, a, b, c) //
+                .transferToHost(DataTransferMode.EVERY_EXECUTION, c);
+
+        try (TornadoExecutionPlan executionPlan = new TornadoExecutionPlan(tg1.snapshot(), tg2.snapshot())) {
+
+            // Select graph 1 (tg2) to execute
+            // Once selected, every time we call the execute method,
+            // TornadoVM will launch the passed task-graph.
+            executionPlan.withGraph(1).execute();
+            for (int i = 0; i < c.getSize(); i++) {
+                assertEquals(a.get(i) + b.get(i), c.get(i));
+            }
+
+            // Select the graph 0 (tg1) to execute
+            executionPlan.withGraph(0).execute();
+            for (int i = 0; i < c.getSize(); i++) {
+                assertEquals(a.get(i) + b.get(i), c.get(i));
+            }
+
+            // Select all graphs (tg1 and tg2) to execute.
+            // Since we selected individual task-graphs, we should be
+            // able to reverse this action and invoke all task-graph
+            // again. This is achieved with the `withAllGraphs` from the
+            // execution plan.
+            executionPlan.withAllGraphs().execute();
+        }
+    }
+
+    /**
+     * Test Multi-Graphs in an execution plan. Based on the {@link #test06()}, this test checks
+     * task-graphs with different tasks and data.
+     *
+     * @throws TornadoExecutionPlanException
+     */
+    @Test
+    public void test07() throws TornadoExecutionPlanException {
+        int numElements = 16;
+        IntArray a = new IntArray(numElements);
+        IntArray b = new IntArray(numElements);
+        IntArray c = new IntArray(numElements);
+
+        Random r = new Random();
+        for (int i = 0; i < a.getSize(); i++) {
+            a.set(i, r.nextInt());
+            b.set(i, r.nextInt());
+        }
+
+        TaskGraph tg1 = new TaskGraph("s0") //
+                .transferToDevice(DataTransferMode.EVERY_EXECUTION, a, b) //
+                .task("t0", TestHello::add, a, b, c) //
+                .transferToHost(DataTransferMode.EVERY_EXECUTION, c);
+
+        // Set dependency from graph tg1 to tg2 
+
+        TaskGraph tg2 = new TaskGraph("s1") //
+                .transferToDevice(DataTransferMode.EVERY_EXECUTION, c) //
+                .task("t1", TestHello::compute, c) //
+                .transferToHost(DataTransferMode.EVERY_EXECUTION, c);
+
+        try (TornadoExecutionPlan executionPlan = new TornadoExecutionPlan(tg1.snapshot(), tg2.snapshot())) {
+
+            // Select graph 0 (tg1) and run
+            executionPlan.withGraph(0).execute();
+            for (int i = 0; i < c.getSize(); i++) {
+                assertEquals(a.get(i) + b.get(i), c.get(i));
+            }
+
+            // Select the graph 1 (tg2) and run
+            executionPlan.withGraph(1).execute();
+            for (int i = 0; i < c.getSize(); i++) {
+                assertEquals((a.get(i) + b.get(i)) * 2, c.get(i));
+            }
+        }
+    }
+
     // CHECKSTYLE:ON
 }