[DispatchCreation] Fuse reshape op chains along with set_encoding ops

[Max191]

There can be reshape ops in between set_encoding ops and their producer dispatch region ops when we fuse encoding ops into dispatches, so this PR pulls any reshape op chains into the producer dispatch along with the set_encoding op. This enables more producer fusions for set_encoding in the data tiling fusion path.

[jtuyls]

LGTM

[Max191]

I'm going to convert this to a draft now, because @jtuyls found some issues with llama.

[hanhanW]

Let me put a blocker, just in case.

[jtuyls]

Just a quick update on the Llama issue: it's caused by a dominance problem, but after fixing that, one of the new dispatches still fails somewhere in the LLVM backend. I think I have a fix/workaround for that as well, but it's in LLVM and I need to look at it a bit more to see whether it makes sense and create an upstream PR.

[Max191]

Just a quick update on the Llama issue: it's caused by a dominance problem, but after fixing that, one of the new dispatches still fails somewhere in the LLVM backend. I think I have a fix/workaround for that as well, but it's in LLVM and I need to look at it a bit more to see whether it makes sense and create an upstream PR.

Nice! What is the dispatch that's failing?

[jtuyls]

Nice! What is the dispatch that's failing?

Here it is: https://gist.github.com/jtuyls/5590340b86fedbbd83f310f45904eee4. It doesn't look too special to me but it fails with:

iree/third_party/llvm-project/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp:3628: const llvm::TargetRegisterClass* llvm::SIRegisterInfo::getEquivalentSGPRClass(const llvm::TargetRegisterClass*) const: Assertion `SRC && "Invalid register class size"' failed.

The repro command is:

iree-compile --compile-from=executable-sources --compile-to=hal module_prefill_bs4\$async_dispatch_3.mlir --iree-llvmgpu-test-combine-layout-transformation

[jtuyls]

I need two fixes to get this to work with Llama:

• [Flow] Add support for moving operations with dependencies into dispatch regions #21399
• [AMDGPU] Fix sgpr to vreg_1 copy analysis llvm/llvm-project#149181

So I think we could land this PR? @Max191 @hanhanW

[Max191]

So I think we could land this PR?

Sure, I think we can just land it, since this is on an experimental path still, and your PR is already up. I think it will be easier that way.

EDIT: Alternatively, I could cherry-pick your fix onto this PR and list you as a coauthor if we want to avoid landing this without the fix. I'll leave it up to you which you prefer.

[jtuyls]

So I think we could land this PR?

Sure, I think we can just land it, since this is on an experimental path still, and your PR is already up. I think it will be easier that way.

EDIT: Alternatively, I could cherry-pick your fix onto this PR and list you as a coauthor if we want to avoid landing this without the fix. I'll leave it up to you which you prefer.

I am ok landing this as is and then I can rebase my PR and add a test for that as well. Llama isn't working on main yet anyway, I need an integrate as well to get the fix from llvm/llvm-project#149181 and until we put in some e2e llama-like tests I think it will likely be broken a couple more times.

[Max191]

@hanhanW looks like we are ready to land this now then.