Auto-generate spatial mapping?

[r09g]

How should the mapping files look (for example tpu_like.py) if one wanted to use the SpatialMappingGeneratorStage to automatically generate the spatial mappings? I tried removing the entire line on spatial mapping but received an error in the model parsing:

mapping = {
    "default": {
        "core_allocation": 1,
        "memory_operand_links": {'O': 'O', 'W': 'I2', 'I': 'I1'}
    },
}

runfile('/Users/rhyang/Documents/ai_codesign/wrapper.py')
Traceback (most recent call last):
  File "/Users/rhyang/Documents/ai_codesign/main_onnx.py", line 48, in <module>
    mainstage.run()
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/stages/Stage.py", line 53, in run
    for cme, extra_info in self.list_of_callables[0](self.list_of_callables[1:], **self.kwargs).run():
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/stages/ONNXModelParserStage.py", line 16, in run
    self.onnx_model_parser.run()
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/io/onnx/model.py", line 57, in run
    workload = self.parse_workload_from_onnx_model_and_mapping()
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/io/onnx/model.py", line 97, in parse_workload_from_onnx_model_and_mapping
    node_obj = parser.run()
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/io/onnx/conv.py", line 21, in run
    layer_node = self.generate_layer_node_for_conv()
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/io/onnx/conv.py", line 169, in generate_layer_node_for_conv
    node_attrs = get_layer_node_input_format(kernel_shape, strides, dilations, groups, padding,
  File "/Users/rhyang/Documents/ai_codesign/zigzag/classes/io/onnx/conv.py", line 119, in get_layer_node_input_format
    d["spatial_mapping"] =  node_mapping["spatial_mapping"]
KeyError: 'spatial_mapping'

Then I tried adding the spatial mapping line in tpu_like.py but just having an empty dict:

mapping = {
    "default": {
        "core_allocation": 1,
        "spatial_mapping": {},
        "memory_operand_links": {'O': 'O', 'W': 'I2', 'I': 'I1'}
    },
}

there was another error from the SpatialMappingConversionStage on no mappings provided by the user. I found that in line 68 and 70 of this conversion stage, the if conditions to check whether the user provided spatial mappings were to check if the "spatial mapping" is an instance of dict or list. However an empty dict still is an instance of a dict so this will not go into automatic spatial mapping generation.

Thanks.

[LY-Mei]

Hi, you can remove the "spatial_mapping" in the mapping file, and at the same time, replace the SpatialMappingConversionStage with SpatialMappingGeneratorStage in the main file. In this case, the framework will search for the best spatial mapping instead of using the user-defined one.