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()]
