diff --git a/examples/BuddyLeNet/README.md b/examples/BuddyLeNet/README.md index b9b0c44a5f..a742e2451d 100644 --- 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 index e77a53686d..d876c0bd43 100644 --- 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 index 7bfd2e8f98..256437f154 100644 --- 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