This is a follow-up to Linalg RFC and next step in adding native scaled contraction support to MLIR.
Summary
This RFC proposes a new Vector operation vector.scaled_contract to allow representing contractions in which the inputs A and B are scaled using additional scale values.
The operation primarily targets lower-level hardware intrinsics implementing Open Compute Project (OCP) Microscaling Formats (MX) Specification.
Prior Art
For the general overview of the MX format and scaled contraction definition, I recommend the introduction section of the Linalg RFC.
Relevant hardware intrinsics and lowering targets for vector.scaled_contract:
-
XeGPU dialect scaled mma op [link]
-
NVVM dialect block mma op [link]
-
x86 new ACE intrinsics [spec, LLVM PR] - see spec’s section 10 (p. 51) for overview and section 14 (p. 76) for compute instructions
Op: vector.scaled_contract
A new Vector op that supports general contraction semantics with extra scale inputs together with flexible scale indexing maps which can be used to represent various scaling schemes.
The aim is to create a pair of operations analogous to linalg.contract and vector.contract that bridge the memory (tensor, memref) and register (vector) worlds and serve as counterpart representations of general contractions in their respective abstraction levels.
The proposed vector.scaled_contract extends vector.contract and closely mirrors the scaling semantics of linalg.scaled_contract, serving as its register-level counterpart.
Similarly to linalg.scaled_contract, block scaling can be represented using expanded shapes:
%D = vector.scaled_contract
{indexing_maps = [
affine_map<(m, n, ko, kb) -> (m, ko, kb)>,
affine_map<(m, n, ko, kb) -> (ko)>,
affine_map<(m, n, ko, kb) -> (n, ko, kb)>,
affine_map<(m, n, ko, kb) -> (ko)>,
affine_map<(m, n, ko, kb) -> (m, n)>],
iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
%A, %scaleA, %B, %scaleB, %C
: vector<16x2x32xf8E5M2>, vector<2xf8E8M0FNU>,
vector<32x2x32xf8E4M3FN>, vector<2xf8E8M0FNU> into vector<16x32xf32>
or semi-affine indexing maps that mirror the linalg.scaled_contract semantics:
%D = vector.scaled_contract
{indexing_maps = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (k floordiv 32)>,
affine_map<(m, n, k) -> (n, k)>,
affine_map<(m, n, k) -> (k floordiv 32)>,
affine_map<(m, n, k) -> (m, n)>],
iterator_types = ["parallel", "parallel", "reduction"]}
%A, %scaleA, %B, %scaleB, %C
: vector<16x64xf8E5M2>, vector<2xf8E8M0FNU>,
vector<32x64xf8E4M3FN>, vector<2xf8E8M0FNU> into vector<16x32xf32>
Initially, the op adapts the same data type restrictions as its Linalg counterpart namely:
The output element type must be floating-point. Input element types must be
integer or floating-point, with bitwidth no greater than the output element
type’s bitwidth.
to prevent over-generalization of scaling semantics. This can be relaxed later as needed.
Additionally, unlike vector.contract today, the new op supports different element data types for its operands to align with commonly supported mixed low-precision variants such as: f8E5M2 * f8E4M3FN.
Combination function kind attribute could be skipped for now to start with the default add or kind could be added to more closely align with vector.contract.
Ingress
Direct vectorization from linalg.scaled_contract.
Lowering
The main goal of vector.scaled_contract is to serve as an intermediate stage toward scaled hardware intrinsics. For example, the operation can directly lowered into xegpu.dpas_mx given compatible shapes and types. Alternatively, further layout pre-processing can be performed locally at the time of lowering to align with hardware instruction requirements like, for example, online VNNI packing is applied today for x86 AMX compute ops.
As a fallback, vector.scaled_contract can be rewritten into a sequence of standard vector operations. An example rewrite strategy:
- partition K by scale block
- compute low-precision contraction per block
- multiply by combined scale
- accumulate final result
Given a semi-affine block scaled example:
func.func @vec_block_scaling_maps(
%A: vector<16x64xf8E5M2>, %scaleA: vector<2xf8E8M0FNU>,
%B: vector<32x64xf8E5M2>, %scaleB: vector<2xf8E8M0FNU>,
%C: vector<16x32xf32>) -> vector<16x32xf32> {
%D = vector.scaled_contract
{indexing_maps = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (k floordiv 32)>,
affine_map<(m, n, k) -> (n, k)>,
affine_map<(m, n, k) -> (k floordiv 32)>,
affine_map<(m, n, k) -> (m, n)>],
iterator_types = ["parallel", "parallel", "reduction"]}
%A, %scaleA, %B, %scaleB, %C
: vector<16x64xf8E5M2>, vector<2xf8E8M0FNU>,
vector<32x64xf8E5M2>, vector<2xf8E8M0FNU> into vector<16x32xf32>
return %D : vector<16x32xf32>
}
can be decomposed as:
func.func @vec_block_scaling_decomposed(
%A: vector<16x64xf8E5M2>, %scaleA: vector<2xf8E8M0FNU>,
%B: vector<32x64xf8E5M2>, %scaleB: vector<2xf8E8M0FNU>,
%C: vector<16x32xf32>) -> vector<16x32xf32> {
%zero = arith.constant dense<0.0> : vector<16x32xf32>
// Extract scale factors.
%sA0 = vector.extract %scaleA[0]
: f8E8M0FNU from vector<2xf8E8M0FNU>
%sA1 = vector.extract %scaleA[1]
: f8E8M0FNU from vector<2xf8E8M0FNU>
%sB0 = vector.extract %scaleB[0]
: f8E8M0FNU from vector<2xf8E8M0FNU>
%sB1 = vector.extract %scaleB[1]
: f8E8M0FNU from vector<2xf8E8M0FNU>
// Convert scales to accumulator type.
%sA0f = arith.extf %sA0 : f8E8M0FNU to f32
%sA1f = arith.extf %sA1 : f8E8M0FNU to f32
%sB0f = arith.extf %sB0 : f8E8M0FNU to f32
%sB1f = arith.extf %sB1 : f8E8M0FNU to f32
%scale0 = arith.mulf %sA0f, %sB0f : f32
%scale1 = arith.mulf %sA1f, %sB1f : f32
// Contraction per scaled block.
// K block 0 : [0, 32)
%A0 = vector.extract_strided_slice %A
offsets = [0, 0], sizes = [16, 32], strides = [1, 1]
: vector<16x64xf8E5M2> to vector<16x32xf8E5M2>
%B0 = vector.extract_strided_slice %B
offsets = [0, 0], sizes = [32, 32], strides = [1, 1]
: vector<32x64xf8E5M2> to vector<32x32xf8E5M2>
%part0 = vector.contract
{indexing_maps = [
affine_map<(m,n,k) -> (m,k)>,
affine_map<(m,n,k) -> (n,k)>,
affine_map<(m,n,k) -> (m,n)>],
iterator_types = ["parallel","parallel","reduction"],
kind = #vector.kind<add>}
%A0, %B0, %zero
: vector<16x32xf8E5M2>, vector<32x32xf8E5M2>
into vector<16x32xf32>
%scale0vec = vector.broadcast %scale0 : f32 to vector<16x32xf32>
%part0scaled = arith.mulf %part0, %scale0vec
: vector<16x32xf32>
// K block 1 : [32, 64)
%A1 = vector.extract_strided_slice %A
offsets = [0, 32], sizes = [16, 32], strides = [1, 1]
: vector<16x64xf8E5M2> to vector<16x32xf8E5M2>
%B1 = vector.extract_strided_slice %B
offsets = [0, 32], sizes = [32, 32], strides = [1, 1]
: vector<32x64xf8E5M2> to vector<32x32xf8E5M2>
%part1 = vector.contract
{indexing_maps = [
affine_map<(m,n,k) -> (m,k)>,
affine_map<(m,n,k) -> (n,k)>,
affine_map<(m,n,k) -> (m,n)>],
iterator_types = ["parallel","parallel","reduction"],
kind = #vector.kind<add>}
%A1, %B1, %zero
: vector<16x32xf8E5M2>, vector<32x32xf8E5M2>
into vector<16x32xf32>
%scale1vec = vector.broadcast %scale1 : f32 to vector<16x32xf32>
%part1scaled = arith.mulf %part1, %scale1vec
: vector<16x32xf32>
// Final accumulation.
%sum = arith.addf %part0scaled, %part1scaled
: vector<16x32xf32>
%res = arith.addf %C, %sum
: vector<16x32xf32>
return %res : vector<16x32xf32>
}
Alternatives
- General Intrinsic Op
In the related Linalg RFC, @Groverkss proposed a more general operation vector.inner_intrinsic:
My personal interpretation: The inner_intrinsic operation is a flexible container that can represent many different operations going beyond only scaled contraction. That would allow to limit number of specialized named operations in Vector dialect. Patterns can verify op’s exact semantics by checking attribute and/or body contents.
With clearly defined boundaries to avoid too abstract generic implementation, it could be a valid alternative. However, I have not analyzed this case thoroughly as it goes way beyond my target scope.
Happy to hear more feedback and experience from others for such approach.
- Pattern Matching
Similar alternative option presented also in the Linalg RFC.
Instead of a new operation, scaled contraction could be represented as a sub-graph of operations.
It requires more expensive pattern matching, and it creates an odd representation gap where high-level (Linalg) and low-level (hardware intrinsics) capture scaled contraction in a single operation.
By preserving whole scaled contraction semantics in a single operation, it will be easier to connect Vector dialect with its surrounding abstractions.
- Extending vector.contract
Existing vector.contract op could be extended to represent scaled contraction by optionally accepting extra operands and indexing maps for scales e.g., scale_a = [%scale_val, #scale_map] etc.
This option is the most invasive and would require add extra checks to existing code. However, it eliminates need for a new operation.
Overall, I’d avoid going this route due to possible unforeseen interactions with the existing code.
Implementation Plan
Proposed steps:
-
Land the new operation according to reached consensus.
-
Add Linalg vectorization.
-
Add lowering to XeGPU dialect.
-
Later, add lowering to x86 dialect when ACE is available in LLVM.
I would leave out the decomposition rewrite fallback until it is needed.
=====================================================
CC @banach-space @dcaballe @Groverkss @Jianhui-Li @rengolin @arunt1204