NVlabs / timeloop

Timeloop performs modeling, mapping and code-generation for tensor algebra workloads on various accelerator architectures.
https://timeloop.csail.mit.edu/
BSD 3-Clause "New" or "Revised" License
303 stars 99 forks source link

Systolic array dataflow representation #228

Open joon9885 opened 8 months ago

joon9885 commented 8 months ago

Hello.

I am trying to run the energy calculations of the conventional output stationary systolic array where the input activations and weights are skewed to exploit the data reuse by moving the data to the neighbor (try to use the link-transfer of Timeloop).

The input activations and weights are unrolled and skewed and then sent to a vector and continuously fed to the systolic array.

I saw that the Timeloop support skew feature, but not sure that is same as the skew I am talking about.

Attached is the dataflow that I want to simulate (which is the figure from here).

스크린샷 2023-10-22 오후 10 04 36

So, is it possible to simulate the skewed dataflow? Then, does the data movement between Activation memory and Systolic array be modeled as multicast? I believe there should be difference between data movement from the activation memory to the PEs at the edge, and between PEs.

Thank you.

angshuman-parashar commented 8 months ago

Right, so skew can transform a multicast into link transfers (or vice-versa). Timeloop supports spatial skew, which means the space_id at each level can be shuffled around using a linear function of the current level's spatial and temporal loop indices.

Some dataflows can be described simply using a spatial skew. I am unable to review your specific dataflow, but if you can describe the behavior by setting the space_id of each PE as described above, I can help you with the syntax to express this in the YAML.

However, not all dataflows can be described using spatial skew, and need the time-coordinate to be shifted/shuffled as well. At present the skew interface does not support this. That said, I believe Timeloop's new internal polyhedral analysis core can in principle model temporal skew. @gilbertmike could you please comment on whether (a) this is correct and (b) it's possible to tap into this in any way?

If this fails and your dataflow does need temporal skew, you can model an unskewed/spatial-only-skewed version of the dataflow, but change the multicast hop-count calculations to model peer-to-peer transfers instead of a parent-to-child multicast. We can provide further guidance with this approach.

joon9885 commented 8 months ago

Thank you for the comments.

The dataflow represented in the model indeed need the temporal skew, so I think the method in the last paragraph would be plausible.

The dataflow of the input is that, for each input feature map section selected by the weight kernel size, the section is unrolled to 1D and sent to the PEs. The corresponding weights are also unrolled so that for each cycle, the MAC operation for convolution is conducted. For spatial-X, the number of weights (M) is divided to each column of PE array, and for spatial-Y, each combination of output feature element (P and Q) is divided to each row of PE array. For example, if we have 3X3 PE, then PE at (1, 1) is for (P1Q1M1), PE at (1, 2) is for (P1Q1M2), PE at (2, 1) is for (P2Q1M1) and so on. For better visibility, I updated the dataflow picture above.

EMJzero commented 1 month ago

Hi, I apologize for writing in this old issue, but I would be interested in knowing the way spatial skew can be modeled in Timeloop. I already found the parser for "skew" constraints here: https://github.com/NVlabs/timeloop/blob/c116d0b35b91d99001021bccb02e935a184bbda1/src/mapping/parser.cpp#L725 However, I was unable to find an explanation for any of the options there present. Henceforth, could I ask you for guidance in understanding such feature?

Unlike @joon9885, I am interested only in spatial skew, as my goal is to model a systolic array with an output-stationary dataflow, hence I would like to model the spatial skew of its input matrices within Timeloop. My apologies in case I missed out on any related documentation, and thank you in advance for any help provided.

PS: I also found that the current Timeloop-v4 python frontend does not seem to recognize the "skew" constraint type. Is there another way to handle skew since v4?

angshuman-parashar commented 1 month ago

The new ISL-based nest analysis should be able to handle spatial skew, but we don't have the parser to read a spatial skew descriptor from the user and translate it into the equivalent ISL mapping relations. @gilbertmike may be able to elaborate.

Paging @tanner-andrulis for the v4 incompatibility issue with skew.

tanner-andrulis commented 1 month ago

Is there a skew spec available anywhere? If so, I can integrate it into the v4 parser in the next day or two