conv-allocator 3.17.1__py3-none-any.whl

conv_allocator/__init__.py

@@ -0,0 +1,18 @@
+# -------------------------------------------------------------------------------
+# (c) Copyright 2025 Sony Semiconductor Israel, Ltd. All rights reserved.
+#
+#      This software, in source or object form (the "Software"), is the
+#      property of Sony Semiconductor Israel Ltd. (the "Company") and/or its
+#      licensors, which have all right, title and interest therein, You
+#      may use the Software only in accordance with the terms of written
+#      license agreement between you and the Company (the "License").
+#      Except as expressly stated in the License, the Company grants no
+#      licenses by implication, estoppel, or otherwise. If you are not
+#      aware of or do not agree to the License terms, you may not use,
+#      copy or modify the Software. You may use the source code of the
+#      Software only for your internal purposes and may not distribute the
+#      source code of the Software, any part thereof, or any derivative work
+#      thereof, to any third party, except pursuant to the Company's prior
+#      written consent.
+#      The Software is the confidential information of the Company.
+# -------------------------------------------------------------------------------

conv_allocator/allocator.py

@@ -0,0 +1,323 @@
+# -------------------------------------------------------------------------------
+# (c) Copyright 2025 Sony Semiconductor Israel, Ltd. All rights reserved.
+#
+#      This software, in source or object form (the "Software"), is the
+#      property of Sony Semiconductor Israel Ltd. (the "Company") and/or its
+#      licensors, which have all right, title and interest therein, You
+#      may use the Software only in accordance with the terms of written
+#      license agreement between you and the Company (the "License").
+#      Except as expressly stated in the License, the Company grants no
+#      licenses by implication, estoppel, or otherwise. If you are not
+#      aware of or do not agree to the License terms, you may not use,
+#      copy or modify the Software. You may use the source code of the
+#      Software only for your internal purposes and may not distribute the
+#      source code of the Software, any part thereof, or any derivative work
+#      thereof, to any third party, except pursuant to the Company's prior
+#      written consent.
+#      The Software is the confidential information of the Company.
+# -------------------------------------------------------------------------------
+import argparse
+import math
+import os
+import csv
+import json
+import time
+from dataclasses import dataclass
+from itertools import combinations
+from pathlib import Path
+
+from ortools.sat.python import cp_model
+from ortools.sat.python.cp_model import FIXED_SEARCH
+from uni_model import GraphDefToUniModelConverter
+from uni_model import UniModelToGraphDefConverter
+from collections import defaultdict
+from typing import List, Dict, Set
+import json
+from uni_model import UniModelPath
+
+from conv_allocator.um_accessors import has_scratch, has_allocation, get_allocation, get_schedule, get_scratch_cut, get_cut, \
+    set_allocation
+
+
+@dataclass(eq=False)
+class Block:
+    layer_name: str
+    output_index: int
+    dynamic_memory: int
+    static_memory: int
+    offset: int
+    tensor_id: str
+
+    def __repr__(self):
+        return (f"Block({self.layer_name}, {self.output_index}, "
+                f"tensorId={self.tensor_id}, dynMem={self.dynamic_memory}, offset={self.offset})")
+
+    def __eq__(self, other):
+        return isinstance(other, Block) and (self.tensor_id == other.tensor_id)
+
+    def __hash__(self):
+        return hash((self.layer_name, self.output_index))
+
+
+# Filter out layers without output and allocations
+def filter_relevant_nodes(nodes):
+    excluded_classes = {'UniLayerConst'}
+
+    def is_relevant(node):
+        class_name = node.__class__.__name__
+        if class_name in excluded_classes:
+            return False
+        return True
+
+    return [node for node in nodes if is_relevant(node)]
+
+
+def validate_allocation_data(nodes):
+    for node in nodes:
+        name = getattr(node, 'name', '<unknown>')
+
+        if not has_allocation(node):
+            raise ValueError(f"Node {name} missing 'allocations'")
+
+        try:
+            loaded_allocation = get_allocation(node)
+        except Exception:
+            raise ValueError(f"Node {name} has invalid 'allocations'")
+
+        if loaded_allocation.schedule is None:
+            raise ValueError(f"Node {name} missing 'schedule'")
+
+        if loaded_allocation.outputSizes is None:
+            raise ValueError(f"Node {name} missing 'outputSizes'")
+
+
+def down_scale(out, scale):
+    return math.ceil(out.dynamicMemory/scale)
+
+def parse_blocks_from_node(node, scale):
+    allocations_data = get_allocation(node)
+    blocks = [
+        Block(
+            out.layerName,
+            out.outputIndex,
+            down_scale(out, scale),
+            out.staticMemory,
+            out.offset,
+            out.tensorId
+        )
+        for out in allocations_data.outputSizes
+    ]
+    return blocks
+
+
+def create_fake_scratch_node(real_node):
+    class FakeScratchNode:
+        def __init__(self, source_node):
+            self.source_real_node = source_node
+            self.name = f"{source_node.name}_scratch"
+
+    return FakeScratchNode(real_node)
+
+
+def build_graph_cuts(unimodel, scale):
+    graph = unimodel.uni_graphs[0]
+    filtered_nodes = filter_relevant_nodes(graph.layer_nodes)
+    validate_allocation_data(filtered_nodes)
+
+    # Step 1: Sort by schedule
+    sorted_nodes = sorted(filtered_nodes, key=get_schedule)
+
+    # Step 2: Parse blocks from sorted nodes
+    all_blocks_by_node = {}
+    blocks_map = {}
+    for node in sorted_nodes:
+        blocks = parse_blocks_from_node(node, scale)
+        all_blocks_by_node[node] = blocks
+        for block in blocks:
+            blocks_map[block] = node
+
+    # Step 3: Handle scratch layers, preserving order
+    scratch_fake_nodes = []
+    ordered_nodes_with_scratch = []
+
+    for node in sorted_nodes:
+        ordered_nodes_with_scratch.append(node)
+        if has_scratch(node):
+            scratch_node = create_fake_scratch_node(node)
+            scratch_fake_nodes.append(scratch_node)
+            ordered_nodes_with_scratch.append(scratch_node)
+
+    # Step 4: Use precomputed graph cuts from node metadata
+    tensor_id_to_block = {block.tensor_id: block for block in blocks_map}
+    graph_cuts = []
+
+    for node in ordered_nodes_with_scratch:
+        if hasattr(node, 'source_real_node'):  # Scratch fake node
+            real_node = node.source_real_node
+            scratch_ids = get_scratch_cut(real_node)
+            block_set = {tensor_id_to_block[tid] for tid in scratch_ids}
+        else:
+            cut_ids = get_cut(node)
+            block_set = {tensor_id_to_block[tid] for tid in cut_ids}
+
+        graph_cuts.append(block_set)
+
+    return graph_cuts, blocks_map
+
+
+def solve_allocation(graph_cuts, blocks, time_limit):
+    """
+    Solve the memory allocation problem using CP-SAT solver.
+
+    Args:
+        graph_cuts: List of sets, each set contains block IDs that need to be in memory simultaneously
+        blocks: Dictionary mapping block ID to its size
+
+    Returns:
+        Dictionary mapping block ID to (start_time, end_time)
+    """
+
+    start_time = time.time()
+
+    # Estimate horizon (safe upper bound on end time)
+    horizon = int(sum(block.dynamic_memory for block in blocks) * 2)
+
+    model = cp_model.CpModel()
+
+    # Create variables for each block
+    block_dict = {}
+    for block in blocks.keys():
+        block_id = block.tensor_id
+        block_size = block.dynamic_memory
+
+        start_var = model.NewIntVar(0, horizon, f"start_{block_id}")
+        end_var = model.NewIntVar(0, horizon, f"end_{block_id}")
+        interval_var = model.NewIntervalVar(start_var, block_size, end_var, f"interval_{block_id}")
+
+        block_dict[block] = (start_var, end_var, interval_var)
+
+    # Solver setup
+    solver = cp_model.CpSolver()
+    solver.parameters.max_time_in_seconds = time_limit
+    solver.parameters.random_seed = 17
+    solver.parameters.randomize_search = False
+    solver.parameters.use_absl_random = False
+    solver.parameters.num_workers = 1
+
+    # fixed strategy - works with fixed_search
+    # activation_vars = [start for start, _, _ in block_dict.values()]
+    # model.AddDecisionStrategy(
+    #     activation_vars,
+    #     cp_model.CHOOSE_LOWEST_MIN,
+    #     cp_model.SELECT_MIN_VALUE
+    # )
+    # solver.parameters.search_branching = cp_model.FIXED_SEARCH
+
+
+    # Use AddNoOverlap per graph cut
+    for cut in graph_cuts:
+        cut_intervals = [block_dict[block][2] for block in cut]
+        model.AddNoOverlap(cut_intervals)
+
+    # Objective: minimize makespan (max end address)
+    obj_var = model.NewIntVar(0, horizon, "makespan")
+    model.AddMaxEquality(obj_var, [end for _, end, _ in block_dict.values()])
+    model.Minimize(obj_var)
+
+    # Solution callback
+    class SolutionCallback(cp_model.CpSolverSolutionCallback):
+        def __init__(self):
+            super().__init__()
+            self._solution_count = 0
+            self.latest_solution = None
+
+        def on_solution_callback(self):
+            self._solution_count += 1
+            self.latest_solution = {
+                block_id: (self.Value(start), self.Value(end))
+                for block_id, (start, end, _) in block_dict.items()
+            }
+
+    callback = SolutionCallback()
+
+    status = solver.Solve(model, callback)
+    solve_time = time.time() - start_time
+
+    results = callback.latest_solution if callback.latest_solution else {
+        block_id: (solver.Value(start), solver.Value(end))
+        for block_id, (start, end, _) in block_dict.items()
+    } if status in [cp_model.OPTIMAL, cp_model.FEASIBLE] else {}
+
+    if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
+        peak_memory = max(end for _, end in results.values())
+        print(f"Status: {solver.StatusName(status)}")
+        # print(f"Objective value: {solver.ObjectiveValue()}")
+        # print(f"Peak memory usage: {peak_memory}")
+        return results, solver, peak_memory, solve_time, True
+    else:
+        print(f"No solution found. Status: {solver.StatusName(status)}")
+        return results, solver, 0, solve_time, False
+
+def main():
+    parser = argparse.ArgumentParser(description="Memory allocation solver.")
+    parser.add_argument("--input_path", required=True, help="Path to input file")
+    parser.add_argument("--output_path", required=True, help="Path to output file")
+    parser.add_argument("--time_limit", type=int, required=True, help="Time limit in seconds")
+    parser.add_argument("--scale", type=int, required=True, help="block size scale factor for address alignment.")
+    args = parser.parse_args()
+
+    scale = args.scale
+    uni_model_loaded = GraphDefToUniModelConverter.convert(args.input_path)
+    graph_cuts, blocks_map = build_graph_cuts(uni_model_loaded, scale)
+    results_scaled, allocation, max_address_scaled, solve_time, solved = solve_allocation(graph_cuts, blocks_map, args.time_limit)
+    max_address = max_address_scaled * scale
+
+    if solved:
+        layer_to_block_map = defaultdict(list)
+        for k, v in blocks_map.items():
+            layer_to_block_map[v].append(k)
+
+        block_id_to_offset = defaultdict(list)
+        for block in results_scaled:
+            offset = results_scaled[block][0] * scale
+            block_id_to_offset[block.tensor_id] = offset
+
+        for node in uni_model_loaded.uni_graphs[0].layer_nodes:
+            if has_allocation(node):
+                # Load allocation object from accessor
+                alloc_data = get_allocation(node)
+
+                # Update the outputSizes offsets inside AllocationDataPerLayer
+                for output_alloc in alloc_data.outputSizes:
+                    output_alloc.offset = block_id_to_offset[output_alloc.tensorId]
+
+                # Write updated allocation object back to node
+                set_allocation(node, alloc_data)
+
+        # debug print addresses
+        # for node in uni_model_loaded.uni_graphs[0].layer_nodes:
+        #     if has_allocation(node):
+        #         alloc_data = get_allocation(node)
+        #         for output_alloc in alloc_data.outputSizes:
+        #             output_alloc.offset = block_id_to_offset[output_alloc.tensorId]
+        #             print(f"tid: {output_alloc.tensorId}, offset: {output_alloc.offset}")
+
+        print("Solution found")
+        print(f"max_address: {max_address}")
+        print(f"solve_time: {solve_time}")
+    else:
+        print("No solution found")
+
+    path = Path(args.output_path)
+
+    directory = str(path.parent)
+    ext = ".um.pb"
+    filename_wo_ext = os.path.basename(path)
+    if filename_wo_ext.endswith(ext):
+        filename_wo_ext = filename_wo_ext[:-len(ext)]
+
+    uniPath = UniModelPath(directory, filename_wo_ext)
+    UniModelToGraphDefConverter.convert(uni_model_loaded, uniPath)
+
+if __name__ == "__main__":
+    main()

