Visualising Pallas Load/Stores (even with BlockSpecs)

[oliverdutton]

Pallas_Visualisation

A tool to visualise the load/store ops in Pallas functions. Similar to triton-viz which the visualisation

You can use any pallas machinery for writing your kernel, e.g. BlockSpecs or the grid, and visualise it.

You can view indexing for a particular index on the grid or an overview of the whole grid, such as this add op split over four blocks.

The tracing is done by adding MemRefs which track the indexing, on which all load/store ops are duplicated to operate on and then logged in an external callback by editing the jaxpr.

As a followup

Would anyone be interested in:

1. a Pallas version of Triton-Puzzles but with Pallas explicitly in mind from the start.
   or
2. BlockSpecs that support masking (so a boundary mask is added to any loads/stores and the 'other' kwarg is set on a per-MemRef rather than per-load basis) and implicit tiling (for the subset of cases where we are just working with blocks which partition the arrays, to reduce the amount of code required)