Document for different specializations and how they affect shapes

[facebookincubator]

Document for different specializations and how they affect shapes

[antinucleon]

the explicitly reshape introduced for this reason:

In low-level math component such as cuBLAS/CUTLASS etc, a gemm is a strictly 2D problem, eg RCR variance:

Y: [M, N] -> gemm_rcr(X: [M, K], W: [N, K])

In pytorch or other framework, there is sugar for ND problem, eg

Y: [B, S, 4H] -> torch.functional.linear(X: [B, S, H], W: [4H, H])
# in low level, X is reshaped into [B * S, H], output Y is initially [B * S, 4H], then reshaped into [B, S, 4H] 

To lower this syntax sugar to actual low level implementation, we insert reshape in AIT

In AIT all reshape ops are 0 cost.

You can try the visualize tool as well: https://facebookincubator.github.io/AITemplate/tutorial/how_to_visualize.html

[dashesy]

Thanks! That is what I was looking for.