def test_check_attribute_match(): global op1, op2 op2.attributes.add(TestAttribute(op2)) assert not check_attribute_match(op1, TestAttribute) assert check_attribute_match(op2, TestAttribute)
def _to_serializable_(self): placeholders = self.get_all_placeholders() unresolved_value_lists = [] for kernel in self.kernels: unresolved_value_list = kernel.exec_info.unresolved_value_list or [] unresolved_value_lists.append([{ "offset": v[0], "placeholder": v[1] } for v in unresolved_value_list]) return { "weight_encoding": self.constants_encoding, "memory_layout": self.memory_layout, "placeholders": placeholders, "unresolved_value_lists": unresolved_value_lists, "inputs": [ v.parameters["name"] for v in self.inputs if not traverse.check_attribute_match(v, Constant) ], "outputs": [v.parameters["name"] for v in self.outputs], "licenses": self.licenses }
def _to_serializable_(self): # placeholders = self.get_all_placeholders() placeholders = [] return { "converted_at": int(datetime.timestamp(datetime.now())), "inputs": [ v.parameters["name"] for v in self.inputs if not traverse.check_attribute_match(v, Constant) ], "outputs": [v.parameters["name"] for v in self.outputs], "memory_layout": self.memory_layout, "weight_encoding": self.constants_encoding, "placeholders": placeholders, "shader_sources": self.concat_kernel_sources(), "exec_infos": [kernel.exec_info for kernel in self.kernels], "constants_map": self.constants_map, "licenses": self.licenses, }
def __init__(self, op: Operator): self.delegate = lambda exp: exp # type: Callable[[str], str] self.has_inline = traverse.check_attribute_match(op, PostInlineInplace) if self.has_inline: post_inline_inplace = op.get_attribute(PostInlineInplace)[ 0] # type: PostInlineInplace if post_inline_inplace.injected is not None: self.delegate = post_inline_inplace.injected.injector
def _to_serializable_(self): return { "weight_allocation": self.constants_layout, "weight_encoding": self.constants_encoding, "variable_allocation": self.variables_layout, "inputs": [v.parameters["name"] for v in self.inputs if not traverse.check_attribute_match(v, Constant)], "outputs": [v.parameters["name"] for v in self.outputs], "licenses": self.licenses }
def _to_serializable_(self): return { "kernel_source": self.concat_kernel_sources(), "exec_infos": [kernel.exec_info for kernel in self.kernels], "weight_allocation": self.constants_layout, "weight_encoding": self.constants_encoding, "variable_allocation": self.variables_layout, "inputs": [v.parameters["name"] for v in self.inputs if not traverse.check_attribute_match(v, Constant)], "outputs": [v.parameters["name"] for v in self.outputs], "licenses": self.licenses }
def _to_serializable_(self): placeholders = self.get_all_placeholders() return { "kernel_source": self.concat_kernel_sources(), "exec_infos": [kernel.exec_info for kernel in self.kernels], "weight_encoding": self.constants_encoding, "memory_layout": self.memory_layout, "placeholders": placeholders, "inputs": [ self.memory_layout[v].name for v in self.inputs if not traverse.check_attribute_match(v, Constant) ], "outputs": [self.memory_layout[v].name for v in self.outputs], "licenses": self.licenses }
def _to_serializable_(self): # placeholders = self.get_all_placeholders() placeholders = [] return { "inputs": [ v.parameters["name"] for v in self.inputs if not traverse.check_attribute_match(v, Constant) ], "outputs": [v.parameters["name"] for v in self.outputs], "memory_layout": self.memory_layout, "weight_encoding": self.constants_encoding, "placeholders": placeholders, "shader_sources": self.concat_kernel_sources(), "exec_infos": [kernel.exec_info for kernel in self.kernels], "variables": { v.name: { "variable_size": v.size, "allocation_name": a.name } for v, a in self.allocations.items() }, # type: Variable, generator.WebGLAllocation "allocations": { a.name: { "allocation_size": a.size, "channel_mode": a.channel_mode.name } for a in self.allocations.values() }, "constants_map": self.constants_map, "licenses": self.licenses, }
def get_lifetime(graph: Graph, ops: List[Operator], variables: List[Variable]): LIFETIME_FOREVER = len(ops) + 1 lifetime = {} # type: Dict[Variable, Tuple[int, int]] retain_count = {v: 0 for v in variables} # type: Dict[Variable, int] allocated = set() # type: Set[Variable] for var in variables: if isinstance(var, ConstantVariable): lifetime[var] = (0, LIFETIME_FOREVER) allocated.add(var) for var in graph.inputs: lifetime[var] = (0, LIFETIME_FOREVER) allocated.add(var) for t, op in enumerate(ops): for var in op.outputs.values(): if isinstance(var, ConstantVariable): continue if var not in allocated: flag_allocated = False if flags.optimize.OPTIMIZE and flags.optimize.OPTIMIZE_INPLACE_OPERATION \ and not flag_allocated \ and traverse.check_attribute_match(op, Inplace): # Inplace optimization inplace = op.get_attribute(Inplace)[0] # type: Inplace v_in = inplace.get_input( ) # Use memory allocated for input variable v_out = inplace.get_output() common_axes = set(v_in.order.axes).intersection( set(v_out.order.axes)) if len(v_in.input_to) == 1 and \ all(v_in.stride_dict[axis] == v_out.stride_dict[axis] for axis in common_axes) and \ v_in.output_from is not None and \ len(v_out.input_to) != 0: while "inplace_src" in v_in.parameters: v_in = v_in.parameters["inplace_src"] var.parameters["inplace_src"] = v_in retain_count[v_in] += len(var.input_to) allocated.add(var) flag_allocated = True if not flag_allocated: lifetime[var] = (t, LIFETIME_FOREVER) retain_count[var] = len(var.input_to) allocated.add(var) flag_allocated = True if not flag_allocated: raise ValueError("[Allocator] Memory Allocation Failed.") for var in op.inputs.values(): if isinstance(var, ConstantVariable) or var in graph.inputs: continue while "inplace_src" in var.parameters: var = var.parameters["inplace_src"] if retain_count[var] == 0: # var is temporally workspace memory lifetime[var] = (t, t + 1) else: retain_count[var] -= 1 if retain_count[var] == 0: # `t + 1` means that `var` will be released AFTER `op` will be finished. lifetime[var] = (lifetime[var][0], t + 1) for var in graph.outputs: lifetime[var] = (lifetime[var][0], LIFETIME_FOREVER) return lifetime
def _analyse_variable_lifetime(graph: Graph, ops: List[Operator], variables: List[Variable]): # 計算グラフを辿りながら、retain回数をカウントし、生存期間のテーブルを作成する LIFETIME_FOREVER = len(ops) + 1 analysis_table: Dict[Variable, AllocationAnalysisData] = {} retain_count: Dict[Variable, int] = {v: 0 for v in variables} allocated: Set[Variable] = set() for var in graph.inputs: if isinstance(var, ConstantVariable): continue analysis_table[var] = AllocationAnalysisData(var, 0, LIFETIME_FOREVER) for t, op in enumerate(ops): for var in op.outputs.values(): if isinstance(var, ConstantVariable): continue if var not in allocated: # 新規に確保する flag_allocated = False if flags.optimize.OPTIMIZE and flags.optimize.OPTIMIZE_INPLACE_OPERATION \ and not flag_allocated \ and traverse.check_attribute_match(op, Inplace): # Inplace処理 inplace = op.get_attribute(Inplace)[0] # type: Inplace # 入力のメモリをそのまま使う v_in = inplace.get_input() while "inplace_src" in v_in.parameters: v_in = v_in.parameters["inplace_src"] var.parameters["inplace_src"] = v_in retain_count[v_in] += len(var.input_to) allocated.add(var) flag_allocated = True if not flag_allocated: # 新しくメモリを確保 analysis_table[var] = AllocationAnalysisData( var, t, LIFETIME_FOREVER) retain_count[var] = len(var.input_to) allocated.add(var) flag_allocated = True if not flag_allocated: raise ValueError( "[WebGPUAllocator] Memory Allocation Failed.") for var in op.inputs.values(): if isinstance(var, ConstantVariable): continue while "inplace_src" in var.parameters: var = var.parameters["inplace_src"] retain_count[var] -= 1 if retain_count[var] == 0: # 解放はこの層が終わってから! analysis_table[var].end = t + 1 return [x for x in analysis_table.values()]