def test_small_graph():
subgraph = TESubgraph([], None)
part_a = InlinePart(
subgraph,
[
Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
Propagator(
[[0, 1, 0], [1, 0, 0], [0, 0, 1]],
[-1, -1],
),
],
)
part_b = InlinePart(
subgraph,
[
Propagator(
[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
[0, 0],
),
],
)
tensor_1 = Tensor([10, 10], "uint8")
tensor_2 = Tensor([9, 9], "uint8")
tensor_3 = Tensor([10, 10], "uint8")
tensor_4 = 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 = 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_inline_part():
subgraph = TESubgraph([], None)
part = InlinePart(
subgraph,
[
Propagator(
[[0, 1, 0], [1, 0, 0], [0, 0, 1]],
[0, 0],
),
],
)
output_stripe_config = StripeConfig([2, 4], [8, 8], [2, 4], [1, 2], [4, 2],
[0, 0])
input_stripe_config = 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 match_ethosu_conv2d(output_tensor):
"""Match a Tensor Expression corresponding to an NPU Conv2D.
If the Tensor Expression matches, an EthosuPart will be created that models the
matched Tensor Expression. Otherwise, None will be returned.
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
Returns
-------
Union[None, EthosuPart]
The created EthosuPart if there was a match, otherwise None.
"""
write = output_tensor
if write.op.name != "ethosu_write":
return None
convert_to_nhcwb16 = write.op.input_tensors[0]
if convert_to_nhcwb16.op.name != "ethosu_convert_to_nhcwb16":
return None
conv2d = convert_to_nhcwb16.op.input_tensors[0]
if conv2d.op.name != "ethosu_conv2d":
return None
pad = conv2d.op.input_tensors[0]
if pad.op.name != "ethosu_pad":
return None
convert_to_nhwc = pad.op.input_tensors[0]
if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
return None
read = convert_to_nhwc.op.input_tensors[0]
if read.op.name != "ethosu_read":
return None
input_tensors = [
read.op.input_tensors[0],
conv2d.op.input_tensors[1],
conv2d.op.input_tensors[2],
]
subgraph = TESubgraph(input_tensors, output_tensor)
propagators = [
write.op.attrs["ifm_propagator"],
write.op.attrs["weights_propagator"],
write.op.attrs["bias_propagator"],
]
# TODO(@jacobbohlin) Both the output_quantum and quantum_cycles here are placeholders,
# needs true implementation.
if convert_to_nhcwb16.op.attrs["layout"] == "NHWC":
output_quantum = [1, 2, 2, 1]
else:
output_quantum = [1, 2, 1, 2, 1]
quantum_cycles = 1000
return EthosuPart(subgraph, propagators, output_quantum, quantum_cycles)
def match_ethosu_inline(output_tensor, device_config):
"""Match a Tensor Expression corresponding to an operator that will be inlined.
If the Tensor Expression matches, an InlinePart will be created that models the
matched Tensor Expression. Otherwise, None will be returned. This matcher is
naive and assumes nothing about the compute of the Tensor Expression. Therefore,
the resulting InlinePart will have full-tensor dependencies (i.e. each output
element depends on every input element).
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
device_config : EthosuDeviceConfig
Target device configuration
Returns
-------
Union[None, InlinePart]
The created InlinePart if there was a match, otherwise None.
"""
if output_tensor.op.name not in INLINE_OPS:
return None
input_tensors = output_tensor.op.input_tensors
propagators = []
output_dims = len(output_tensor.shape)
for input_tensor in input_tensors:
input_dims = len(input_tensor.shape)
transform_matrix = np.zeros((input_dims + 1, output_dims + 1))
for i, axis in enumerate(input_tensor.shape):
transform_matrix[i, output_dims] = int(axis)
transform_matrix[input_dims, output_dims] = 1
offset_vector = np.zeros(input_dims, dtype="int64")
propagators.append(
Propagator(
transform_matrix.tolist(),
offset_vector.tolist(),
))
subgraph = TESubgraph(input_tensors, output_tensor)
return InlinePart(
subgraph,
propagators,
)
def match_ethosu_unary_elementwise(output_tensor, device_config):
"""Match a Tensor Expression corresponding to an NPU Unary Elementwise.
If the Tensor Expression matches, an EthosuPart will be created that models the
matched Tensor Expression. Otherwise, None will be returned.
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
device_config : EthosuDeviceConfig
Target device configuration
Returns
-------
Union[None, EthosuPart]
The created EthosuPart if there was a match, otherwise None.
"""
write = output_tensor
if write.op.name != "ethosu_write":
return None
convert_to_nhcwb16 = write.op.input_tensors[0]
if convert_to_nhcwb16.op.name != "ethosu_convert_to_nhcwb16":
return None
unary_elementwise = convert_to_nhcwb16.op.input_tensors[0]
if unary_elementwise.op.name != "ethosu_unary_elementwise":
return None
pad = unary_elementwise.op.input_tensors[0]
if pad.op.name != "ethosu_pad":
return None
convert_to_nhwc = pad.op.input_tensors[0]
if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
return None
read = convert_to_nhwc.op.input_tensors[0]
if read.op.name != "ethosu_read":
return None
input_tensors = [
read.op.input_tensors[0],
]
subgraph = TESubgraph(input_tensors, output_tensor)
propagators = [
write.op.attrs["ifm_propagator"],
]
ifm_dtype = input_tensors[0].dtype
ofm_dtype = output_tensor.dtype
output_layout = convert_to_nhcwb16.op.attrs["layout"]
input_layout = convert_to_nhwc.op.attrs["layout"]
output_quantum = device_config.get_output_quantum(output_layout)
block_config = device_config.get_elementwise_block_config(
propagators[0],
None,
unary_elementwise.op.attrs,
output_tensor.shape,
output_layout,
input_layout,
None,
ifm_dtype,
ofm_dtype,
)
return EthosuPart(
subgraph,
propagators,
output_quantum,
1,
block_config,
)
def match_ethosu_conv2d(output_tensor, device_config):
"""Match a Tensor Expression corresponding to an NPU Conv2D.
If the Tensor Expression matches, an EthosuPart will be created that models the
matched Tensor Expression. Otherwise, None will be returned.
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
device_config : EthosuDeviceConfig
Target device configuration
Returns
-------
Union[None, EthosuPart]
The created EthosuPart if there was a match, otherwise None.
"""
write = output_tensor
if write.op.name != "ethosu_write":
return None
convert_to_nhcwb16 = write.op.input_tensors[0]
if convert_to_nhcwb16.op.name != "ethosu_convert_to_nhcwb16":
return None
conv2d = convert_to_nhcwb16.op.input_tensors[0]
if conv2d.op.name != "ethosu_conv2d":
return None
pad = conv2d.op.input_tensors[0]
if pad.op.name != "ethosu_pad":
return None
upscale = pad.op.input_tensors[0]
if upscale.op.name != "ethosu_upscale":
return None
convert_to_nhwc = upscale.op.input_tensors[0]
if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
return None
read = convert_to_nhwc.op.input_tensors[0]
if read.op.name != "ethosu_read":
return None
input_tensors = [
read.op.input_tensors[0],
conv2d.op.input_tensors[1],
conv2d.op.input_tensors[2],
]
subgraph = TESubgraph(input_tensors, output_tensor)
propagators = [
write.op.attrs["ifm_propagator"],
write.op.attrs["weights_propagator"],
write.op.attrs["bias_propagator"],
]
ifm_dtype = input_tensors[0].dtype
ofm_dtype = output_tensor.dtype
ifm_channels = int(input_tensors[0].shape[3])
ofm_channels, kernel_height, kernel_width = (
int(axis) for axis in input_tensors[1].shape[0:3])
kernel_elements = kernel_height * kernel_width
is_part_kernel = device_config.is_partkernel(conv2d.op.name, ifm_channels,
ifm_dtype, kernel_elements)
subkernels = len(
device_config.get_kernel_steps(conv2d.op.name, kernel_height,
kernel_width, ifm_dtype,
is_part_kernel))
output_layout = convert_to_nhcwb16.op.attrs["layout"]
input_layout = convert_to_nhwc.op.attrs["layout"]
output_quantum = device_config.get_output_quantum(output_layout)
valid_block_configs = device_config.get_valid_block_configs(
propagators[0],
conv2d.op.attrs,
output_tensor.shape,
ofm_channels,
ifm_channels,
output_layout,
input_layout,
ifm_dtype,
ofm_dtype,
kernel_height,
kernel_width,
)
return EthosuPart(
subgraph,
propagators,
output_quantum,
subkernels,
valid_block_configs,
1,
)
def match_ethosu_pooling(output_tensor, device_config):
"""Match a Tensor Expression corresponding to an NPU Pooling.
If the Tensor Expression matches, an EthosuPart will be created that models the
matched Tensor Expression. Otherwise, None will be returned.
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
device_config : EthosuDeviceConfig
Target device configuration
Returns
-------
Union[None, EthosuPart]
The created EthosuPart if there was a match, otherwise None.
"""
write = output_tensor
if write.op.name != "ethosu_write":
return None
convert_to_nhcwb16 = write.op.input_tensors[0]
if convert_to_nhcwb16.op.name != "ethosu_convert_to_nhcwb16":
return None
pool2d = convert_to_nhcwb16.op.input_tensors[0]
if pool2d.op.name != "ethosu_pooling":
return None
pad = pool2d.op.input_tensors[0]
if pad.op.name != "ethosu_pad":
return None
convert_to_nhwc = pad.op.input_tensors[0]
if convert_to_nhwc.op.name != "ethosu_convert_to_nhwc":
return None
read = convert_to_nhwc.op.input_tensors[0]
if read.op.name != "ethosu_read":
return None
input_tensors = [
read.op.input_tensors[0],
]
subgraph = TESubgraph(input_tensors, output_tensor)
propagators = [
write.op.attrs["ifm_propagator"],
]
ifm_dtype = input_tensors[0].dtype
ofm_dtype = output_tensor.dtype
ifm_channels = int(input_tensors[0].shape[3])
ofm_channels = ifm_channels
pool_shape_h = int(pool2d.op.attrs["pool_shape_h"])
pool_shape_w = int(pool2d.op.attrs["pool_shape_w"])
subkernels = len(
device_config.get_kernel_steps(pool2d.op.name, pool_shape_h, pool_shape_w, ifm_dtype)
)
output_layout = convert_to_nhcwb16.op.attrs["layout"]
input_layout = convert_to_nhwc.op.attrs["layout"]
output_quantum = device_config.get_output_quantum(output_layout)
valid_block_configs = device_config.get_valid_block_configs(
propagators[0],
pool2d.op.attrs,
output_tensor.shape,
ofm_channels,
ifm_channels,
output_layout,
input_layout,
ifm_dtype,
ofm_dtype,
pool_shape_h,
pool_shape_w,
)
return EthosuPart(
subgraph,
propagators,
output_quantum,
subkernels,
valid_block_configs,
)
def match_ethosu_identity(output_tensor, device_config):
"""Match a Tensor Expression corresponding to an NPU identity.
If the Tensor Expression matches, an EthosuPart will be created that models the
matched Tensor Expression. Otherwise, None will be returned.
Parameters
----------
output_tensor : tvm.te.Tensor
The tensor to attempt to match with.
device_config : EthosuDeviceConfig
Target device configuration
Returns
-------
Union[None, EthosuPart]
The created EthosuPart if there was a match, otherwise None.
"""
write = output_tensor
if write.op.name != "ethosu_write":
return None
identity = write.op.input_tensors[0]
if identity.op.name != "ethosu_identity":
return None
read = identity.op.input_tensors[0]
if read.op.name != "ethosu_read":
return None
input_tensors = [
read.op.input_tensors[0],
]
subgraph = TESubgraph(input_tensors, output_tensor)
propagators = [
write.op.attrs["ifm_propagator"],
]
ifm_dtype = input_tensors[0].dtype
ofm_dtype = output_tensor.dtype
input_tensors_shape = input_tensors[0].shape
length = len(input_tensors_shape)
assert length <= 4, "Input tensor shape must be <= 4 for the identity operator"
channels = int(input_tensors_shape[length - 1]) if length >= 3 else 1
subkernels = len(
device_config.get_kernel_steps(identity.op.name, 1, 1, ifm_dtype))
input_layout = output_layout = "NHWC"
output_quantum = device_config.get_output_quantum(output_layout)
valid_block_configs = device_config.get_valid_block_configs(
propagators[0],
identity.op.attrs,
output_tensor.shape,
channels,
channels,
output_layout,
input_layout,
ifm_dtype,
ofm_dtype,
1,
1,
)
return EthosuPart(
subgraph,
propagators,
output_quantum,
subkernels,
valid_block_configs,
)