conv_allocator/allocator_test.py

@@ -0,0 +1,57 @@
+# -------------------------------------------------------------------------------
+# (c) Copyright 2025 Sony Semiconductor Israel, Ltd. All rights reserved.
+#
+#      This software, in source or object form (the "Software"), is the
+#      property of Sony Semiconductor Israel Ltd. (the "Company") and/or its
+#      licensors, which have all right, title and interest therein, You
+#      may use the Software only in accordance with the terms of written
+#      license agreement between you and the Company (the "License").
+#      Except as expressly stated in the License, the Company grants no
+#      licenses by implication, estoppel, or otherwise. If you are not
+#      aware of or do not agree to the License terms, you may not use,
+#      copy or modify the Software. You may use the source code of the
+#      Software only for your internal purposes and may not distribute the
+#      source code of the Software, any part thereof, or any derivative work
+#      thereof, to any third party, except pursuant to the Company's prior
+#      written consent.
+#      The Software is the confidential information of the Company.
+# -------------------------------------------------------------------------------
+import sys
+import argparse
+from ortools.linear_solver import pywraplp
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser(description='Allocation optimizer')
+    parser.add_argument('--source', type=str, help='Source filename', default='')
+    parser.add_argument('--target', type=str, help='Target filename', default='')
+    parser.add_argument('--time', type=int, help='Allowed time in seconds', default=0)
+
+    return parser.parse_args(args)
+
+def hello_world(source_filename, target_filename, allowed_time):
+    print(f"Hello from conv-allocator!")
+    print(f"Source: {source_filename}")
+    print(f"Target: {target_filename}")
+    print(f"Allowed time: {allowed_time} seconds")
+
+    solver = pywraplp.Solver.CreateSolver('GLOP')
+    x = solver.NumVar(0, 10, 'x')
+    y = solver.NumVar(0, 10, 'y')
+    solver.Add(x + y <= 20)
+    solver.Maximize(x + y)
+    status = solver.Solve()
+    if status == pywraplp.Solver.OPTIMAL:
+        print(f"x = {x.solution_value()}, y = {y.solution_value()}")
+    else:
+        print("No solution found")
+
+def main(args=None):
+    if args is None:
+        # Use sys.argv if no args provided
+        args = sys.argv[1:]
+
+    parsed_args = parse_args(args)
+    hello_world(parsed_args.source, parsed_args.target, parsed_args.time)
+
+if __name__ == "__main__":
+    main()

