[TIR][Transform] Introduce new InjectPermutedLayout pass

[Ubospica]

This PR introduces the new InjectPermutedLayout pass, partially adapted from the old version.

This pass will permute the visiting indices of the shared memory with xor operations to avoid bank conflicts. For permutation details see p39 of these slides.

The current interface is:

with T.block("..."):
    T.block_attr({"permuted_layout": 1})
    # load or store to shared memory

The load/store operation could be:

• BufferLoad node
• BufferStore node
• ptx_ldmatrix call
• mma_store call

And the InjectPermutedLayout pass is injected in the lowering process in driver_api.cc. So after this pass in the building stage, the visiting indices will be permuted.

Notably this pass works after the CompactBufferAllocation pass because it relies on the shape of the shared memory in every threadblock.

cc @spectrometerHBH @vinx13 @Hzfengsy

[Ubospica]

@tvm-bot rerun

[Ubospica]

@tvm-bot rerun