Add rock.blockwise_load_tile to encapsulate the logic of loads

[dhernandez0]

Motivation

The purpose of this PR is to encapsulate the logic of loads from memory (+ store to LDS and load from LDS). The goal is to be able to use the same logic for attention.

Technical Details

• make LDS inputs optional for BlockwiseGemmAccelOp
• GemmLoadTileType to describe loading behaviors (Default, DoubleBuffer, BypassLDS)
• new BlockwiseLoadTileOp to encapsulate the loading logic
• BlockwiseLoadTileToThreadwisePass pass to lower BlockwiseLoadTileOp to the desired loading behavior (based on GemmLoadTileType)
• ToBlockwise: refactor gemm and attention to use common functions and use BlockwiseLoadTileOp to encapsulate loading for gemm.

Test Plan

• Add tests for new pass
• Add test for ToBlockwise (as it has changed significantly)
• Compare performance with develop branch. To verify no slowdown is introduced by this PR.

Test Result

Tests pass

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[justinrosner]

Do we want for this to also implement MemoryEffectsOpInterface?

[dhernandez0]

that's a good point, GemmOp and ThreadwiseAccelGemmOp do implement it, not sure why this one doesn't. I'll do that.

[dhernandez0]

and BlockwiseGemmOp implements it as well. This looks like a bug

[dhernandez0]

done

[justinrosner]

Can we have a more detailed description of what this pass does here?

[dhernandez0]

done

[pabloantoniom]

Code LGTM! Did you check that there are is no performance variation?

[dhernandez0]

Code LGTM! Did you check that there are is no performance variation?

Yes, I just finished doing that. Most of the tier1 gemm list problem configs generate the exact same assembly. For some problem configs (~18%), there are some with minor assembly changes but no performance difference.