def test_ignore_internal_bound_flag():
# +-2-3-4-5-+
# 1-+ +-6-
# +---------+
n1 = Node("n1")
n2 = Node("n2")
n3 = Node("n3")
n4 = Node("n4")
n5 = Node("n5")
n6 = Node("n6")
n2.append_prev(n1)
n3.append_prev(n2)
n4.append_prev(n3)
n5.append_prev(n4)
n6.append_prev(n1)
n6.append_prev(n5)
# noinspection PyTypeChecker
graph = Graph([n1, n3], [n4, n6])
result = listup_nodes(graph,
ignore_internal_input_bound=False,
ignore_internal_output_bound=False)
assert len(result) == 4 and set(result) == {n1, n3, n4, n6}
assert result.index(n1) < result.index(n6)
assert result.index(n3) < result.index(n4)
result = listup_nodes(graph,
ignore_internal_input_bound=False,
ignore_internal_output_bound=True)
assert len(result) == 5 and set(result) == {n1, n3, n4, n5, n6}
assert result.index(n1) < result.index(n6)
assert result.index(n3) < result.index(n4)
assert result.index(n4) < result.index(n5)
assert result.index(n5) < result.index(n6)
result = listup_nodes(graph,
ignore_internal_input_bound=True,
ignore_internal_output_bound=False)
assert len(result) == 5 and set(result) == {n1, n2, n3, n4, n6}
assert result.index(n1) < result.index(n2)
assert result.index(n2) < result.index(n3)
assert result.index(n3) < result.index(n4)
assert result.index(n1) < result.index(n6)
result = listup_nodes(graph,
ignore_internal_input_bound=True,
ignore_internal_output_bound=True)
assert len(result) == 6 and set(result) == {n1, n2, n3, n4, n5, n6}
assert result.index(n1) < result.index(n2)
assert result.index(n2) < result.index(n3)
assert result.index(n3) < result.index(n4)
assert result.index(n4) < result.index(n5)
assert result.index(n5) < result.index(n6)
assert result.index(n1) < result.index(n6)
def generate(cls, graph: Graph, **kwargs):
graph, _ = WebGLOptimizeRule().optimize(graph)
if flags.DEBUG:
traverse.dump(graph)
with open("cg.dot", "w") as f:
f.write(traverse.dump_dot(graph))
memory_layout = allocate(graph)
constants_map = {}
for constant in traverse.filter_nodes(traverse.listup_nodes(graph), ConstantVariable): # type: ConstantVariable
constants_map[constant.name] = {
"byte_offset": memory_layout[constant].offset * 4,
"size": constant.size
}
constant_encoder = ConstantEncoder.get_encoder(kwargs.get("constant_encoder_name", None))
constants_bytes = constant_encoder.encode(memory_layout)
kernels = cls.generate_kernels(graph)
descriptor = GraphDescriptor(
kernels=kernels,
memory_layout=memory_layout,
inputs=graph.inputs,
outputs=graph.outputs,
constants_encoding=constant_encoder.name,
constants_map=constants_map,
licenses=graph.licenses
)
return GraphExecutionData(graph, descriptor, constants_bytes)
def allocate(graph: Graph) -> WebGLMemoryLayout:
nodes = traverse.listup_nodes(graph)
operators = traverse.filter_nodes(nodes, Operator) # type: List[Operator]
variables = traverse.filter_nodes(nodes, Variable) # type: List[Variable]
for i, v in enumerate(variables):
if v.name is None:
v.name = _name("v")
dynamic_constants = traverse.filter_nodes([v for v in variables if not Placeholder.check_resolved(v.size)], ConstantVariable)
assert len(dynamic_constants) == 0, f"ConstantVariable with unresolved placeholder shape is detected: f{dynamic_constants}"
allocations = _get_allocations(graph, operators, variables)
_optimize_buffer_reuse(allocations)
variable_allocations = {v: allocations[v] for v in variables if not isinstance(v, ConstantVariable)}
constant_allocations = {v: allocations[v] for v in variables if isinstance(v, ConstantVariable)}
data = _update_constant_offset(constant_allocations)
allocations = variable_allocations
allocations.update(constant_allocations)
layout = WebGLMemoryLayout(allocations, data)
return layout
def generate(cls, graph: Graph, **kwargs):
data_dict = {} # type: Dict[int, Tuple[GraphDescriptor, bytes]]
for max_texture_size in [4096, 8192, 16384]:
config.WEBGL_MAX_TEXTURE_SIZE = max_texture_size
graph, _ = WebGLOptimizeRule().optimize(graph)
memory_layout = allocate(graph)
constants_map = {}
for constant in traverse.filter_nodes(traverse.listup_nodes(graph), ConstantVariable): # type: ConstantVariable
constants_map[constant.name] = {
"byte_offset": memory_layout[constant].offset * 4,
"size": constant.size
}
constant_encoder = ConstantEncoder.get_encoder(kwargs.get("constant_encoder_name", None))
constants_bytes = constant_encoder.encode(memory_layout)
kernels = cls.generate_kernels(graph)
descriptor = GraphDescriptor(
kernels=kernels,
memory_layout=memory_layout,
inputs=graph.inputs,
outputs=graph.outputs,
constants_encoding=constant_encoder.name,
constants_map=constants_map,
licenses=graph.licenses
)
data_dict[max_texture_size] = (descriptor, constants_bytes)
return GraphExecutionData(graph, data_dict)
def allocate(graph: Graph) -> MemoryLayout:
nodes = traverse.listup_nodes(graph)
operators = traverse.filter_nodes(nodes, Operator) # type: List[Operator]
variables = traverse.filter_nodes(nodes, Variable) # type: List[Variable]
for i, v in enumerate(variables):
if v.name is None:
v.name = _name("v")
dynamic_constants = traverse.filter_nodes([v for v in variables if not Placeholder.check_resolved(v.size)], ConstantVariable)
assert len(dynamic_constants) == 0, f"ConstantVariable with unresolved placeholder shape is detected: f{dynamic_constants}"
allocations = _get_allocations(graph, operators, variables)
_optimize_inplace(operators, allocations)
variable_allocations = {v: allocations[v] for v in variables if not isinstance(v, ConstantVariable)}
constant_allocations = {v: allocations[v] for v in variables if isinstance(v, ConstantVariable)}
_update_offset(variable_allocations)
_optimize_buffer_reuse(variable_allocations)
data = _update_constant_offset(constant_allocations)
for allocation in set(variable_allocations.values()):
allocation.offset += data.size
allocations = variable_allocations
allocations.update(constant_allocations)
layout = MemoryLayout(allocations, data)
if flags.VISUALIZE_MEMORY_ALLOCATION:
_visualize_allocation(operators, variables, layout)
return layout
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
global _rgba_support_operators
flag_changed = False
for node in traverse.listup_nodes(graph):
if node.has_attribute(ChannelMode):
continue
if isinstance(node, ConvertRtoRGBA) or isinstance(node, ConvertRGBAtoR):
continue
flag_changed = True
node.attributes.add(ChannelMode(node, ChannelModeEnum.R))
if isinstance(node, Operator):
node.attributes.add(SupportedChannelMode(node, ChannelModeEnum.R))
if node.__class__ not in _rgba_support_operators:
continue
variables = list(node.inputs.values()) + list(node.outputs.values())
if not all(v.order == variables[0].order for v in variables):
continue
if not all(v.shape == variables[0].shape for v in variables):
continue
node.attributes.add(SupportedChannelMode(node, ChannelModeEnum.RGBA))
return graph, flag_changed
def optimize(self, graph: Graph):
flag_changed = False
for v in traverse.filter_nodes(traverse.listup_nodes(graph), SplitTarget): # type: Variable
axis = _choose_split_axis(v)
_split_axis(v, axis, graph)
flag_changed = True
return graph, flag_changed
def test_listup_nodes_hidden_output():
v0 = Variable((1, 1), OrderNC)
op1 = Operator("op1")
v1 = Variable((1, 2), OrderNC)
op2 = TestOperator("op2")
v2 = Variable((1, 3), OrderNC)
op1.append_input("v0", v0)
op1.append_output("v1", v1)
op2.append_input("v1", v1)
op2.append_output("v2", v2)
graph = Graph([v0], [v1, v2]) # outputs hidden variable
nodes = listup_nodes(graph)
assert tuple(nodes) == (v0, op1, v1, op2, v2), str(nodes)
def optimize(self, graph: Graph) -> Tuple[Graph, bool]:
if not (flags.optimize.OPTIMIZE
and flags.optimize.CONCAT_SCALAR_OPERATION):
return graph, False
flag_changed = False
nodes = traverse.listup_nodes(graph)
filtered_nodes = traverse.filter_nodes(
nodes, ScalarAffine) # type: List[ScalarAffine]
while len(filtered_nodes) > 0:
op = filtered_nodes.pop()
if op.scale == 1 and op.bias == 0:
remove_operator(op)
flag_changed = True
filtered_nodes = traverse.filter_nodes(
nodes, ScalarAdd) # type: List[ScalarAdd]
while len(filtered_nodes) > 0:
op = filtered_nodes.pop()
if op.value == 0:
remove_operator(op)
flag_changed = True
filtered_nodes = traverse.filter_nodes(
nodes, ScalarMul) # type: List[ScalarMul]
while len(filtered_nodes) > 0:
op = filtered_nodes.pop()
if op.value == 1:
remove_operator(op)
flag_changed = True
filtered_nodes = traverse.filter_nodes(
nodes, ScalarPow) # type: List[ScalarPow]
while len(filtered_nodes) > 0:
op = filtered_nodes.pop()
if op.value == 1:
remove_operator(op)
flag_changed = True
return graph, flag_changed
def test_listup_nodes_residual():
global graph, op1, op2, op3, v0, v1, v2, v3
nodes = listup_nodes(graph)
assert tuple(nodes) == (v0, op1, v1, op2, v2, op3, v3)
