1//===- LowerContractionToSVEI8MMPattern.cpp - Contract to I8MM --*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements lowering patterns from vector.contract to operations
10// that map to instructions from the SVE FEAT_I8MM extension.
11//
12// TODO: There may be opportunities to unify this with a similar pattern
13// for Neon. See:
14// https://github.com/llvm/llvm-project/issues/145559
15// LowerContractionToNeonI8MMPattern.cpp
16//
17//===----------------------------------------------------------------------===//
18
19#include "mlir/Dialect/Arith/IR/Arith.h"
20#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
21#include "mlir/Dialect/ArmSVE/Transforms/Transforms.h"
22#include "mlir/Dialect/Func/IR/FuncOps.h"
23#include "mlir/Dialect/Vector/IR/VectorOps.h"
24#include "mlir/IR/AffineMap.h"
25#include "mlir/IR/PatternMatch.h"
26
27#include "mlir/Dialect/UB/IR/UBOps.h"
28
29#define DEBUG_TYPE "lower-contract-to-arm-sve-i8mm"
30
31using namespace mlir;
32
33namespace {
34// Get the operand of a `vector.contract`. This function is intended to abstract
35// away from the particular way a value is extended before feeding it into the
36// `vector.contract` - via zero-extend or an explicit or implicit sign-extend
37// (for implicit sign-extension see `vector.contract` documentation).
38//
39// The template parameter `Op` indicates the extension operation (explicit or
40// implicit) for which we are checking.
41//
42// Return success only for extensions from `i8` to `i32`.
43template <typename Op>
44std::optional<Value> getExtOperand(Value v) {
45
46 static_assert(llvm::is_one_of<Op, arith::ExtSIOp, arith::ExtUIOp>::value,
47 "Must be instantiated with either sign- or zero- extension op");
48
49 // If the operand is not defined by an explicit extend operation of the
50 // accepted operation type allow for an implicit sign-extension.
51 auto extOp = dyn_cast_or_null<Op>(v.getDefiningOp());
52 if (!extOp) {
53 if constexpr (std::is_same<Op, arith::ExtSIOp>::value) {
54 auto vTy = cast<VectorType>(Val: v.getType());
55 if (!vTy.getElementType().isSignlessInteger(width: 8))
56 return {};
57 return v;
58 }
59 return {};
60 }
61
62 // If the operand is defined by an explicit extend operation of the accepted
63 // operation type, check it's extended from `i8` to `i32`.
64 auto inOp = extOp.getIn();
65 auto inTy = dyn_cast<VectorType>(inOp.getType());
66 if (!inTy || !inTy.getElementType().isSignlessInteger(8))
67 return {};
68
69 auto outTy = dyn_cast<VectorType>(extOp.getType());
70 if (!outTy || !outTy.getElementType().isSignlessInteger(32))
71 return {};
72
73 return inOp;
74}
75
76// Designate the operation (resp. instruction) used to do sub-tile matrix
77// multiplications.
78enum class MMLA {
79 Signed, // smmla
80 Unsigned, // ummla
81 Mixed, // usmmla
82 MixedSwapped // usmmla with LHS and RHS swapped
83};
84
85// Create the matrix mulitply and accumulate operation according to `op`.
86Value createMMLA(PatternRewriter &rewriter, MMLA op, Location loc,
87 mlir::VectorType accType, Value acc, Value lhs, Value rhs) {
88 switch (op) {
89 case MMLA::Signed:
90 return rewriter.create<arm_sve::SmmlaOp>(location: loc, args&: accType, args&: acc, args&: lhs, args&: rhs);
91 case MMLA::Unsigned:
92 return rewriter.create<arm_sve::UmmlaOp>(location: loc, args&: accType, args&: acc, args&: lhs, args&: rhs);
93 case MMLA::Mixed:
94 return rewriter.create<arm_sve::UsmmlaOp>(location: loc, args&: accType, args&: acc, args&: lhs, args&: rhs);
95 case MMLA::MixedSwapped:
96 // The accumulator comes transposed and the result will be transposed
97 // later, so all we have to do here is swap the operands.
98 return rewriter.create<arm_sve::UsmmlaOp>(location: loc, args&: accType, args&: acc, args&: rhs, args&: lhs);
99 }
100}
101
102/// Lower a contraction operation that performs a matrix multiplication
103/// of two 8-bit integer matrix tiles with logical dimensions <Mx8> and <8x[N]>
104/// for the left-hand side and the right-hand side, respectively,
105/// yielding a <Mx[N]> 32-bit integer result.
106///
107/// The operands' shapes are such that the operands can be evenly split into
108/// sub-tiles with dimensions as expected by the targeted FEAT_I8MM
109/// instructions. The intent is that M and N are chosen (by higher level
110/// transforms) in such a way as to maximise register usage. The main use case
111/// we envision as of now is MMT4D, thus the RHS operand is expected
112/// pre-transposed.
113///
114/// The matrix multiplication is performed by unrolling the usual tiled matrix
115/// multiplication algorithm using sub-tiles with dimensions <2x8> for the LHS,
116/// <8x[2]> for the RHS, and <2x[2]> for the result and the input accumulator.
117///
118/// One way to illustrate the operation is as follows:
119///
120/// RHS<8x[N]>: <8x[2]> <8x[2]> ... <8x[2]>
121/// +-----------------------------
122/// LHS<Mx8>: <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
123/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
124/// ... | ... ... ... ...
125/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
126///
127/// The RHS operand is unpacked into N/2 values, each representing a sequence of
128/// VSCALE number of sub-tiles with dimensions <8x2>.
129/// The LHS operand is initially unpacked into M/2 values, each representing a
130/// sub-tile with dimensions <2x8>, and then each such sub-tile is replicated
131/// VSCALE times.
132/// Multiplying thus replicated LHS sub-tile by the corresponding RHS sub-tile
133/// correctly computes an entire result sub-tile.
134class LowerContractionToSVEI8MMPattern
135 : public OpRewritePattern<vector::ContractionOp> {
136public:
137 using OpRewritePattern::OpRewritePattern;
138 LogicalResult matchAndRewrite(vector::ContractionOp op,
139 PatternRewriter &rewriter) const override {
140
141 Location loc = op.getLoc();
142 mlir::VectorType lhsType = op.getLhsType();
143 mlir::VectorType rhsType = op.getRhsType();
144
145 // Check the rank the types so we can safely examine their dimensions.
146 if (lhsType.getRank() != 2 || rhsType.getRank() != 2)
147 return rewriter.notifyMatchFailure(arg&: op, msg: "non-matching operand shape");
148
149 auto M = lhsType.getDimSize(idx: 0);
150 auto N = rhsType.getDimSize(idx: 0);
151 auto K = rhsType.getDimSize(idx: 1);
152
153 // Check the operands have the expected shape:
154 // * for LHS: fixed vector MxK
155 // * for RHS: scalable vector [N]xK
156 // * K == 8
157 // * M and N even and at least 2
158 if (lhsType.isScalable() || !rhsType.getScalableDims()[0] ||
159 rhsType.getScalableDims()[1] || lhsType.getDimSize(idx: 1) != K || K != 8 ||
160 M < 2 || M % 2 != 0 || N < 2 || N % 2 != 0 ||
161 !rhsType.getScalableDims()[0])
162 return rewriter.notifyMatchFailure(arg&: op, msg: "non-matching operand shape");
163
164 // Check permutation maps. For now only accept
165 // lhs: (d0, d1, d2) -> (d0, d2)
166 // rhs: (d0, d1, d2) -> (d1, d2)
167 // acc: (d0, d1, d2) -> (d0, d1)
168 // This corresponds to matrix multiplication with transposed RHS.
169 if (op.getIndexingMapsArray()[0] !=
170 AffineMap::getMultiDimMapWithTargets(numDims: 3, targets: ArrayRef{0u, 2u},
171 context: op.getContext()) ||
172 op.getIndexingMapsArray()[1] !=
173 AffineMap::getMultiDimMapWithTargets(numDims: 3, targets: ArrayRef{1u, 2u},
174 context: op.getContext()) ||
175 op.getIndexingMapsArray()[2] !=
176 AffineMap::getMultiDimMapWithTargets(numDims: 3, targets: ArrayRef{0u, 1u},
177 context: op.getContext()))
178 return rewriter.notifyMatchFailure(arg&: op, msg: "non-matching permutation maps");
179
180 // Check iterator types for matrix multiplication.
181 auto itTypes = op.getIteratorTypesArray();
182 if (itTypes.size() != 3 || itTypes[0] != vector::IteratorType::parallel ||
183 itTypes[1] != vector::IteratorType::parallel ||
184 itTypes[2] != vector::IteratorType::reduction)
185 return rewriter.notifyMatchFailure(
186 arg&: op, msg: "iterator types do not correspond to matrix multiplication");
187
188 // Check the combining kind is addition.
189 if (op.getKind() != vector::CombiningKind::ADD)
190 return rewriter.notifyMatchFailure(arg&: op,
191 msg: "combining kind is not an addition");
192
193 // Check the output is a vector of i32 elements.
194 auto outTy = dyn_cast<VectorType>(Val: op.getResultType());
195 if (!outTy || outTy.getElementType() != rewriter.getI32Type())
196 return rewriter.notifyMatchFailure(arg&: op,
197 msg: "output type is not a vector of i32");
198
199 // Check inputs are sign-/zero- extensions from i8 to i32. Get the values
200 // before the extension. All four signed/unsigned combinations for input
201 // operands are supported, but they are lowered to different operations.
202 // Determine which is the appropriate operation to lower to.
203 MMLA mmlaOp = MMLA::Signed;
204 auto maybeLhs = getExtOperand<arith::ExtSIOp>(v: op.getLhs());
205 if (!maybeLhs) {
206 mmlaOp = MMLA::Unsigned;
207 maybeLhs = getExtOperand<arith::ExtUIOp>(v: op.getLhs());
208 }
209 if (!maybeLhs)
210 return rewriter.notifyMatchFailure(
211 arg&: op, msg: "LHS is not a sign- or zero- extended i8");
212
213 auto maybeRhs = getExtOperand<arith::ExtSIOp>(v: op.getRhs());
214 if (maybeRhs) {
215 if (mmlaOp == MMLA::Unsigned)
216 mmlaOp = MMLA::Mixed;
217 } else {
218 if (mmlaOp == MMLA::Signed)
219 mmlaOp = MMLA::MixedSwapped;
220 maybeRhs = getExtOperand<arith::ExtUIOp>(v: op.getRhs());
221 }
222 if (!maybeRhs)
223 return rewriter.notifyMatchFailure(
224 arg&: op, msg: "RHS is not a sign- or zero- extended i8");
225
226 // One-dimensional vector types for arm_sve.*mmla
227 auto nxv16i8 = VectorType::get(/*shape=*/16, elementType: rewriter.getI8Type(),
228 /*scalableDims=*/{true});
229 auto nxv4i32 = VectorType::get(/*shape=*/4, elementType: rewriter.getI32Type(),
230 /*scalableDims=*/{true});
231
232 // Extract LHS sub-tiles with logicall shape <2x8>.
233 SmallVector<Value> lhsTile;
234 for (int64_t i = 0; i < M; i += 2) {
235 // Extract two consecutive rows of the LHS tile.
236 auto r0 = rewriter.create<vector::ExtractOp>(location: loc, args&: *maybeLhs,
237 args: ArrayRef<int64_t>{i});
238 auto r1 = rewriter.create<vector::ExtractOp>(location: loc, args&: *maybeLhs,
239 args: ArrayRef<int64_t>{i + 1});
240 // Concatenate to obtain a 16 x i8 flattened sub-tile.
241 auto t = rewriter.create<vector::ShuffleOp>(
242 location: loc, args&: r0, args&: r1,
243 args: llvm::ArrayRef<int64_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
244 14, 15});
245 // Turn it into a scalable vector.
246 auto s = rewriter.create<vector::ScalableInsertOp>(
247 location: loc, args&: t, args: rewriter.create<ub::PoisonOp>(location: loc, args&: nxv16i8), args: 0);
248 // Replicate the sub-tile VSCALE times to fill the entire vector.
249 auto r = rewriter.create<arm_sve::DupQLaneOp>(location: loc, args&: s, args: 0);
250 lhsTile.push_back(Elt: r);
251 }
252
253 // "Flatten" the RHS tile from <[N]x8> to <[8*N]>.
254 auto rhs = rewriter.create<vector::ShapeCastOp>(
255 location: maybeRhs->getLoc(),
256 args: VectorType::get(/*shape=*/8 * N, elementType: rewriter.getI8Type(),
257 /*scalableDims=*/{true}),
258 args&: *maybeRhs);
259
260 // Extract the RHS sub-tiles with logical shape <8x[2]>.
261 SmallVector<Value> rhsTile;
262 for (int64_t j = 0; j < N; j += 2)
263 rhsTile.push_back(
264 Elt: rewriter.create<vector::ScalableExtractOp>(location: loc, args&: nxv16i8, args&: rhs, args: j * 8));
265
266 // Handy types for packing/unpacking of the accumulator tile.
267 auto accRowTy = VectorType::get(/*shape=*/N, elementType: rewriter.getI32Type(),
268 /*scalableDims=*/{true});
269 auto accRowX2Ty = VectorType::get(/*shape=*/2 * N, elementType: rewriter.getI32Type(),
270 /*scalableDims=*/{true});
271 auto accRow64Ty = VectorType::get(/*shape=*/N / 2, elementType: rewriter.getI64Type(),
272 /*scalableDims=*/{true});
273 auto accRowX264Ty = VectorType::get(/*shape=*/N, elementType: rewriter.getI64Type(),
274 /*scalableDims=*/{true});
275
276 // Extract and pack the ACC sub-tiles.
277 SmallVector<Value> accTile;
278 for (int64_t i = 0; i < M; i += 2) {
279 // Extract two consecutive rows of the accumulator tile.
280 auto r0 = rewriter.create<vector::ExtractOp>(location: loc, args: op.getAcc(),
281 args: ArrayRef<int64_t>{i});
282 auto r1 = rewriter.create<vector::ExtractOp>(location: loc, args: op.getAcc(),
283 args: ArrayRef<int64_t>{i + 1});
284 Value accTileVec;
285 if (mmlaOp == MMLA::MixedSwapped) {
286 // We need to swap the positions of the LHS and RHS (since we don't have
287 // a signed * unsigned operation), but then each individual 2x2 tile of
288 // the acumulator and (later) the result need to be transposed.
289 accTileVec = rewriter.create<vector::InterleaveOp>(location: loc, args&: r0, args&: r1);
290 } else {
291 // Bitcast them to 64-bit elements, so subsequent
292 // interleave/deinterleave work on pairs of 32-bit numbers.
293 auto r0I64 = rewriter.create<vector::BitCastOp>(location: loc, args&: accRow64Ty, args&: r0);
294 auto r1I64 = rewriter.create<vector::BitCastOp>(location: loc, args&: accRow64Ty, args&: r1);
295
296 // Interleave the rows, effectively flattening each 2x2 tile into 4
297 // consecutive elements.
298 auto intrI64 = rewriter.create<vector::InterleaveOp>(location: loc, args&: r0I64, args&: r1I64);
299
300 // Bitcast back to 32-bit elements.
301 accTileVec =
302 rewriter.create<vector::BitCastOp>(location: loc, args&: accRowX2Ty, args&: intrI64);
303 }
304 // Extract ACC sub-tiles.
305 for (int64_t j = 0; j < N; j += 2)
306 accTile.push_back(Elt: rewriter.create<vector::ScalableExtractOp>(
307 location: loc, args&: nxv4i32, args&: accTileVec, args: j * 2));
308 }
309
310 // Emit sub-tile matrix multiplications.
311 SmallVector<Value> outTile;
312 for (int64_t i = 0; i < M / 2; ++i)
313 for (int64_t j = 0; j < N / 2; ++j) {
314 Value mmla = createMMLA(rewriter, op: mmlaOp, loc, accType: nxv4i32,
315 acc: accTile[i * N / 2 + j], lhs: lhsTile[i], rhs: rhsTile[j]);
316 outTile.push_back(Elt: mmla);
317 }
318
319 // Unpack the OUT sub-tiles and insert into the result.
320 Value result = rewriter.create<ub::PoisonOp>(location: loc, args: op.getResultType());
321 for (int64_t i = 0; i < M / 2; ++i) {
322 // Collect a number of sub-tiles in a row.
323 Value row = rewriter.create<ub::PoisonOp>(location: loc, args&: accRowX2Ty);
324 for (int64_t j = 0; j < N / 2; ++j)
325 row = rewriter.create<vector::ScalableInsertOp>(
326 location: loc, args&: outTile[i * N / 2 + j], args&: row, args: j * 4);
327
328 // Unpack the row to obtain two rows of the output. If we have the out
329 // sub-tiles transposed we obtain two consecutive output rows by
330 // separating even and odd elements, i.e. a simple deinterleave.
331 // Otherwise, the interleave is by pairs.
332 Value out0, out1;
333 if (mmlaOp == MMLA::MixedSwapped) {
334 auto tmp = rewriter.create<vector::DeinterleaveOp>(location: loc, args&: row);
335 out0 = tmp.getRes1();
336 out1 = tmp.getRes2();
337 } else {
338 // Deinterleave by pairs.
339 auto row64 = rewriter.create<vector::BitCastOp>(location: loc, args&: accRowX264Ty, args&: row);
340 auto deintr64 = rewriter.create<vector::DeinterleaveOp>(location: loc, args&: row64);
341
342 // Bitcast back into 32-bit elements and insert into the result.
343 out0 = rewriter.create<vector::BitCastOp>(location: loc, args&: accRowTy,
344 args: deintr64.getRes1());
345 out1 = rewriter.create<vector::BitCastOp>(location: loc, args&: accRowTy,
346 args: deintr64.getRes2());
347 }
348 result = rewriter.create<vector::InsertOp>(location: loc, args&: out0, args&: result, args: i * 2);
349 result = rewriter.create<vector::InsertOp>(location: loc, args&: out1, args&: result, args: i * 2 + 1);
350 }
351
352 rewriter.replaceOp(op, newValues: result);
353 return success();
354 }
355};
356
357} // namespace
358
359void mlir::populateLowerContractionToSVEI8MMPatternPatterns(
360 RewritePatternSet &patterns) {
361 MLIRContext *context = patterns.getContext();
362 patterns.add<LowerContractionToSVEI8MMPattern>(arg&: context, /*benefit=*/args: 2);
363}
364

source code of mlir/lib/Dialect/ArmSVE/Transforms/LowerContractionToSVEI8MMPattern.cpp