[mlir][vector][PSA] Disallow pure reduction dims in vector.contract

MLIR
1

Let me draw the context first. In the following example, there are two reduction dimensions:

  • r1 is a “pure reduction” dimension (only present in lhs),
  • r2 is a “contraction” dim (present in both lhs and rhs).
%res = vector.contract {
  indexing_maps = [
    affine_map<(r1, p, r2) -> (r1, r2)>,
    affine_map<(r1, p, r2) -> (p, r2)>,
    affine_map<(r1, p, r2) -> (p)>
  ],
  iterator_types = ["reduction", "parallel", "reduction"],
  kind = #vector.kind<add>
} %lhs, %rhs, %acc : vector<7x3xi32>, vector<2x3xi32> into vector<2xi32>

In #237709, I am aligning vector.contract with linalg.contract by disallowing “pure reduction” dims like r1.

Why is this happening?

There is already some rationale in the PR regarding the MLIR implementation - lowering support for this feature is currently limited to unit-size “pure reduction” dims (effectively, no support).

I also checked corresponding operations in other dialects and couldn’t find support for “pure reduction” dims there, see e.g.:

Both are very well documented references.

The only example I found that supports “pure reduction” dims was ONNX’s Einsum. However, ONNX-MLIR uses decomposition (cf. inference_backend.py):

Limited to the types supported by ReduceSum and MatMul (which we decompose to in most cases) (…)

I feel that Vector should also opt for decomposition.

Isn’t Vector loosing expressability?

No. For the example above, after my PR, you will still be able to represent the computation in Vector, but more explicitly:

%zero = arith.constant dense<0> : vector<3xi32>

// Reduce LHS over r1.
%reduced_lhs = vector.multi_reduction <add>, %lhs, %zero [0]
  : vector<7x3xi32> to vector<3xi32>

// res[p] = acc[p] + sum_r2 reduced_lhs[r2] * rhs[p, r2].
%res = vector.contract {
  indexing_maps = [
    affine_map<(p, r2) -> (r2)>,
    affine_map<(p, r2) -> (p, r2)>,
    affine_map<(p, r2) -> (p)>
  ],
  iterator_types = ["parallel", "reduction"],
  kind = #vector.kind<add>
} %reduced_lhs, %rhs, %acc : vector<3xi32>, vector<2x3xi32> into vector<2xi32>

Why PSA instead of RFC?

Comments are very welcome, and please let me know if I missed anything. As I haven’t found a use case that requires keeping “pure reduction” dims inside vector.contract, I am treating this primarily as an announcement. That said, if you rely on this feature, I’d particularly like to hear from you.

Thank you for taking a look!
-Andrzej

2

This is related to [mlir][vector] Switch from vector.multi_reduction to vector.contract · Issue #234238 · llvm/llvm-project · GitHub

I support this proposal, as was discussed in the issue above. Thanks for doing this!

3

All in all, I support this PSA, too. I’ll leave comments in the dedicated PRs and issues.
Thank you, @banach-space!