[mlir] Vectorize directly to a named contraction in e2e tests

[banach-space]

Following #147296, we should try simplify our lowering pipelines for e2e tests. Specifically,

• We should lower directly to named contractions, i.e. vector.contract, rather than via vector.multi_reduction.

Below are links to two examples. Specifically, lowering steps that we should be able to avoid if we lowered directly to vector.contract.

llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir

Lines 82 to 86 in 0d989b2

	// Step 3: Simplify
	// vector.multi_reduction --> vector.contract
	// Generates a 6-dim vector.contract with the dim matching the original MMT4D Op
	// and with the following split into parallel and reduction dims:
	// * parallel, parallel, reduction, parallel, parallel, reduction

llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/ArmSVE/matmul.mlir

Lines 105 to 111 in 0d989b2

	// Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns)
	transform.apply_patterns to %func {
	transform.apply_patterns.vector.reduction_to_contract
	transform.apply_patterns.vector.transfer_permutation_patterns
	transform.apply_patterns.vector.lower_masked_transfers
	transform.apply_patterns.vector.sink_ops
	} : !transform.op<"func.func">

Another example is being added here: #157815

Note, in order to enable lowering directly to vector.contract, update:

  transform.structured.vectorize %op_to_vectorize  vector_sizes  [...]  : !transform.any_op

as

  transform.structured.vectorize %op_to_vectorize  vector_sizes  [...]   {create_named_contraction} : !transform.any_op

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Following #147296, we should try simplify our lowering pipelines for [e2e tests](https://github.com/llvm/llvm-project/tree/0d989b2aefe8d8f66981f1d39f56ce0214b5de86/mlir/test/Integration/Dialect/Linalg/CPU). Specifically, * We should lower directly to named contractions, i.e. `vector.contract`, rather than via `vector.multi_reduction`.

Below are links to two examples. Specifically, lowering steps that we should be able to avoid if we lowered directly to vector.contract.

llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/mmt4d.mlir

Lines 82 to 86 in 0d989b2

	// Step 3: Simplify
	// vector.multi_reduction --> vector.contract
	// Generates a 6-dim vector.contract with the dim matching the original MMT4D Op
	// and with the following split into parallel and reduction dims:
	// * parallel, parallel, reduction, parallel, parallel, reduction

llvm-project/mlir/test/Integration/Dialect/Linalg/CPU/ArmSVE/matmul.mlir

Lines 105 to 111 in 0d989b2

	// Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns)
	transform.apply_patterns to %func {
	transform.apply_patterns.vector.reduction_to_contract
	transform.apply_patterns.vector.transfer_permutation_patterns
	transform.apply_patterns.vector.lower_masked_transfers
	transform.apply_patterns.vector.sink_ops
	} : !transform.op<"func.func">

[egebeysel]

Just a quick observation here from my side, the vectorize as named contraction path does not seem to use assumeDynamicDimsMatchVecSizes , see here. So that might be the missing pattern :) I had an ad-hoc patch to address that here.