conv_allocator/um_accessors.py

@@ -0,0 +1,113 @@
+# -------------------------------------------------------------------------------
+# (c) Copyright 2025 Sony Semiconductor Israel, Ltd. All rights reserved.
+#
+#      This software, in source or object form (the "Software"), is the
+#      property of Sony Semiconductor Israel Ltd. (the "Company") and/or its
+#      licensors, which have all right, title and interest therein, You
+#      may use the Software only in accordance with the terms of written
+#      license agreement between you and the Company (the "License").
+#      Except as expressly stated in the License, the Company grants no
+#      licenses by implication, estoppel, or otherwise. If you are not
+#      aware of or do not agree to the License terms, you may not use,
+#      copy or modify the Software. You may use the source code of the
+#      Software only for your internal purposes and may not distribute the
+#      source code of the Software, any part thereof, or any derivative work
+#      thereof, to any third party, except pursuant to the Company's prior
+#      written consent.
+#      The Software is the confidential information of the Company.
+# -------------------------------------------------------------------------------
+import json
+from dataclasses import dataclass, field
+from json import JSONEncoder, dumps
+from typing import Optional, List
+from uni_model import UniModel, UniGraph, UniLayer
+
+ALLOCATIONS = 'allocations'
+HAS_SCRATCH = 'hasScratch'
+OUTPUT_SIZES = 'outputSizes'
+
+
+def has_scratch(layer: UniLayer):
+    return layer.extended_attr.get(HAS_SCRATCH, False)
+
+
+def has_allocation(layer: UniLayer):
+    return layer.extended_attr.get(ALLOCATIONS, False)
+
+
+def get_allocation(layer: UniLayer):
+    return AllocationDataPerLayer.decode_from_json(layer.extended_attr.get(ALLOCATIONS))
+
+
+
+
+def get_schedule(node: UniLayer):
+    if has_allocation(node):
+        return get_allocation(node).schedule
+    else:
+        raise ValueError(f"Node {node.name} missing 'allocation'")
+
+
+def get_scratch_cut(node: UniLayer):
+    if has_allocation(node):
+        return get_allocation(node).scratchCutTensorIds
+    else:
+        raise ValueError(f"Node {node.name} missing 'allocation'")
+
+
+def get_cut(node: UniLayer):
+    if has_allocation(node):
+        return get_allocation(node).cutTensorIds
+    else:
+        raise ValueError(f"Node {node.name} missing 'allocation'")
+
+
+@dataclass
+class OutputAllocation:
+    tensorId: str
+    layerName: str
+    outputIndex: int
+    dynamicMemory: int
+    staticMemory: int
+    offset: Optional[int] = None
+
+
+@dataclass
+class AllocationDataPerLayer:
+    schedule: int
+    dynamicMemory: int
+    cutTensorIds: list[str] = field(default_factory=list)
+    scratchCutTensorIds: list[str] = field(default_factory=list)
+    outputSizes: list[OutputAllocation] = field(default_factory=list)
+
+    def encode_to_json(self) -> str:
+        d = self.__dict__.copy()
+        d[OUTPUT_SIZES] = [o.__dict__ for o in d[OUTPUT_SIZES]]
+        return json.dumps(d)
+
+    @staticmethod
+    def decode_from_json(j: str) -> 'AllocationDataPerLayer':
+        d = json.loads(j)
+        d[OUTPUT_SIZES] = [OutputAllocation(**o) for o in d[OUTPUT_SIZES]]
+        return AllocationDataPerLayer(**d)
+
+
+def set_allocation(layer: UniLayer, alloc_data: AllocationDataPerLayer):
+    layer.extended_attr[ALLOCATIONS] = alloc_data.encode_to_json()
+
+
+if __name__ == "__main__":
+    a = AllocationDataPerLayer(
+        schedule=4,
+        dynamicMemory=9,
+        cutTensorIds=["t1", "t2"],
+        scratchCutTensorIds=["t2"],
+        outputSizes=[
+            OutputAllocation(tensorId="t1", layerName="layerX", outputIndex=0, dynamicMemory=10, staticMemory=5)
+        ]
+    )
+    encoded = a.encode_to_json()
+    print("Serialized:", encoded)
+
+    decoded = AllocationDataPerLayer.decode_from_json(encoded)
+    print("Deserialized:", decoded)

