1 //===- ArmSMEToSCF.cpp - Convert ArmSME to SCF dialect ----------*- 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 of ArmSME operations to SCF. 10 // 11 //===----------------------------------------------------------------------===// 12 #include "mlir/Conversion/ArmSMEToSCF/ArmSMEToSCF.h" 13 14 #include "mlir/Dialect/Arith/IR/Arith.h" 15 #include "mlir/Dialect/ArmSME/IR/ArmSME.h" 16 #include "mlir/Dialect/ArmSME/Utils/Utils.h" 17 #include "mlir/Dialect/SCF/IR/SCF.h" 18 #include "mlir/Pass/Pass.h" 19 #include "mlir/Transforms/DialectConversion.h" 20 21 namespace mlir { 22 #define GEN_PASS_DEF_CONVERTARMSMETOSCF 23 #include "mlir/Conversion/Passes.h.inc" 24 } // namespace mlir 25 26 using namespace mlir; 27 28 namespace { 29 /// Returns adjusted (1-D or 2-D) `indices` for a tile slice as follows: 30 /// rank 1: (indices[0] + (tileSliceIndex * tileSliceNumElts)) 31 /// rank 2: (indices[0] + tileSliceIndex, indices[1]) 32 SmallVector<Value, 2> getMemrefIndices(ValueRange indices, unsigned rank, 33 Value tileSliceIndex, 34 Value tileSliceNumElts, Location loc, 35 PatternRewriter &rewriter) { 36 assert((rank == 1 || rank == 2) && "memref has unexpected rank!"); 37 SmallVector<Value, 2> outIndices; 38 39 auto tileSliceOffset = tileSliceIndex; 40 if (rank == 1) 41 tileSliceOffset = 42 rewriter.create<arith::MulIOp>(loc, tileSliceOffset, tileSliceNumElts); 43 44 auto baseIndexPlusTileSliceOffset = 45 rewriter.create<arith::AddIOp>(loc, indices[0], tileSliceOffset); 46 outIndices.push_back(baseIndexPlusTileSliceOffset); 47 48 if (rank == 2) 49 outIndices.push_back(indices[1]); 50 51 return outIndices; 52 } 53 54 /// Creates an scf.for for the load/store of an ArmSME tile. 55 FailureOr<scf::ForOp> createLoadStoreForOverTileSlices( 56 PatternRewriter &rewriter, Location loc, VectorType tileType, 57 ValueRange memrefIndices, int memrefRank, Value mask, Value initTile, 58 function_ref<Value(/*index=*/Value, ValueRange, /*predicate=*/Value, 59 /*currentTile=*/Value)> 60 makeLoopBody) { 61 PatternRewriter::InsertionGuard guard(rewriter); 62 63 auto minTileSlices = rewriter.create<arith::ConstantIndexOp>( 64 loc, arm_sme::getSMETileSliceMinNumElts(tileType.getElementType())); 65 auto vscale = 66 rewriter.create<vector::VectorScaleOp>(loc, rewriter.getIndexType()); 67 auto predicateType = 68 VectorType::get(tileType.getDimSize(1), rewriter.getI1Type(), true); 69 70 // This describes both the number of ZA tile slices and the number of 71 // elements in a vector of SVL bits for a given element type (SVL_B, 72 // SVL_H, ..., SVL_Q). 73 auto numTileSlices = 74 rewriter.create<arith::MulIOp>(loc, minTileSlices, vscale); 75 76 Value predicate; 77 Value upperBound; 78 if (mask) { 79 auto createMaskOp = mask.getDefiningOp<vector::CreateMaskOp>(); 80 if (!createMaskOp) 81 return rewriter.notifyMatchFailure( 82 loc, "unsupported mask op, only 'vector.create_mask' is " 83 "currently supported"); 84 85 auto maskDim0 = createMaskOp.getOperands()[0]; 86 auto maskDim1 = createMaskOp.getOperands()[1]; 87 88 // The upper bound of the loop must be clamped at `numTileSlices` as 89 // `vector.create_mask` allows operands to be greater than the size of a 90 // dimension. 91 auto numRowI64 = rewriter.create<arith::IndexCastOp>( 92 loc, rewriter.getI64Type(), maskDim0); 93 auto numTileSlicesI64 = rewriter.create<arith::IndexCastOp>( 94 loc, rewriter.getI64Type(), numTileSlices); 95 auto upperBoundI64 = 96 rewriter.create<arith::MinSIOp>(loc, numRowI64, numTileSlicesI64); 97 upperBound = rewriter.create<arith::IndexCastOp>( 98 loc, rewriter.getIndexType(), upperBoundI64); 99 100 predicate = 101 rewriter.create<vector::CreateMaskOp>(loc, predicateType, maskDim1); 102 } else { 103 upperBound = numTileSlices; 104 // No mask. Create an 'all true' predicate for the tile slice. 105 predicate = rewriter.create<arith::ConstantOp>( 106 loc, DenseElementsAttr::get(predicateType, true)); 107 } 108 109 bool hasCarriedArgs = bool(initTile); 110 auto lowerBound = rewriter.create<arith::ConstantIndexOp>(loc, 0); 111 auto step = rewriter.create<arith::ConstantIndexOp>(loc, 1); 112 auto forOp = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound, step, 113 hasCarriedArgs ? ValueRange{initTile} 114 : ValueRange{}); 115 116 rewriter.setInsertionPointToStart(forOp.getBody()); 117 Value tileSliceIndex = forOp.getInductionVar(); 118 119 auto adjustedIndices = getMemrefIndices( 120 memrefIndices, memrefRank, tileSliceIndex, numTileSlices, loc, rewriter); 121 auto nextTile = makeLoopBody( 122 tileSliceIndex, adjustedIndices, predicate, 123 /*currentTile=*/hasCarriedArgs ? forOp.getRegionIterArg(0) : Value{}); 124 125 assert(bool(nextTile) == hasCarriedArgs); 126 if (nextTile) 127 rewriter.create<scf::YieldOp>(loc, nextTile); 128 129 return forOp; 130 } 131 132 FailureOr<scf::ForOp> createLoadStoreForOverTileSlices( 133 PatternRewriter &rewriter, Location loc, VectorType tileType, 134 ValueRange memrefIndices, int memrefRank, Value mask, 135 function_ref<void(/*index=*/Value, ValueRange, /*predicate=*/Value)> 136 makeLoopBody) { 137 return createLoadStoreForOverTileSlices( 138 rewriter, loc, tileType, memrefIndices, memrefRank, mask, Value{}, 139 [&](Value index, ValueRange adjustedIndices, Value predicate, 140 Value) -> Value { 141 makeLoopBody(index, adjustedIndices, predicate); 142 return {}; 143 }); 144 } 145 146 /// Lower `arm_sme.tile_load` without a mask, or with a mask and a zero pad. 147 /// 148 /// With a mask: 149 /// 150 /// BEFORE: 151 /// ```mlir 152 /// %pad = arith.constant 0 : i32 153 /// %mask = vector.create_mask %num_rows, %num_cols : vector<[4]x[4]xi1> 154 /// %tile = arm_sme.tile_load %src[%c0, %c0], %pad, %mask : 155 /// memref<?x?xi32>, vector<[4]x[4]xi32> 156 /// ``` 157 /// 158 /// AFTER: 159 /// ```mlir 160 /// %init_tile = arm_sme.zero : vector<[4]x[4]xi32> 161 /// %mask_cols = vector.create_mask %num_cols : vector<[4]xi1> 162 /// %loop_rows = arith.minsi %num_rows, %svl_s : index 163 /// %tile = scf.for %tile_slice_idx = %c0 to %loop_rows step %c1 164 /// iter_args(%iter_tile = %init_tile) -> (vector<[4]x[4]xi32>) { 165 /// %tile_update = arm_sme.load_tile_slice 166 /// %src[%tile_slice_idx], %num_cols, %iter_tile, %tile_slice_idx : 167 /// memref<?x?xi32>, vector<[1]xi32>, vector<[4]x[4]xi32> 168 /// scf.yield %tile_update : vector<[4]x[4]xi32> 169 /// } 170 /// ``` 171 /// 172 /// Without a mask the lowering is pretty much identical. The only difference is 173 /// %mask_cols becomes an all-true mask, and %loop_rows becomes %svl_s. 174 /// 175 /// NOTE: Only mask of 'vector.create_mask' op is currently supported. 176 struct TileLoadOpConversion : public OpRewritePattern<arm_sme::TileLoadOp> { 177 using OpRewritePattern<arm_sme::TileLoadOp>::OpRewritePattern; 178 179 LogicalResult matchAndRewrite(arm_sme::TileLoadOp tileLoadOp, 180 PatternRewriter &rewriter) const override { 181 auto loc = tileLoadOp.getLoc(); 182 auto tileType = tileLoadOp.getVectorType(); 183 auto mask = tileLoadOp.getMask(); 184 185 Value initTile; 186 if (mask) { 187 auto padOp = tileLoadOp.getPadding(); 188 assert(padOp && "expected padding when masking!"); 189 190 auto constPadOp = padOp.getDefiningOp<arith::ConstantOp>(); 191 if (!constPadOp || constPadOp.getValue() != 192 rewriter.getZeroAttr(tileType.getElementType())) 193 return rewriter.notifyMatchFailure( 194 tileLoadOp, "op has non-zero pad, needs non-zero pad pattern"); 195 196 // Initialize tile with zero to satisfy padding. Inactive cols will be 197 // zeroed anyway since the loads use zeroing predication. For inactive 198 // rows however, no load will occur so these need to be zeroed. 199 initTile = rewriter.create<arm_sme::ZeroOp>(loc, tileType); 200 } else { 201 initTile = rewriter.create<arm_sme::GetTileOp>(loc, tileType); 202 } 203 204 // Create a loop to load the active tile slices from memory. 205 auto forOp = createLoadStoreForOverTileSlices( 206 rewriter, loc, tileType, tileLoadOp.getIndices(), 207 tileLoadOp.getMemRefType().getRank(), mask, initTile, 208 [&](Value tileSliceIndex, ValueRange memrefIndices, Value predicate, 209 Value currentTile) -> Value { 210 // Create 'arm_sme.load_tile_slice' to load tile slice from memory 211 // into tile. 212 return rewriter.create<arm_sme::LoadTileSliceOp>( 213 loc, tileType, tileLoadOp.getBase(), predicate, currentTile, 214 memrefIndices, tileSliceIndex, tileLoadOp.getLayout()); 215 }); 216 217 if (failed(forOp)) 218 return forOp; 219 220 // Replace 'arm_sme.tile_load' with the result. 221 rewriter.replaceOp(tileLoadOp, forOp->getResult(0)); 222 223 return success(); 224 } 225 }; 226 227 /// Lower `arm_sme.tile_load` with mask and non-zero pad. 228 /// 229 /// BEFORE: 230 /// ```mlir 231 /// %mask = vector.create_mask %num_rows, %num_cols : vector<[4]x[4]xi1> 232 /// %tile = arm_sme.tile_load %src[%c0, %c0], %pad, %mask : 233 /// memref<?x?xi32>, vector<[4]x[4]xi32> 234 /// ``` 235 /// 236 /// AFTER: 237 /// ```mlir 238 /// ... 239 /// %pad_1d = vector.splat %pad : vector<[4]xi32> 240 /// %tile = scf.for %tile_slice_idx = %c0 to %svl_s step %c1 241 /// iter_args(%iter_tile = %init_tile) -> (vector<[4]x[4]xi32>) { 242 /// ... 243 /// %mask_1d = vector.create_mask <combined_mask> : vector<[4]xi1> 244 /// %slice = vector.maskedload %base[%tile_slice_idx, %c0], %mask_1d, %pad_1d 245 /// : memref<?x?xi32>, vector<[4]xi1>, 246 /// vector<[4]xi32> into vector<[4]xi32> 247 /// // Insert slice into tile 248 /// %tile_update = arm_sme.insert_tile_slice 249 /// %slice, %iter_tile[%tile_slice_idx] : 250 /// vector<[4]xi32> into vector<[4]x[4]xi32> 251 /// scf.yield %tile_update : vector<[4]x[4]xi32> 252 /// } 253 /// ``` 254 struct TileLoadOpWithMaskAndPadNonZeroConversion 255 : public OpRewritePattern<arm_sme::TileLoadOp> { 256 using OpRewritePattern<arm_sme::TileLoadOp>::OpRewritePattern; 257 258 LogicalResult matchAndRewrite(arm_sme::TileLoadOp tileLoadOp, 259 PatternRewriter &rewriter) const override { 260 OpBuilder::InsertionGuard g(rewriter); 261 auto loc = tileLoadOp.getLoc(); 262 auto tileType = tileLoadOp.getVectorType(); 263 auto tileElementType = tileType.getElementType(); 264 265 auto maskOp = tileLoadOp.getMask(); 266 if (!maskOp) 267 return rewriter.notifyMatchFailure( 268 tileLoadOp, "op has no mask, needs unmasked pattern"); 269 270 auto padOp = tileLoadOp.getPadding(); 271 assert(padOp && "expected padding when masking!"); 272 273 auto createMaskOp = maskOp.getDefiningOp<vector::CreateMaskOp>(); 274 if (!createMaskOp) 275 return rewriter.notifyMatchFailure( 276 tileLoadOp, "unsupported mask op, only 'vector.create_mask' is " 277 "currently supported"); 278 279 auto constPadOp = padOp.getDefiningOp<arith::ConstantOp>(); 280 if (constPadOp && 281 constPadOp.getValue() == rewriter.getZeroAttr(tileElementType)) 282 return rewriter.notifyMatchFailure( 283 tileLoadOp, "op has constant zero pad, needs zero pad pattern"); 284 285 auto numRows = createMaskOp.getOperands()[0]; 286 auto numCols = createMaskOp.getOperands()[1]; 287 288 auto numColsI32 = rewriter.create<arith::IndexCastUIOp>( 289 loc, rewriter.getI32Type(), numCols); 290 291 auto initTile = rewriter.create<arm_sme::GetTileOp>(loc, tileType); 292 293 // Create a loop that loads each ZA tile slice from memory. 294 auto step = rewriter.create<arith::ConstantIndexOp>(loc, 1); 295 auto minTileSlices = rewriter.create<arith::ConstantIndexOp>( 296 loc, arm_sme::getSMETileSliceMinNumElts(tileElementType)); 297 auto vscale = 298 rewriter.create<vector::VectorScaleOp>(loc, rewriter.getIndexType()); 299 auto lowerBound = rewriter.create<arith::ConstantIndexOp>(loc, 0); 300 auto numTileSlices = 301 rewriter.create<arith::MulIOp>(loc, minTileSlices, vscale); 302 auto forOp = rewriter.create<scf::ForOp>(loc, lowerBound, numTileSlices, 303 step, ValueRange{initTile}); 304 305 rewriter.setInsertionPointToStart(forOp.getBody()); 306 307 auto tileSliceIndex = forOp.getInductionVar(); 308 auto currentTile = forOp.getRegionIterArg(0); 309 310 // Combine masks. 311 auto rowIsActive = rewriter.create<arith::CmpIOp>( 312 loc, arith::CmpIPredicate::ult, tileSliceIndex, numRows); 313 auto rowIsActiveI32 = rewriter.create<arith::ExtSIOp>( 314 loc, rewriter.getI32Type(), rowIsActive); 315 auto mask = rewriter.create<arith::AndIOp>(loc, rowIsActiveI32, numColsI32); 316 auto maskIndex = 317 rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), mask); 318 auto predicateType = 319 VectorType::get(tileType.getDimSize(1), rewriter.getI1Type(), true); 320 auto maskOp1D = rewriter.create<vector::CreateMaskOp>( 321 loc, predicateType, maskIndex.getResult()); 322 323 auto memrefIndices = getMemrefIndices( 324 tileLoadOp.getIndices(), tileLoadOp.getMemRefType().getRank(), 325 tileSliceIndex, numTileSlices, loc, rewriter); 326 327 // Splat pad into 1-D vector matching type of tile slice. 328 VectorType tileSliceType = VectorType::Builder(tileType).dropDim(0); 329 auto pad1DOp = rewriter.create<vector::SplatOp>(loc, tileSliceType, padOp); 330 331 auto loadSlice = rewriter.create<vector::MaskedLoadOp>( 332 loc, tileSliceType, tileLoadOp.getBase(), memrefIndices, maskOp1D, 333 /*passthru=*/pad1DOp); 334 335 // Create 'arm_sme.insert_tile_slice' to insert slice into tile. 336 auto insertSlice = rewriter.create<arm_sme::InsertTileSliceOp>( 337 loc, tileType, loadSlice->getResult(0), currentTile, tileSliceIndex, 338 tileLoadOp.getLayout()); 339 rewriter.create<scf::YieldOp>(loc, insertSlice.getResult()); 340 341 rewriter.setInsertionPointAfter(forOp); 342 343 // Replace 'arm_sme.tile_load' with the result. 344 rewriter.replaceOp(tileLoadOp, forOp.getResult(0)); 345 346 return success(); 347 } 348 }; 349 350 /// Lower `arm_sme.tile_store` to a loop over the tile slices and store each 351 /// slice using `arm_sme.store_tile_slice`. 352 /// 353 /// BEFORE: 354 /// ```mlir 355 /// arm_sme.tile_store %tile, %dest[%c0, %c0] layout<vertical> 356 /// : memref<?x?xi32>, vector<[4]x[4]xi32 357 /// ``` 358 /// 359 /// AFTER: 360 /// ```mlir 361 /// %vscale = vector.vscale 362 /// %c0 = arith.constant 0 : index 363 /// %c1 = arith.constant 1 : index 364 /// %min_svl_s = arith.constant 4 : index 365 /// %svl_s = arith.muli %min_svl_s, %vscale : index 366 /// scf.for %tile_slice_idx = %c0 to %svl_s step %c1 { 367 /// arm_sme.store_tile_slice %tile, %tile_slice_idx, %dest[%tile_slice_idx], 368 /// layout<vertical> : memref<?x?xi32>, vector<[4]x[4]xi32> 369 /// } 370 /// ``` 371 struct TileStoreOpConversion : public OpRewritePattern<arm_sme::TileStoreOp> { 372 using OpRewritePattern<arm_sme::TileStoreOp>::OpRewritePattern; 373 374 LogicalResult matchAndRewrite(arm_sme::TileStoreOp tileStoreOp, 375 PatternRewriter &rewriter) const override { 376 // Create a loop that stores each active ZA tile slice from memory. 377 return createLoadStoreForOverTileSlices( 378 rewriter, tileStoreOp.getLoc(), tileStoreOp.getVectorType(), 379 tileStoreOp.getIndices(), tileStoreOp.getMemRefType().getRank(), 380 tileStoreOp.getMask(), 381 [&](Value tileSliceIndex, ValueRange memrefIndices, Value predicate) { 382 rewriter.replaceOpWithNewOp<arm_sme::StoreTileSliceOp>( 383 tileStoreOp, tileStoreOp.getValueToStore(), tileSliceIndex, 384 predicate, tileStoreOp.getBase(), memrefIndices, 385 tileStoreOp.getLayout()); 386 }); 387 } 388 }; 389 390 } // namespace 391 392 void mlir::populateArmSMEToSCFConversionPatterns(RewritePatternSet &patterns) { 393 patterns.add<TileLoadOpConversion, TileLoadOpWithMaskAndPadNonZeroConversion, 394 TileStoreOpConversion>(patterns.getContext()); 395 } 396 397 namespace { 398 399 struct ConvertArmSMEToSCFPass 400 : public impl::ConvertArmSMEToSCFBase<ConvertArmSMEToSCFPass> { 401 void runOnOperation() override { 402 RewritePatternSet patterns(&getContext()); 403 ConversionTarget target(getContext()); 404 populateArmSMEToSCFConversionPatterns(patterns); 405 target.addLegalDialect<arm_sme::ArmSMEDialect, vector::VectorDialect, 406 arith::ArithDialect, scf::SCFDialect>(); 407 target.addIllegalOp<arm_sme::TileLoadOp, arm_sme::TileStoreOp>(); 408 if (failed(applyPartialConversion(getOperation(), target, 409 std::move(patterns)))) 410 signalPassFailure(); 411 } 412 }; 413 414 } // namespace 415 416 std::unique_ptr<Pass> mlir::createConvertArmSMEToSCFPass() { 417 return std::make_unique<ConvertArmSMEToSCFPass>(); 418 } 419