Add Readme and custum partition

buddy-compiler · Oct 30, 2024 · 1d6e63d · 1d6e63d
1 parent 2594899
commit 1d6e63d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 8 deletions.
diff --git a/examples/BuddyLeNet/README.md b/examples/BuddyLeNet/README.md
@@ -50,11 +50,11 @@ $ export LENET_EXAMPLE_PATH=${BUDDY_MLIR_BUILD_DIR}/../examples/BuddyLeNet/
 
 ```bash
 $ cmake -G Ninja .. -DBUDDY_LENET_EXAMPLES=ON
-$ ninja buddy-lenet-run
+$ ninja buddy-lenet-{DEVICE_TYPE}-run
 $ cd bin
-$ ./buddy-lenet-run
+$ ./buddy-lenet-{DEVICE_TYPE}-run
 ```
-
+{DEVICE_TYPE} is device type, currently supporting cpu, heter.
 ## Debug the Lowering Pass Pipeline with Fake Parameters.
 
 ```bash

diff --git a/examples/BuddyLeNet/buddy-lenet-import.py b/examples/BuddyLeNet/buddy-lenet-import.py
@@ -36,6 +36,7 @@
 from buddy.compiler.graph.type import DeviceType
 from buddy.compiler.ops import tosa, gpu
 from buddy.compiler.graph.json_decoder import json_to_graph
+from buddy.compiler.graph.operation import *
 from model import LeNet
 
 # Retrieve the LeNet model path from environment variables.
@@ -73,9 +74,20 @@
 
 if type == "cpu":
     pattern_list = [simply_fuse]
-else:
-    pattern_list = [custom_partition]
-graph.fuse_ops(pattern_list)
+    graph.fuse_ops(pattern_list)
+elif type == "heter":
+    group = []
+    for i, op in enumerate(graph._body):
+        if isinstance(op, PlaceholderOp) or isinstance(op, OutputOp) or i == 25:
+            continue
+        group.append(op)
+        subgraph_name = "subgraph0"
+        graph.group_map_device[subgraph_name] = DeviceType.CPU
+        graph.op_groups[subgraph_name] = group
+    new_group = [graph._body[25]]
+    subgraph_name = "subgraph1"
+    graph.group_map_device[subgraph_name] = DeviceType.GPU
+    graph.op_groups[subgraph_name] = new_group
 path_prefix = os.path.dirname(os.path.abspath(__file__))
 driver = GraphDriver(graph)
 driver.subgraphs[0].lower_to_top_level_ir()

diff --git a/frontend/Python/graph/transform/fuse_ops.py b/frontend/Python/graph/transform/fuse_ops.py
@@ -86,10 +86,10 @@ def custom_partition(graph: Graph):
         if isinstance(op, PlaceholderOp) or isinstance(op, OutputOp) or i == 25:
             continue
         group.append(op)
-        subgraph_name = "subgraph1"
+        subgraph_name = "subgraph0"
         graph.group_map_device[subgraph_name] = DeviceType.CPU
         graph.op_groups[subgraph_name] = group
     new_group = [graph._body[25]]
-    subgraph_name = "subgraph0"
+    subgraph_name = "subgraph1"
     graph.group_map_device[subgraph_name] = DeviceType.GPU
     graph.op_groups[subgraph_name] = new_group