Update: I'm actively working on this. What I've found out so far is that activating causal masking affects the propagation of layout information in the kernel function. There are now two dependent loops, and the current block's indices are evaluated in the computation, so we need to support additional operations (tt.make_range, arith.select and arith.cmpi) related to the mask computation.

Update: Still WIP. Locally, I have extended the perf-attn branch enough to reach LLVM-IR generation, though there are a couple of illegal instruction being created, which I believe originate from the TTGIR-level.

Update: Still WIP. I'm investigating how to handle "row vectors" (e.g. tensor<1x64>) correctly.

WIP.

Update: PR is up, still working on tests.

Update: The linked PRs are under review.

Update (for 9/9/24): Two outstanding PRs to land, then I need to check if E2E execution works on llvm-target.

Update (for 9/9/24): Two outstanding PRs to land, then I need to check if E2E execution works on llvm-target.

FYI, flash attention is not yet functional on llvm-target, you can apply #2061 when testing.

Update: I'm investigating an outstanding issue with the layout propagation at the beginning of the advanced path.

Update: Unchanged, still investigating

Update: Still working on this.