Skip to content

Commit ae74220

Browse files
KoolJBlackKoolJBlack
committed
implement vector unroll for i8mm
1 parent ab76052 commit ae74220

File tree

2 files changed

+130
-32
lines changed

2 files changed

+130
-32
lines changed

‎mlir/lib/Dialect/ArmNeon/Transforms/LowerContractionToSMMLAPattern.cpp

Lines changed: 80 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
#include "mlir/Dialect/ArmNeon/Transforms.h"
1717
#include "mlir/Dialect/Func/IR/FuncOps.h"
1818
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
19+
#include "mlir/Dialect/Utils/IndexingUtils.h"
1920
#include "mlir/Dialect/Vector/IR/VectorOps.h"
21+
#include "mlir/IR/AffineMap.h"
2022
#include "mlir/IR/PatternMatch.h"
2123
#include "mlir/Support/LogicalResult.h"
2224
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -36,19 +38,17 @@ static Type matchContainerType(Type element, Type container) {
3638
return element;
3739
}
3840

39-
/// Lowering from a single vector::contractOp directly to the arm neon smmla
40-
/// intrinsic. The shapes of the contract and intrinsic must match.
41+
/// Lowering from a vector::contractOp arm neon smmla intrinsic. This up to an
42+
/// 8x8x8 vector contract that is tiled (up to 16) smmla instructions with
43+
/// unrolling. If no unrolling is necessary, a single smmla instruction is
44+
/// emitted.
4145
class LowerContractionToSMMLAPattern
4246
: public OpRewritePattern<vector::ContractionOp> {
4347
public:
4448
using OpRewritePattern::OpRewritePattern;
4549
LogicalResult matchAndRewrite(vector::ContractionOp op,
4650
PatternRewriter &rewriter) const override {
4751
Location loc = op.getLoc();
48-
Value lhs = op.getLhs();
49-
Value rhs = op.getRhs();
50-
Value res = op.getAcc();
51-
5252
// Check index maps that represent M N K in contract.
5353
auto indexingMaps = op.getIndexingMapsArray();
5454
if (llvm::any_of(indexingMaps, [](mlir::AffineMap affineMap) {
@@ -57,7 +57,6 @@ class LowerContractionToSMMLAPattern
5757
})) {
5858
return failure();
5959
}
60-
6160
// Check iterator types for contract.
6261
auto iteratorTypes = op.getIteratorTypesArray();
6362
if (iteratorTypes.size() != 3 ||
@@ -66,22 +65,24 @@ class LowerContractionToSMMLAPattern
6665
iteratorTypes[2] != vector::IteratorType::reduction) {
6766
return failure();
6867
}
69-
70-
// Check the tile size by mapping the dimensions of the contract.
68+
// Infer tile sizes from operands; Note: RHS is not transposed.
7169
mlir::VectorType lhsType = op.getLhsType();
7270
mlir::VectorType rhsType = op.getRhsType();
7371
auto dimM = lhsType.getDimSize(0);
7472
auto dimN = rhsType.getDimSize(0);
7573
auto dimK = lhsType.getDimSize(1);
76-
if (rhsType.getDimSize(1) != dimK || dimM != 2 || dimN != 2 || dimK != 8) {
74+
75+
// Unrolling patterns can handle [(2|4|8), (2|4|8), 8] shaped inputs for
76+
// tiling.
77+
if (dimM % 2 != 0 || dimM > 8 || dimN % 2 != 0 || dimN > 8 || dimK != 8) {
7778
return failure();
7879
}
7980

8081
// Check two extsi inputs Rhs Lhs for contract.
8182
arith::ExtSIOp origLhsExtOp =
82-
dyn_cast_or_null<arith::ExtSIOp>(lhs.getDefiningOp());
83+
dyn_cast_or_null<arith::ExtSIOp>(op.getLhs().getDefiningOp());
8384
arith::ExtSIOp origRhsExtOp =
84-
dyn_cast_or_null<arith::ExtSIOp>(rhs.getDefiningOp());
85+
dyn_cast_or_null<arith::ExtSIOp>(op.getRhs().getDefiningOp());
8586
if (!origLhsExtOp || !origRhsExtOp) {
8687
return failure();
8788
}
@@ -113,26 +114,73 @@ class LowerContractionToSMMLAPattern
113114
return failure();
114115
}
115116

116-
// Collapse to 1D vectors required by smmla intrinsic
117-
auto collapsedInputType = VectorType::get(
118-
{16}, extsiLhs.getType().cast<ShapedType>().getElementType());
119-
auto collapsedOutputType =
120-
VectorType::get({4}, res.getType().cast<ShapedType>().getElementType());
121-
auto collapsedLhs = rewriter.createOrFold<vector::ShapeCastOp>(
122-
extsiLhs.getLoc(), collapsedInputType, extsiLhs);
123-
auto collapsedRhs = rewriter.createOrFold<vector::ShapeCastOp>(
124-
extsiRhs.getLoc(), collapsedInputType, extsiRhs);
125-
auto collapsedRes = rewriter.createOrFold<vector::ShapeCastOp>(
126-
res.getLoc(), collapsedOutputType, res);
127-
128-
// Replace the contract with a neon op
129-
auto smmlaOp = rewriter.createOrFold<arm_neon::SmmlaOp>(
130-
op.getLoc(), collapsedRes.getType(), collapsedRes, collapsedLhs,
131-
collapsedRhs);
132-
133-
// Reshape output back to 2D
134-
rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getResultType(),
135-
smmlaOp);
117+
// Initial accumulator for the final result. This is the un-tiled result if
118+
// tiling is done.
119+
Value result = rewriter.create<arith::ConstantOp>(
120+
loc, op.getResultType(), rewriter.getZeroAttr(op.getResultType()));
121+
122+
SmallVector<int64_t> unrolledSize = *op.getShapeForUnroll();
123+
SmallVector<int64_t> smmlaShape{2, 2, 8};
124+
SmallVector<int64_t> loopOrder{0, 1, 2};
125+
for (SmallVector<int64_t> offsets :
126+
StaticTileOffsetRange(unrolledSize, smmlaShape, loopOrder)) {
127+
128+
// Helper to compute the new shape of each operand and extract the slice.
129+
auto extractOperand = [&](Value operand, AffineMap permutationMap,
130+
ArrayRef<int64_t> operandOffsets) {
131+
SmallVector<int64_t> operandShape =
132+
applyPermutationMap(permutationMap, ArrayRef<int64_t>(smmlaShape));
133+
SmallVector<int64_t> operandStrides(operandOffsets.size(), 1);
134+
return rewriter.createOrFold<vector::ExtractStridedSliceOp>(
135+
loc, operand, operandOffsets, operandShape, operandStrides);
136+
};
137+
138+
// Extract tiled lhs, rhs, and acc
139+
AffineMap lhsPermutationMap = op.getIndexingMapsArray()[0];
140+
SmallVector<int64_t> lhsOffsets =
141+
applyPermutationMap(lhsPermutationMap, ArrayRef<int64_t>(offsets));
142+
auto tiledLhs = extractOperand(extsiLhs, lhsPermutationMap, lhsOffsets);
143+
AffineMap rhsPermutationMap = op.getIndexingMapsArray()[1];
144+
SmallVector<int64_t> rhsOffsets =
145+
applyPermutationMap(rhsPermutationMap, ArrayRef<int64_t>(offsets));
146+
auto tiledRhs = extractOperand(extsiRhs, rhsPermutationMap, rhsOffsets);
147+
AffineMap accPermutationMap = op.getIndexingMapsArray()[2];
148+
SmallVector<int64_t> accOffsets =
149+
applyPermutationMap(accPermutationMap, ArrayRef<int64_t>(offsets));
150+
auto tiledAcc =
151+
extractOperand(op.getAcc(), accPermutationMap, accOffsets);
152+
153+
// Collapse tiled operands to 1D vectors required by smmla intrinsic
154+
auto collapsedInputType = VectorType::get(
155+
tiledLhs.getType().cast<ShapedType>().getNumElements(),
156+
tiledLhs.getType().cast<ShapedType>().getElementType());
157+
auto collapsedOutputType = VectorType::get(
158+
{4}, tiledAcc.getType().cast<ShapedType>().getElementType());
159+
auto collapsedLhs = rewriter.createOrFold<vector::ShapeCastOp>(
160+
tiledLhs.getLoc(), collapsedInputType, tiledLhs);
161+
auto collapsedRhs = rewriter.createOrFold<vector::ShapeCastOp>(
162+
tiledRhs.getLoc(), collapsedInputType, tiledRhs);
163+
auto collapsedRes = rewriter.createOrFold<vector::ShapeCastOp>(
164+
tiledAcc.getLoc(), collapsedOutputType, tiledAcc);
165+
166+
// Insert contract op
167+
auto smmlaOp = rewriter.createOrFold<arm_neon::SmmlaOp>(
168+
op.getLoc(), collapsedRes.getType(), collapsedRes, collapsedLhs,
169+
collapsedRhs);
170+
171+
// Reshape output back to 2D
172+
Value tiledRes = rewriter.createOrFold<vector::ShapeCastOp>(
173+
smmlaOp.getLoc(), tiledAcc.getType(), smmlaOp);
174+
175+
// Insert the tiled result back into the non tiled result of the
176+
// contract op.
177+
SmallVector<int64_t> strides(
178+
tiledRes.getType().cast<ShapedType>().getRank(), 1);
179+
result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
180+
loc, tiledRes, result, accOffsets, strides);
181+
}
182+
183+
rewriter.replaceOp(op, result);
136184
return success();
137185
}
138186
};

