Tiled Matmul for perfect tiling (abstract parallel model) - #152
Conversation
Squashes marco/closure-name-context-dependencies into marco/parallel-impl.
387ea06 to
04155f9
Compare
mstn
changed the title
There was a problem hiding this comment.
Thanks for this! Some questions/comments below. Main one: do we have a benchmark for tiled matmul for issue #124 ?
| `malloc`, and every block receives its own shared-memory arena managed by the | ||
| GPU runtime. The requested total must still fit the device's per-block shared | ||
| memory limit; this hardware limit is currently checked by CUDA/HIP rather than | ||
| encoded in the Hex types. |
There was a problem hiding this comment.
What happens if we exceed the limit? Do we get a crash?
There was a problem hiding this comment.
it should follow the existing trapping policy, but checking...
| b-only-layout = gpu.shared(tile-count, empty) | ||
| shared-layout = gpu.shared(tile-count, b-only-layout) | ||
|
|
||
| grid = gpu.grid.2d( |
There was a problem hiding this comment.
are only 2d grids possible? Can we express batched matmuls?
There was a problem hiding this comment.
I'd add more layouts for grids when needed. I don't think this is relevant for batched matmuls or is it?
| (arr gpu.launch_params : 0 -> 1) | ||
| # Nominal GPU levels. Their dimensions let lowering distinguish kernel | ||
| # launch from block-local cooperative execution. | ||
| (arr gpu.grid.2d : 4 -> 1) |
There was a problem hiding this comment.
I thought these were previously definitions, I'm not clear on what is baked in to the implementation in terms of shape dimensionality
There was a problem hiding this comment.
No, they were not, at least in my implementation. The reason is that they are like a sort of enum/constant, it cannot be defined in terms of other parallel primitives, we need to know that it is a gpu grid, does it make sense?
| ```text | ||
| host<gpu.grid.2d, ...> -> catena_gpu_grid_host_t | ||
| worker<gpu.grid.2d, ...> -> catena_gpu_grid_worker_t | ||
| worker<gpu.block.2d, ...> -> catena_gpu_block_worker_t |
| (def program identity : | ||
| ([a.] [.a]) | ||
| -> | ||
| ([a.] [.a]) |
There was a problem hiding this comment.
Please change this so that it uses the more concise
identity : [a] -> [a]
the more examples of this pattern in the codebase, the more codex propagates it
| ({gpu.state [.b]} * [coargs.]) | ||
| # Build the type-level description of a two-dimensional GPU block and | ||
| # preserve its runtime dimensions for coordinate calculations. | ||
| (arr gpu.block.2d : {[ |
There was a problem hiding this comment.
Similar to above comment, does this mean we need a compiler primitive for every dimensionality of tiling?
There was a problem hiding this comment.
I am afraid so, but in theory we have 1d, 2d, and 3d. In practice 1d and 2d. A thing I was considering is to pass a shape, but I don't know if it brings benefits beyond more complexity.
| @@ -0,0 +1,445 @@ | |||
| # Perfectly tiled f32 matrix multiplication. | |||
There was a problem hiding this comment.
Do we have a benchmark for this vs. a reference gemm?
There was a problem hiding this comment.
Not in this PR, see offline discussion.
380f912 to
91d7258
Compare
2 participants
Summary
Tiled Matmul for perfect tiling using an abstract model for context-based parallelism.
Details
See doc for details.
Limitations
foldinstead ofreduce