def test_small_graph():
subgraph = cs.TESubgraph([], None)
part_a = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
cs.Propagator(
[[0, 1, 0], [1, 0, 0], [0, 0, 1]],
[-1, -1],
),
],
)
part_b = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([10, 10], "uint8")
tensor_2 = cs.Tensor([9, 9], "uint8")
tensor_3 = cs.Tensor([10, 10], "uint8")
tensor_4 = cs.Tensor([10, 10], "uint8")
part_a.set_input(0, tensor_1)
part_a.set_input(1, tensor_2)
part_a.set_output(tensor_3)
tensor_1.add_consumer(part_a)
tensor_2.add_consumer(part_a)
tensor_3.add_producer(part_a)
part_b.set_input(0, tensor_3)
part_b.set_output(tensor_4)
tensor_3.add_consumer(part_b)
tensor_4.add_producer(part_b)
assert part_a.input_tensors == [tensor_1, tensor_2]
assert part_a.output_tensor == tensor_3
assert part_b.input_tensors == [tensor_3]
assert part_b.output_tensor == tensor_4
assert tensor_1.producers == []
assert tensor_1.consumers == [part_a]
assert tensor_2.producers == []
assert tensor_2.consumers == [part_a]
assert tensor_3.producers == [part_a]
assert tensor_3.consumers == [part_b]
assert tensor_4.producers == [part_b]
assert tensor_4.consumers == []
graph = cs.CascaderGraph([tensor_1, tensor_2], [tensor_4])
assert graph.input_tensors == [tensor_1, tensor_2]
assert graph.output_tensors == [tensor_4]
assert graph.part_order == [part_b, part_a]
for i, part in enumerate(graph.part_order):
assert graph.get_part_id(part) == i
def test_generate_graph_plans(SRAM, DRAM):
num_part_groups = 3
stripe_factors = 4
max_plan_size = 10
subgraph = cs.TESubgraph([], None)
part_a = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
cs.Propagator(
[[0, 1, 0], [1, 0, 0], [0, 0, 1]],
[-1, -1],
),
],
)
part_b = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([10, 10], "int8")
tensor_2 = cs.Tensor([9, 9], "int8")
tensor_3 = cs.Tensor([10, 10], "int8")
tensor_4 = cs.Tensor([10, 10], "int8")
part_a.set_input(0, tensor_1)
part_a.set_input(1, tensor_2)
part_a.set_output(tensor_3)
tensor_1.add_consumer(part_a)
tensor_2.add_consumer(part_a)
tensor_3.add_producer(part_a)
part_b.set_input(0, tensor_3)
part_b.set_output(tensor_4)
tensor_3.add_consumer(part_b)
tensor_4.add_producer(part_b)
graph = cs.CascaderGraph([tensor_1, tensor_2], [tensor_4])
home_map = {
tensor_1: [SRAM, DRAM],
tensor_2: [SRAM],
tensor_3: [SRAM],
tensor_4: [SRAM, DRAM],
}
options = make_options(
cascade_region=SRAM,
stripe_factors=stripe_factors,
max_plan_size=max_plan_size,
)
closed_plans = _generate_graph_plans(graph, home_map, options)
assert len(closed_plans) == num_part_groups
def test_generate_output_stripe_configs_disable_striping(stripe_factors):
subgraph = cs.TESubgraph([], None)
part_1 = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[2, 0, 0], [0, 2, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([800, 800], "uint8")
tensor_2 = cs.Tensor([400, 400], "uint8")
part_1.set_input(0, tensor_1)
part_1.set_output(tensor_2)
tensor_1.add_consumer(part_1)
tensor_2.add_producer(part_1)
assert (
len(
_generate_output_stripe_configs(
part_1, stripe_factors, enable_striping=False, multi_dimensional=False
)
)
== 1
)
def test_generate_output_stripe_configs_single_dimension():
stripe_factors = 3
subgraph = cs.TESubgraph([], None)
part_1 = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[2, 0, 0], [0, 2, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([800, 800], "uint8")
tensor_2 = cs.Tensor([400, 400], "uint8")
part_1.set_input(0, tensor_1)
part_1.set_output(tensor_2)
tensor_1.add_consumer(part_1)
tensor_2.add_producer(part_1)
expected_stripe_configs = {
cs.StripeConfig([400, 1], [400, 400], [400, 1], [2, 1], [1, 400], [0, 0]),
cs.StripeConfig([400, 200], [400, 400], [400, 200], [2, 1], [1, 2], [0, 0]),
cs.StripeConfig([1, 400], [400, 400], [1, 400], [1, 2], [400, 1], [0, 0]),
cs.StripeConfig([200, 400], [400, 400], [200, 400], [1, 2], [2, 1], [0, 0]),
cs.StripeConfig([400, 400], [400, 400], [400, 400], [1, 2], [1, 1], [0, 0]),
}
output_stripe_configs = _generate_output_stripe_configs(
part=part_1, stripe_factors=stripe_factors, enable_striping=True, multi_dimensional=False
)
assert len(output_stripe_configs) == len(expected_stripe_configs)
assert set(output_stripe_configs) == expected_stripe_configs
def test_inline_part():
subgraph = cs.TESubgraph([], None)
part = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[0, 1, 0], [1, 0, 0], [0, 0, 1]],
[0, 0],
),
],
)
output_stripe_config = cs.StripeConfig([2, 4], [8, 8], [2, 4], [1, 2],
[4, 2], [0, 0])
input_stripe_config = cs.StripeConfig([4, 2], [8, 8], [4, 2], [2, 1],
[2, 4], [0, 0])
assert part.input_tensors == [None]
assert part.output_tensor == None
assert len(part.propagators) == 1
assert part.in_line == True
assert part.get_stripe_align_hint() == [1, 1]
performance_info = part.get_performance_info(output_stripe_config,
is_rolling=False)
assert performance_info.compute_cycles == 0
assert performance_info.read_bytes == [0]
assert performance_info.write_bytes == 0
input_stripe_configs = part.calculate_input_stripe_configs(
output_stripe_config)
assert len(input_stripe_configs) == 1
assert input_stripe_configs[0] == input_stripe_config
def test_generate_single_plans(SRAM, DRAM):
subgraph = cs.TESubgraph([], None)
part_1 = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[2, 0, 0], [0, 2, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([800, 800], "int8")
tensor_2 = cs.Tensor([400, 400], "int8")
part_1.set_input(0, tensor_1)
part_1.set_output(tensor_2)
tensor_1.add_consumer(part_1)
tensor_2.add_producer(part_1)
home_map = {
tensor_1: [SRAM, DRAM],
tensor_2: [SRAM],
}
options = make_options(cascade_region=SRAM, stripe_factors=1)
output_stripe_configs = _generate_output_stripe_configs(
part_1, options.stripe_factors)
plans = _generate_single_plans(part_1, output_stripe_configs, home_map,
options)
for plan in plans:
assert plan.interior_region == SRAM
assert plan.part_group == frozenset([part_1])
assert set(plan.tensor_configs.keys()) == set([tensor_1, tensor_2])
for open_config in plan.open_configs:
assert open_config.state == cs.TensorConfigState.INTERIOR
def test_generate_output_stripe_configs():
stripe_factors = 3
expected_configs = 13
subgraph = cs.TESubgraph([], None)
part_1 = cs.InlinePart(
subgraph,
[
cs.Propagator(
[[2, 0, 0], [0, 2, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = cs.Tensor([800, 800], "uint8")
tensor_2 = cs.Tensor([400, 400], "uint8")
part_1.set_input(0, tensor_1)
part_1.set_output(tensor_2)
tensor_1.add_consumer(part_1)
tensor_2.add_producer(part_1)
assert len(_generate_output_stripe_configs(
part_1, stripe_factors)) == expected_configs