‎mlir/test/Dialect/ArmNeon/lower-to-arm-neon.mlir

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,53 @@ func.func @test_lower_vector_arm_neon_without_extsi(%lhs: vector<2x8xi32>, %rhs:
4040
%res = vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %rhs, %acc : vector<2x8xi32>, vector<2x8xi32> into vector<2x2xi32>
4141
return %res : vector<2x2xi32>
4242
}
43+
44+
// -----
45+
46+
// CHECK-LABEL: test_lower_vector_arm_neon_unroll
47+
// CHECK-SAME: %[[VAL_0:.*]]: vector<4x8xi8>, %[[VAL_1:.*]]: vector<4x8xi8>, %[[VAL_2:.*]]: vector<4x4xi32>
48+
// CHECK: %[[VAL_3:.*]] = arith.constant dense<0> : vector<4x4xi32>
49+
// CHECK: %[[VAL_4:.*]] = vector.extract_strided_slice %[[VAL_0]] {offsets = [0, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
50+
// CHECK: %[[VAL_5:.*]] = vector.extract_strided_slice %[[VAL_1]] {offsets = [0, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
51+
// CHECK: %[[VAL_6:.*]] = vector.extract_strided_slice %[[VAL_2]] {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xi32> to vector<2x2xi32>
52+
// CHECK: %[[VAL_7:.*]] = vector.shape_cast %[[VAL_4]] : vector<2x8xi8> to vector<16xi8>
53+
// CHECK: %[[VAL_8:.*]] = vector.shape_cast %[[VAL_5]] : vector<2x8xi8> to vector<16xi8>
54+
// CHECK: %[[VAL_9:.*]] = vector.shape_cast %[[VAL_6]] : vector<2x2xi32> to vector<4xi32>
55+
// CHECK: %[[VAL_10:.*]] = arm_neon.intr.smmla %[[VAL_9]], %[[VAL_7]], %[[VAL_8]] : vector<16xi8> to vector<4xi32>
56+
// CHECK: %[[VAL_11:.*]] = vector.shape_cast %[[VAL_10]] : vector<4xi32> to vector<2x2xi32>
57+
// CHECK: %[[VAL_12:.*]] = vector.insert_strided_slice %[[VAL_11]], %[[VAL_3]] {offsets = [0, 0], strides = [1, 1]} : vector<2x2xi32> into vector<4x4xi32>
58+
// CHECK: %[[VAL_13:.*]] = vector.extract_strided_slice %[[VAL_0]] {offsets = [0, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
59+
// CHECK: %[[VAL_14:.*]] = vector.extract_strided_slice %[[VAL_1]] {offsets = [2, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
60+
// CHECK: %[[VAL_15:.*]] = vector.extract_strided_slice %[[VAL_2]] {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xi32> to vector<2x2xi32>
61+
// CHECK: %[[VAL_16:.*]] = vector.shape_cast %[[VAL_13]] : vector<2x8xi8> to vector<16xi8>
62+
// CHECK: %[[VAL_17:.*]] = vector.shape_cast %[[VAL_14]] : vector<2x8xi8> to vector<16xi8>
63+
// CHECK: %[[VAL_18:.*]] = vector.shape_cast %[[VAL_15]] : vector<2x2xi32> to vector<4xi32>
64+
// CHECK: %[[VAL_19:.*]] = arm_neon.intr.smmla %[[VAL_18]], %[[VAL_16]], %[[VAL_17]] : vector<16xi8> to vector<4xi32>
65+
// CHECK: %[[VAL_20:.*]] = vector.shape_cast %[[VAL_19]] : vector<4xi32> to vector<2x2xi32>
66+
// CHECK: %[[VAL_21:.*]] = vector.insert_strided_slice %[[VAL_20]], %[[VAL_12]] {offsets = [0, 2], strides = [1, 1]} : vector<2x2xi32> into vector<4x4xi32>
67+
// CHECK: %[[VAL_22:.*]] = vector.extract_strided_slice %[[VAL_0]] {offsets = [2, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
68+
// CHECK: %[[VAL_23:.*]] = vector.extract_strided_slice %[[VAL_1]] {offsets = [0, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
69+
// CHECK: %[[VAL_24:.*]] = vector.extract_strided_slice %[[VAL_2]] {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x4xi32> to vector<2x2xi32>
70+
// CHECK: %[[VAL_25:.*]] = vector.shape_cast %[[VAL_22]] : vector<2x8xi8> to vector<16xi8>
71+
// CHECK: %[[VAL_26:.*]] = vector.shape_cast %[[VAL_23]] : vector<2x8xi8> to vector<16xi8>
72+
// CHECK: %[[VAL_27:.*]] = vector.shape_cast %[[VAL_24]] : vector<2x2xi32> to vector<4xi32>
73+
// CHECK: %[[VAL_28:.*]] = arm_neon.intr.smmla %[[VAL_27]], %[[VAL_25]], %[[VAL_26]] : vector<16xi8> to vector<4xi32>
74+
// CHECK: %[[VAL_29:.*]] = vector.shape_cast %[[VAL_28]] : vector<4xi32> to vector<2x2xi32>
75+
// CHECK: %[[VAL_30:.*]] = vector.insert_strided_slice %[[VAL_29]], %[[VAL_21]] {offsets = [2, 0], strides = [1, 1]} : vector<2x2xi32> into vector<4x4xi32>
76+
// CHECK: %[[VAL_31:.*]] = vector.extract_strided_slice %[[VAL_0]] {offsets = [2, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
77+
// CHECK: %[[VAL_32:.*]] = vector.extract_strided_slice %[[VAL_1]] {offsets = [2, 0], sizes = [2, 8], strides = [1, 1]} : vector<4x8xi8> to vector<2x8xi8>
78+
// CHECK: %[[VAL_33:.*]] = vector.extract_strided_slice %[[VAL_2]] {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x4xi32> to vector<2x2xi32>
79+
// CHECK: %[[VAL_34:.*]] = vector.shape_cast %[[VAL_31]] : vector<2x8xi8> to vector<16xi8>
80+
// CHECK: %[[VAL_35:.*]] = vector.shape_cast %[[VAL_32]] : vector<2x8xi8> to vector<16xi8>
81+
// CHECK: %[[VAL_36:.*]] = vector.shape_cast %[[VAL_33]] : vector<2x2xi32> to vector<4xi32>
82+
// CHECK: %[[VAL_37:.*]] = arm_neon.intr.smmla %[[VAL_36]], %[[VAL_34]], %[[VAL_35]] : vector<16xi8> to vector<4xi32>
83+
// CHECK: %[[VAL_38:.*]] = vector.shape_cast %[[VAL_37]] : vector<4xi32> to vector<2x2xi32>
84+
// CHECK: %[[VAL_39:.*]] = vector.insert_strided_slice %[[VAL_38]], %[[VAL_30]] {offsets = [2, 2], strides = [1, 1]} : vector<2x2xi32> into vector<4x4xi32>
85+
// CHECK: return %[[VAL_39]] : vector<4x4xi32>
86+
// CHECK: }
87+
func.func @test_lower_vector_arm_neon_unroll(%lhs: vector<4x8xi8>, %rhs: vector<4x8xi8>, %acc : vector<4x4xi32>) -> vector<4x4xi32> {
88+
%lhs_extsi = arith.extsi %lhs : vector<4x8xi8> to vector<4x8xi32>
89+
%rhs_extsi = arith.extsi %rhs : vector<4x8xi8> to vector<4x8xi32>
90+
%res = vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs_extsi, %rhs_extsi, %acc : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>
91+
return %res : vector<4x4xi32>
92+
}

0 commit comments

Comments
 (0)