conv_allocator-3.17.1.dist-info/METADATA

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: conv-allocator
+Version: 3.17.1
+Requires-Dist: ortools==9.9.3963
+Requires-Dist: protobuf==4.25.5
+Requires-Dist: uni-model==10.0.23
+

conv_allocator-3.17.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+conv_allocator/__init__.py,sha256=phGsoy1jdNsCmzQFKZEWB_Thbw633GKNXRHR4TKs_0w,1208
+conv_allocator/allocator.py,sha256=aXp_obbxqtrK3wq57N0WmUY2wUXcn2HdegOuxro_K20,11710
+conv_allocator/allocator_test.py,sha256=sCWtHJKG2kH2f_ppRdNALwhgucK-mWEJyM-JfnGkN3U,2524
+conv_allocator/um_accessors.py,sha256=-6AIZ3x53ILeLBZCsIMI2b-ywslwWGd-QWK7RCy7n9c,3832
+conv_allocator-3.17.1.dist-info/METADATA,sha256=hTlJGQXJS3iwodf0njPo_fbtWENquKMPNaDWjzkE0lk,159
+conv_allocator-3.17.1.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+conv_allocator-3.17.1.dist-info/entry_points.txt,sha256=hLQ1x1lLRtbJWuJG8CtSpOdaN28hbYAVd_07szV-cKU,64
+conv_allocator-3.17.1.dist-info/top_level.txt,sha256=JWue8wV0JXEagm9a0qky3c1sKcImcUrQoJ_JskNFiqQ,15
+conv_allocator-3.17.1.dist-info/RECORD,,

conv_allocator-3.17.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.45.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

conv_allocator-3.17.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+run-allocator = conv_allocator.allocator:main

conv_allocator-3.17.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+conv_allocator