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

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