Systolic Array Extensions (for November Deadline)

[calebmkim]

(These ideas are not necessarily mutually exclusive, just a list of things we could discuss)

Ideas

• Flash Attention
• Softmax calculation that uses a (faster) online calculation for the normalizers (i.e., for each row, calculates both: max value in the row, and $\sum_{i}{e^{i-\text{max of row}}}$ in the same iteration).
• T2S. Halide-inspired programming methodology for spatial architectures. Do we want to a) actually implement a subset of the language features (e.g., actually implement A.tile(i, ii, II) and A.unroll(i, …) commands into Calyx designs)? Or do we want to try to look at what their case studies are doing and see if we can generate a Calyx design that generates the same thing?
• Google TPU has a Programmable Vector Unit instead of a collection of pooling units and hard-coded activation functions. We could try to build a smaller version of this (i.e., a component that can do different post ops based on its input)
• PolySA presents different systolic architectures for MM.
• (Aspirational) Sparse Linear Algebra Integration

Implementation Details

• For anything w/ attention/softmax, we should think about improving the performance of our exponentiation component

[calebmkim]

(Mainly writing this down for my future reference so I don't forget it).

How to make Systolic Array and Post-Op Separate Component

Having the following component signatures: systolic_array(...)(... r1_valid, r1_value, r1_idx, r2_valid, r2_value, r2_idx, etc.) post_op(r1_valid, r1_value, r1_idx, r2_valid, r2_value, r2_idx, etc.) And the following group to control the coordination (the actual group may look slightly different):

// group will be very ugly looking 
group perform_op {
  // something 
  systolic_array.go = !systolic.array.done ? 1'd1;
  // other systolic array input ports 
  post_op.r1_valid = systolic_array.r1_valid; 
  post_op.r1_value = systolic_array.r1_value; 
  // etc. 
  perform_op[done] = post_op.done 
} 

We should not do the following:

par {invoke systolic_array; invoke post_op;}

since there is no guarantee about the relative ordering of threads of par blocks.