[AMD] inThreadTranspose: Transpose between global load and local store for non-TN layouts: part 2 of 4

[jtang10]

inThreadTranpose: part 2 of 4

Introduction

This PR introduces the AMD-specific inThreadTranspose feature to improve shared memory access efficiency for non-TN GEMM and 2nd dotOp in Flash Attention.
The entire feature has been broken into 4 pieces for more reliable integration and this PR is the 2nd of 4.

• Add TTGIR pass for inThreadTranspose w/o activating it
• Add the sharedEncodingAttr update and associate toLLVM change. This should fully enable inThreadTranspose.
• Add the new flag in sharedEncodingAttr w/o using it
• Change the shared linearLayout conversion and shared->dot in amd path to enable it. This should improve the perf of inThreadTranspose.

Feature description

Currently on AMD hardware, if the dot Operand is K-major, we'd use the same vectorization for ds_write as global_load, but won't coalesce on ds_read, resulting in poor shared memory/LDS read efficiency prior to MFMA operation.
This feature, inThreadTranspose, groups multiple global_load together and packs vector across grain to write to LDS with vectorization, so that when the matrix is written into LDS, it's already consecutive on K dimension, and therefore vectorized ds_read is also enabled. This is achieved by v_perm_b32 assembly instruction in AMDGCN, allowing independent register to be contiguous in VGPR space, so that we can write them together into LDS.

PR description

Continuous from the previous PR, this one updates the SharedEncodingAttr on AMD hardware such that SharedEncodingAttr will always guarantee coalesces LDS read, and let the inThreadTranspose to explore if LDS write can be coalesced on non-KContig tensor.

Beyond the TTGIR level update, lowerToLLVM has to be updated in order to properly align the shared memory address to the vector data about to be written. storeDistributedToShared is update to activate special linear layout conversion for blocked layout when it sees blocked -> shared transfer with different order.

In summary, there're three changes:

1. Change SharedEncodingAttr to guarantee coalesces LDS read and swap the sharedEncoding order if the tensor is non-KContig.
2. Activate the previously added blockedToLinearLayoutThreadRake when transferring blocked to shared for non-KContig.
3. A new perVectorCallback to align the shared memory data to the address.