plaidml / tpp-mlir

TPP experimentation on MLIR for linear algebra
https://arxiv.org/abs/2404.15204
Other
111 stars 31 forks source link

IntelAMX TileConfig Hoisting should be anchored on IntelAMXTileConfigDispatchOp #896

Open alheinecke opened 8 months ago

alheinecke commented 8 months ago

The Intel AMX TileConfig Hoisting uses today the AllocaOp for the tileconfig state as anchor op and then attempts to move IntelAMXTileConfigDispatchOp around based on a test if the AllocOp is used by IntelAMXTileConfigDispatchOp.

https://github.com/plaidml/tpp-mlir/blob/main/lib/TPP/Transforms/IntelAMXTileConfigHoisting.cpp#L31

However, we are now attempting to rewrite every alloca and error out if the alloca is not used by IntelAMXTileConfigDispatchOp.

It would a cleaner design if we would rewrite based on IntelAMXTileConfigDispatchOp. However, I see also challenges here, as we have 2 Ops around a BRGEMM, hence the alloca-"Trick" in the current pass made it simple to rewrite.

Let's discuss suggestions on how to make this cleaner.