Add support for linalg.conv1d

[compor]

This PR:

• Adds support for linalg.conv1d
• Adds a simple test case file for the above

[compor]

This is black formatting I'm afraid, but we might want to sync at some point.

[compor]

@jorendumoulin told me it's OK to elide this for now

[EmilySillars]

Zigzag throws an error when I try to feed it the conv_1_d workload:
python xdsl_opt_main.py tests/conv_1d_generic.mlir -p linalg-to-stream && python run_zigzag.py
produces

Traceback (most recent call last):
  File "/home/emily/linalg-to-stream/run_zigzag.py", line 14, in <module>
    answers = zigzag.api.get_hardware_performance_zigzag(workload, accelerator, mapping, "latency", "outputs/my-output.json","outputs/list_of_cmes.pickle")
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/api.py", line 78, in get_hardware_performance_zigzag
    answers = mainstage.run()
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/stages/Stage.py", line 49, in run
    for cme, extra_info in self.list_of_callables[0](
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/stages/MainInputParserStages.py", line 55, in run
    workload = parse_workload_from_path_or_from_module(self.workload, self.mapping)
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/stages/MainInputParserStages.py", line 41, in parse_workload_from_path_or_from_module
    workload = DNNWorkload(workload_copy, mapping)
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/workload/dnn_workload.py", line 32, in __init__
    layer_node = LayerNode(layer_id, layer)
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/workload/layer_node.py", line 118, in __init__
    pr_loop, pr_loop_list, pr_scaling_factors = self.build_pr_funcs()
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/workload/layer_node.py", line 158, in build_pr_funcs
    pr_loop, pr_loop_list, pr_scaling_factors = self.extract_pr_loop_info(
  File "/home/emily/linalg-to-stream/.env/lib/python3.10/site-packages/zigzag/classes/workload/layer_node.py", line 301, in extract_pr_loop_info
    len(scaling_factors) == 2
AssertionError: Please remove any constants in the equation relation iJ=d0+d1.

Not sure if the error means the workload is incorrect, or if the calling of the zigzag api is incorrect. I will keep investigating this, but until it's clear why run_zigzag is failing, I don't think we should merge this PR.

[compor]

OK, this happened before I had venv requirements for stream/zigzag so I couldn't run it.
I don't fully understand the error message, but I'll set it up so I can reproduce this.

Things to try:

• use iJ=1*d0 + 1*d1 (this is not it)
• try stream (not clear what needs to be fixed there according to the README)

[compor]

Update: I can reproduce the assertion

[compor]

It turns out that parsing the workload description is quite finicky since ZigZag cannot parse dimensions dN where N is a digit.
I manually update the generated workload.py to this:

workload = {
    0: {
        "operator_type": "default",
        "equation": "O[x] += I[im] * W[y]",
        "dimension_relations": ["im=x+y"],
        "loop_dim_size": {"X": 16, "Y": 3},
        "operand_precision": {"O": 32, "O_final": 32, "W": 32, "I": 32},
        "operand_source": {"W": [], "I": []},
        "constant_operands": ["I", "W"],
        "padding": {},
    }
}

and I got an outputs dir which seems to be what we want.

Can you please confirm with above worload file that the output makes some sense (my understanding is embryonic on that front ATM).

If so, I'll amend the PR to avoid using d0, d1, etc.

[compor]

Hoping to update this today :)

[EmilySillars]

It turns out that parsing the workload description is quite finicky since ZigZag cannot parse dimensions dN where N is a digit. I manually update the generated workload.py to this:

workload = {
    0: {
        "operator_type": "default",
        "equation": "O[x] += I[im] * W[y]",
        "dimension_relations": ["im=x+y"],
        "loop_dim_size": {"X": 16, "Y": 3},
        "operand_precision": {"O": 32, "O_final": 32, "W": 32, "I": 32},
        "operand_source": {"W": [], "I": []},
        "constant_operands": ["I", "W"],
        "padding": {},
    }
}

and I got an outputs dir which seems to be what we want.

Can you please confirm with above worload file that the output makes some sense (my understanding is embryonic on that front ATM).

If so, I'll amend the PR to avoid using d0, d1, etc.

Hello Chris,
Thank you for looking into this further!

• It's too bad that zigzag cannot parse a dimension variable that contains a digit suffix. I'm opening up an issue to on the ZigZag repo to change this.

1. In the mean time, yes, could you change the code to avoid using d0 and d1 etc? I think this will mean our linalg parser will no longer be able to handle an arbitrary number of loops (unless we start naming loop dimensions with very long strings). Maybe we can just support a maximum of 26 loops, where each loop dimension is a single letter? That seems like the simplest band-aid for right now.
2. Unfortunately, this also affects the loop_dim_size field of the workload object. Loop dimension iterators like 'd1' need to have a corresponding upper case key like 'D1' in the loop_dim_size dictionary. When d1 becomes y, the loop_dim_size corresponding key needs to be Y.
3. Your modified workload looks correct to me, but when the loop iterator names in the workload change, the spatial mapping definition must also change. So to make your modified workload fully work, you need to modify the line in inputs/mapping/snax_gemm.py from

"spatial_mapping": {"PE_ARRAY_D1": ("D0", 8), "PE_ARRAY_D2": ("D1", 8), "PE_ARRAY_D3": ("D2", 8)},

to

"spatial_mapping": {"PE_ARRAY_D1": ("X", 8), "PE_ARRAY_D2": ("Y", 8), "PE_ARRAY_D3": ("D2", 8)},

• To check the output generated by your workload made sense, I added the line print_mapping(cme) to xdsl_opt_main.py file to check that the printed output made sense. (print_mapping(cme) is how I noticed your workload was missing a spatial mapping output)

Let me know if these changes are too tedious, in which case you are welcome to offload them to me :)
Thank you very much and LMK what you think.

[compor]

@EmilySillars thanks for all this. A question regarding 3.: what does "fully work" mean and how was this identified?
Was it just by eyeballing?

[EmilySillars]

@EmilySillars thanks for all this. A question regarding 3.: what does "fully work" mean and how was this identified? Was it just by eyeballing?

"fully work" means the output from zigzag produces both a valid temporal and spatial mapping.

• Without the change I suggest, the spatial mapping is empty: SpatialMapping({'O': [[], [], [], []], 'W': [[], [], []], 'I': [[], [], []]}) which doesn't make sense, because the accelerator can always unroll part of a loop on its mac array.
• For the temporal mapping I kind of eye-balled correctness, but with a little more thinking:

=====================================================================
Temporal Loops                O            W            I            
=====================================================================
for X in [0, 2):              l1           l1           l1           
---------------------------------------------------------------------
  for X in [0, 2):            l1           l1           l1           
---------------------------------------------------------------------
    for X in [0, 2):          l1           l1           l1           
---------------------------------------------------------------------
      for X in [0, 2):        l1           l1           l1           
---------------------------------------------------------------------
        for Y in [0, 3):      rf_32b_O     l1           l1           
---------------------------------------------------------------------

To me this workload looks correct because there is a memory level (l1, rf_32b_O) assigned to each operand, and the loop dimensions in the temporal mapping multiplied together are equal to the original loop dimensions from the workload. For example, the workload input says "loop_dim_size": {"X": 16, "Y": 3}, and looking at zigzag's temporal loop output, 2 * 2 *2 *2 = 16, and 3 = 3.