Add propagation of thread dependent index sequences #292
Conversation
70e166d to
a3b6f19
Compare
8831e4d to
2ef7724
Compare
9a03b35 to
bb59e95
Compare
Signed-off-by: Harsh Menon <harsh@nod-labs.com>
| if dst_op: | ||
| for node in propagated_resolutions: | ||
| get_custom(node).index = dst_op.index |
There was a problem hiding this comment.
When would we use it without dst_op?
There was a problem hiding this comment.
Good question. This was when I was using it in 2 places where dst_op was optional. I can modify this so that it doesn't handle that case.
| thread_dependent_index: dict[IndexSymbol, IndexSequence], | ||
| ) -> dict[IndexSymbol, IndexSequence]: | ||
| combined_index = {k: v for k, v in thread_independent_index.items()} | ||
| for k in combined_index: |
There was a problem hiding this comment.
Would be nice to add some docs, to different examples/cases on when we'd add the start offsets. (perhaps read -> read or offseted_read + mma?)
There was a problem hiding this comment.
In this PR, we split the index assignment into 2 phases. The first is a thread-independent index assignment where we set the indices based on work group constraints and tiling constraints and anything that in general does not have any thread level dependence. Once we have this index then we go through and propagate the thread dependent index which comes either from MMA nodes or if there are no MMA nodes it comes from reads and writes and we then add the thread dependent index to the thread independent index. So we are always going to be adding these two offsets together.
| vector_shapes = ( | ||
| custom.vector_shapes if custom.vector_shapes else source_vector_shapes | ||
| ) | ||
| sources.append((custom, source_index, vector_shapes)) |
There was a problem hiding this comment.
why not use source.index? Seems like the only time sourc.index != source_index is when source.index has underlying indices before getting propagated. But wouldn't we want that?
There was a problem hiding this comment.
So source.index contains the unified index which is the sum of the thread dependent index and the thread independent index. During propagation we only want to propagate the thread dependent index and that's why we only propagate source_index.
This PR adds a major refactor of the index sequence analysis. Specifically,