pub fn recognize_head_matmul(stmts: &[Stmt<'_>]) -> Option<HeadMatmulSpec>Expand description
Recognize Let Y be X matmul W across H heads — the matmul’s nest with a head loop between the row loop and the
contraction: Y[t·(H·b) + h·b + r + 1] += X[t·(H·a) + h·a + e + 1] · W[(h·a + e)·b + r + 1], the loops t, h, e,
r over m, H, a, b.