[VectorDistribution] Configure contraction layouts at linalg level

[Groverkss]

This patch moves layout anchoring for contractions to linalg level. This is primarily motivated by allowing us to decide layouts based on lowering_config.

Future patches will also move the transfer_read anchoring to linalg level.

[qedawkins]

This overall approach looks fine to me w/ a few comments. I would also make sure the right guard rails are in place so that compilation fails if a vector.contraction op fails to end up with a layout (i.e. vector_ext.to_layout gets disconnected from the vector.contract somehow).

[Groverkss]

This overall approach looks fine to me w/ a few comments. I would also make sure the right guard rails are in place so that compilation fails if a vector.contraction op fails to end up with a layout (i.e. vector_ext.to_layout gets disconnected from the vector.contract somehow).

I'm not sure what kind of guard rails you are looking for. I could make the distribution patterns for vector.contract check if it's getting the right layout for the contraction. What do you think?

[qedawkins]

Paging this pipeline back into memory, if a vector.contract fails to get a layout does the pass fail? If so that's good enough for guard rails.

[raikonenfnu]

For documentation purposes: (so I don't need to regain context if I am reading this PR again in the future)

1. Change VectorContractOpInfo to support both linalgOp and vector::ContractionOp using VectorContractOpInfo::fromIndexingMapArray
2. Re-purpose getContractionLayout to work on linalgOp instead of contractOp
3. Refactor the GPUVectorAlloc pass to promote inputs to ToLayoutOp, who have a sharedMemoryConversion attribute set.
4. Move SetContractionAnchors to Tensor level

LMK if anything here is missing. :)

[raikonenfnu]

Overall looks great! Just some last Qs.

[raikonenfnu]

[No action needed here]: I think I'd have preferred two constructors over VectorContractOpInfo::inferFromIndexingMaps, but that's just styling preference.

[raikonenfnu]

Looks like vmfb is not being generated for the MI300 SDXL test as seen on here, maybe we need to remove attention pipeline from the TD tuning script, since we are introducing new way to anchor the layouts.

CC: @saienduri

[Groverkss]

Looks like vmfb is not being generated for the MI300 SDXL test as seen on here, maybe we need to remove attention pipeline from the TD tuning script, since we are introducing new way to anchor the layouts.

CC: @saienduri

That's an issue with the cache, unrelated to this pr. See: https://discord.com/channels/689900678990135345/1166024193599615006/1276556104515584032

[ScottTodd]

#18336 for the cache issue

[raikonenfnu]

Looks like vmfb is not being generated for the MI300 SDXL test as seen on here, maybe we need to remove attention pipeline from the TD tuning script, since we are introducing new way to anchor the layouts.
CC: @saienduri

That's an issue with the cache, unrelated to this pr. See: https://discord.com/channels/689900678990135345/1166024193599615006/1276556104515584032

#18336 for the cache issue

Good to know! thanks guys :)

[Groverkss]

Paging this pipeline back into memory, if a vector.contract fails to get a layout does the pass fail? If so that's good enough for guard rails.

It didn't before, but I just added a check that does this check. So that should be enough for ensuring the layout matches thee intrinsic.

[qedawkins]

Nice, thanks!