Towards optimal matvec vector level lowering

[efric]

This discussion proposes a set of vector lowering related changes motivated by the work on matvec improvements for Qwen MoE decode shapes. The proposal here is intentional about the codegen specific details, please see #24165 for the full context on this work.

Context

The proposed policy for matvec-like per thread vector lowering is (post vector distribute):

if the inputs are naturally <=16 bit:
  lower through a chain dot path
else if inputs are 32 bit or cannot use dot:
  lower through a chain fma path
else:
  keep the normal generic fallback

For full details, please see #24165. Below serves a short description for the semantics of chain dot and chain FMA.

Chain Dot:
In this strategy, we reduce within each row/source vector load with respect to a per thread vector e.g, vector<4x8xf16>.

Using v_dot2 as an example:

Chain v_dot
For each 1D `vector<8xf16>` that we want to arith.mulf -> vector.reduction <add> on:
===>
dot(a[6:8], b[6:8], dot(a[4:6], b[4:6], dot(a[2:4], b[2:4], dot(a[0:2], b[0:2], 0))))

There are no v_dot instructions that accept 32bit operands across any AMDGPU architecture, so when that is the case, we defer to the current lowering, which leverages chain FMA. This is still better than the generic lowering, which does packed mul across each row then v_add to sum the results.

Chain FMA
For each vector load (row), slice across column i, decompose to slices across 2 rows each (if packed is available, if not then 1 row each), then 
===>
fma(a[7], b[7], fma(a[6], b[6], fma(a[5], b[5], fma(a[4], b[4], fma(a[3], b[3], fma(a[2], b[2], fma(a[1], b[1], fma(a[0], b[0])))))

In short, the difference between dot and packed fma in the semantics is the following:

dot (2): a[0] * b[0]  + a[1] * b[1] + c
fma (packed): {a[0] * b[0] + c[0], a[1] * b[1] + c[1]}

Compared to the dot chain, the instruction dependency is twice as long for packed fma (on f16). If packed instructions are available, then because we are retrieving data from different vector loads, we need v_movs in preparation for the pk_fmas.

Early testing results using dot vs fma on RDNA3 for Qwen MoE decode shapes (pre lowering config adjustments)

Shape (M×N×K)	baseline (µs)	dot (µs)	speedup
1×32×2048	37.0	36.8	1.01×
1×128×2048	37.1	37.4	0.99×
1×256×2048	37.0	37.6	0.98×
1×512×2048	38.1	37.7	1.01×
1×768×2048	39.6	38.2	1.04×
1×2048×512	37.8	37.1	1.02×
1×2048×768	38.5	38.5	1.00×
1×2048×4096	44.7	43.3	1.03×
1×2048×5632	44.4	44.7	0.99×
1×4096×2048	49.9	44.4	1.12×
1×8192×2048	54.2	46.4	1.17×
1×151936×2048	806	764	1.05×
1×248320×2048	1253	1226	1.02×
1×6656×16384	301	280	1.07×

However, I expect the benefits to span all architectures which support v_dot

High level design

1. "Natural/Canonical/Unpacked" form while preserving semantics in VectorDistribute for vector.contract/vector.multi_reduction:

Currently, we construct the lowering to chain FMA in LLVMGPUVectorLowering via our ownContractToChainFMA. However, if we follow the VectorContractLowering::OuterProduct strategy, we will actually naturally lower to the same output following: chain emission into fma. The reason why this wasn't picked up even though we pick VectorContractLowering::OuterProduct as our strategy is because post vector distribute, relevant ops such as vector.contract and vector.multi_reduction are in their expanded forms.

%157 = vector.contract {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d1, d2, d4, d5, d7, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2, d3, d4, d5, d6, d7, d8)>, affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d3, d4, d6, d7)>], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %121, %156, %cst_4 : vector<1x1x1x1x1x8xf16>, vector<8x1x1x1x1x1x1x1x8xf16> into vector<8x1x1x1x1x1xf32>

This makes it unable to match the supported strategies in VectorContractLowering::OuterProduct. However, if we were able to able to use its natural form while preserving vector distribution semantics:

 %147 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"], kind = #vector.kind<add>} %111, %146, %cst_0 : vector<8xf16>, vector<8x8xf16> into vector<8xf32>

Then we can remove ContractToChainFMA altogether and depend on upstream for lowering to the chain FMA strategy following VectorContractLowering::OuterProductnaturally.

2. Composable vector.contract lowering

Once the the contract is in the canonical shape, the decision should be made per contract

sub-f32 dot candidate:
  use VectorContractLowering::Dot

natural f32:
  use VectorContractLowering::OuterProduct

fallback:
  use generic progressive lowering

This is because VectorContractLowering::Dot eventually decomposes to a series of 1D vector.reduction which we can then (next section) match to the dot chain. VectorContractLowering::OuterProduct eventually to chain FMA naturally as mentioned above.

The current upstream API selects one preferred strategy globally:

vector::populateVectorContractLoweringPatterns(
    patterns, vector::VectorContractLowering::OuterProduct);

Which isn't enough to express the policy we'd like. The proposed extension is something to the liens of

using VectorContractLoweringFilter =
    std::function<LogicalResult(vector::ContractionOp)>;

LogicalResult acceptAllVectorContractLoweringFilter(vector::ContractionOp op);

void populateVectorContractToDotPatterns(
    RewritePatternSet &patterns,
    VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
    PatternBenefit benefit = 1);

void populateVectorContractToOuterProductPatterns(
    RewritePatternSet &patterns,
    VectorContractLoweringFilter filter = acceptAllVectorContractLoweringFilter,
    PatternBenefit benefit = 1);

And then in IREE we can express this directly as roughly

vector::populateVectorContractToDotPatterns(contractLoweringPatterns, preferDotContract, PatternBenefit(3));
vector::populateVectorContractToOuterProductPatterns(contractLoweringPatterns, vector::acceptAllVectorContractLoweringFilter, PatternBenefit(2));
vector::populateVectorContractGenericLoweringPatterns(contractLoweringPatterns, vector::acceptAllVectorContractLoweringFilter, PatternBenefit(1));

3. Vector Reduction to amdgpu.dot

This is the most straightforward implementation, we will just match 1D vector.reduction <add> into amdgpu.dot chain (of course, if only feasible in terms of vector length, architecture, type, etc).

%lhs_ext = arith.extf %lhs : vector<8xf16> to vector<8xf32>
%rhs_ext = arith.extf %rhs : vector<8xf16> to vector<8xf32>
%mul = arith.mulf %lhs_ext, %rhs_ext : vector<8xf32>
%sum = vector.reduction <add>, %mul : vector<8xf32> into f32

===>

%zero = arith.constant 0.0 : f32
%lhs0 = vector.extract_strided_slice %lhs
  {offsets = [0], sizes = [2], strides = [1]}
  : vector<8xf16> to vector<2xf16>
%rhs0 = vector.extract_strided_slice %rhs
  {offsets = [0], sizes = [2], strides = [1]}
  : vector<8xf16> to vector<2xf16>
%dot0 = amdgpu.dot %lhs0 * %rhs0 + %zero
  : vector<2xf16>, vector<2xf16>, f32
%lhs1 = vector.extract_strided_slice %lhs
  {offsets = [2], sizes = [2], strides = [1]}
  : vector<8xf16> to vector<2xf16>
%rhs1 = vector.extract_strided_slice %rhs
  {offsets = [2], sizes = [2], strides = [1]}
  : vector<8xf16> to vector<2xf16>
%dot1 = amdgpu.dot %lhs1 * %rhs1 + %dot0
  : vector<2xf16>, vector<2xf16>, f32
...

This can then run some time after LLVMGPUVectorLowering, specifically in this block

 if (forROCDL) {
    // This pass needs to run after the LLVMGPUVectorLoweringPass.
    funcPassManager.addPass(amdgpu::createAmdgpuVectorReductionToDotPass);
    funcPassManager.addPass(amdgpu::createAmdgpuMaskedloadToLoadPass);
    // This pass needs to run before the ResolveSwizzleHints pass.
    funcPassManager.addPass(amdgpu::createAmdgpuFoldMemRefOpsPass);
  }

seems reasonable.

cc some folks for feedback on reasonableness of this approach at a high level:
@Groverkss @amd-eochoalo @kuhar @krzysz00

Overall, this direction is more involved than adding a monolithic iree local ContractToChainDot pattern, but I think it is better aligned with the more “composable, intentional” vector lowering direction that Erick and Kunwar has been working on. It also lets us reuse more upstream lowering infrastructure instead of adding another target specific contract lowering path in IREE. Step 1 should also be useful beyond the current matvec contracts as it reveals what may currently be hidden behind the distributed layout form and therefore deferred to generic lowering when it doesn't have to.

[amd-eochoalo]

@efric this is what I understood from here

1. vector.contract operations acting on floating point types with a bitwidth less than 32 bits will be matched by ContractionOpToDotLowering.
2. vector.contract operations acting on f32 types in "natural/canonical/unpacked form" will be matched by ContractionOpToOuterProductOpLowering
3. all other vector.contract operations will be matched by ContractionOpLowering

For case 1, these will also require "natural/canonical/unpacked form" right? (At the very least, what I can see from the LLVM lowerings is that for the matvec case, the operation would need to have the reduction dimension to be innermost.

else if (isParallelIterator(iteratorTypes[0]) &&
             isReductionIterator(iteratorTypes[1])) {
    //
    // One outer parallel, one inner reduction (matvec flavor)
    //

Regarding the PRs I worked on (which attempted to split the ChainToFMA pattern and make them fit into the "optimize", "flatten", "unroll", "lower" pattern. The "optimize" part here could be done I think for both of these, it essentially just establishes the invariant that all reduction dimensions are innermost.

For both of these, we then can flatten them into 2-D and 1-D (for ChainToFMA since it is acting on matmat, then we flatten both operands to 2-D).

For unrolling, we can call lowerParallel to unroll parallel dimensions (or something similar that ideally should be added to the ND legalization pass)

And finally, the lowering for vector.contract would depend on whether it is a matmat (fma) or a matvec and the specific types (dot or fma + other ops).

This is different from what you are proposing, but I would like to hear from @Groverkss as well. My understanding about this legalization pass is that we can also hard code lowerings before the 1-D unrolling as the legalization would have just some "unrollings" which are general and work for all backends. But anything that is backend specific should happen before this 1-D legalization. I'm still not sure about where that split is (backend specific vs general enough).

[efric]

Yes, case 1 will also require that form (what step 1 in the design i propose would fix, which would be when we distribute contract in vector distribute, same for multi_reduction). I think that will be a good move regardless of how we subsequently lower the vectors. Also, the dot lowering is not restricted to floats. Some architectures support packed ints, e.g v_dot4_i32_iu8 on gfx12+.

I don't understand complete the need to do the optimize/flatten/unroll/lower thing manually for "both". We can just delete ChainFMA in IREE since upstream will generate it if we massage the IR right (basically just give 2D contract rather than obfuscated unpacked nested layout form). Same thing for dot if we just select ContractionOpToDotLowering, then we can do the 1D lowering to dot the way we'd like with amdgpu::createAmdgpuVectorReductionToDotPass (I mean we can just keep this in IREE too, but it would be consistent with other backends stories that we provide an upstream path to lower to amdgpu dot intrinsics like x86 or spirv already does).

Or is the idea that we don't want to pass any ND vectors at all anymore to upstream?

Also re:

And finally, the lowering for vector.contract would depend on whether it is a matmat (fma) or a matvec and the specific types (dot or fma + other ops).

it can be matvec for both; that's why i want to extend the strategy selection story. But I guess if we want to keep it all in IREE then we can do something different.

[amd-eochoalo]

Or is the idea that we don't want to pass any ND vectors at all anymore to upstream?

I think this question is a bit ambiguous. "At what point do we not want to pass ND vectors to upstream?" would be a better question. The answer is during LLVM lowering there should be no ND vectors.

. But I guess if we want to keep it all in IREE then we can do something different.

I don't think there is an official "keep everything in IREE" mandate. Just being pragmatic about big changes. I think it would be good to push this upstream, but it is a big change and we would need to coordinate a bit how we upstream it, communicate about it.

[amd-eochoalo]

We had a chat about this. I remember I was told during the meeting that essentially every operation that needs some specialized lowering needs to be done before unrolling and that everything else would just go through the legalization. @efric and I chatted and we agreed that this implies that @efric's plan of using the upstream patterns before legalization should be sufficient and would not interfere with the n-d legalization / unrolling / flattening, etc... cc @Groverkss