Transform to lower packs/unpacks without transpose, except when the packing occurs on a constant

[rolfmorel]

A transform that enables selectively "reverting" packing on linalg.generic operands and unpacking on the generic's results by lowering tensor.pack to tensor.expand_shape with identity permutation and tensor.unpack to tensor.collapse_shape with identity permutation.

The new pass -lower-packs-unpacks-without-transpose does this reverting on all (compatible) packed/unpacked linalg.generic operands/results, except if they are constants. Similarly, the CLI flag --lower-pack-unpack-without-transpose enables this pass in the default pipeline.

[rolfmorel]

@adam-smnk , @rengolin

This is now ready for final review (excepting a clear PR description).

[adam-smnk]

Logic looks solid, I'd still cleanup the tests a bit.

[rolfmorel]

Performance on the lora-fragment benchmark in use for Openvino:

$ python -m xonsh lora-runner.xsh 
CONFIGS [[8], [16], [32], [64], [128], [256], [512], [1024], [2048], [4096], [8192]]
OV NO-MLIR [0.84, 1.38, 2.48, 4.7, 9.62, 19.14, 37.54, 74.51, 147.15, 311.04, 628.24]
OV MLIR [3.82, 4.43, 5.35, 8.06, 13.45, 23.87, 44.33, 85.15, 167.73, 340.55, 690.18]
NO-OV MLIR [4.05, 4.54, 5.57, 8.2, 13.47, 23.31, 44.07, 83.64, 166.25, 333.87, 662.45]
MANUAL MLIR [0.54, 1.1, 2.2, 4.43, 8.82, 17.39, 35.26, 70.28, 138.55, 279.32, 561.73]

• OV NO-MLIR vs. MANUAL MLIR compares plain OV vs TPP-MLIR with this new rewrite.
• NO-OV MLIR vs. MANUAL MLIR compares TPP-MLIR (with runtime packing of all arguments and including weights passed as an argument) vs TPP-MLIR with this new rewrite.
• OV MLIR vs. NO-OV MLIR compares TPP-MLIR vs TPP-MLIR (both without the new rewrite and both with runtime packing of all arguments including the weights as an argument) to infer overhead of calling MLIR-generated code from OV.

EDIT: make more accurate as pointed out by Adam below.
EDIT2: the following are runs which vary the LoRA dimension:

$ python -m xonsh lora-runner.xsh 
lora_dim = 2
CONFIGS [[8], [16], [32], [64], [128], [256], [512], [1024], [2048], [4096], [8192]]
OV NO-MLIR [0.86, 1.41, 2.51, 4.72, 9.46, 18.9, 37.65, 74.65, 147.38, 311.4, 624.89]
OV MLIR [3.84, 4.4, 5.47, 8.06, 13.38, 23.82, 44.67, 85.53, 167.85, 337.54, 693.38]
NO-OV MLIR [4.05, 4.67, 5.59, 8.21, 13.44, 23.28, 44.63, 85.13, 164.08, 329.97, 666.76]
MANUAL MLIR [0.55, 1.1, 2.2, 4.4, 8.83, 17.44, 35.31, 70.46, 141.23, 279.89, 561.65]
lora_dim = 4
CONFIGS [[8], [16], [32], [64], [128], [256], [512], [1024], [2048], [4096], [8192]]
OV NO-MLIR [0.87, 1.4, 2.5, 4.71, 9.55, 19.05, 37.65, 73.86, 147.19, 311.13, 630.02]
OV MLIR [3.67, 4.18, 5.4, 8.09, 13.45, 23.79, 44.17, 85.59, 168.4, 341.18, 689.78]
NO-OV MLIR [4.04, 4.62, 5.65, 8.26, 13.44, 23.75, 43.78, 83.24, 166.19, 330.83, 662.83]
MANUAL MLIR [0.55, 1.1, 2.2, 4.4, 8.81, 17.66, 35.26, 70.48, 138.62, 281.63, 563.89]
lora_dim = 8
CONFIGS [[8], [16], [32], [64], [128], [256], [512], [1024], [2048], [4096], [8192]]
OV NO-MLIR [0.86, 1.41, 2.51, 4.71, 9.6, 19.02, 37.65, 74.19, 147.19, 311.58, 629.79]
OV MLIR [3.66, 4.42, 5.48, 8.11, 13.56, 23.81, 44.69, 84.69, 166.77, 341.29, 687.18]
NO-OV MLIR [3.97, 4.65, 5.63, 8.19, 13.41, 23.29, 44.94, 84.8, 162.98, 332.9, 655.18]
MANUAL MLIR [0.55, 1.1, 2.2, 4.41, 8.81, 17.67, 35.28, 69.54, 138.47, 280.48, 564.5]

[adam-smnk]

`> NO-OV MLIR vs. MANUAL MLIR compares TPP-MLIR vs TPP-MLIR with this new rewrite.

Kind of, NO-OV MLIR still has runtime packing due to weights passed as arguments.

[rolfmorel]

@adam-smnk, @rengolin, in case of no further comments I will merge this tomorrow, Sept 25th.

[rolfmorel]

Thank you for the thorough review, @adam-smnk !

I think I have addressed everything, though will wait a couple more hours before I merge.

[adam-smnk]

Thanks for the tweaks 👍
Looks good