VectorToSCF.cpp Source File (original) (raw)

13 #include

14 #include

15 #include <type_traits>

34 namespace mlir {

35 #define GEN_PASS_DEF_CONVERTVECTORTOSCF

36 #include "mlir/Conversion/Passes.h.inc"

37 }

39 using namespace mlir;

40 using vector::TransferReadOp;

41 using vector::TransferWriteOp;

43 namespace {

46 static const char kPassLabel[] = "__vector_to_scf_lowering__";

49 static bool isTensorOp(VectorTransferOpInterface xferOp) {

50 if (isa(xferOp.getShapedType())) {

51 if (isavector::TransferWriteOp(xferOp)) {

53 assert(xferOp->getNumResults() > 0);

54 }

55 return true;

56 }

57 return false;

58 }

62 template

64 explicit VectorToSCFPattern(MLIRContext *context,

68 LogicalResult checkLowerTensors(VectorTransferOpInterface xferOp,

70 if (isTensorOp(xferOp) && options .lowerTensors) {

72 xferOp, "lowering tensor transfers is disabled");

73 }

74 return success();

75 }

78 };

83 template

84 static std::optional<int64_t> unpackedDim(OpTy xferOp) {

86 assert(xferOp.getTransferRank() > 0 && "unexpected 0-d transfer");

87 auto map = xferOp.getPermutationMap();

88 if (auto expr = dyn_cast(map.getResult(0))) {

89 return expr.getPosition();

90 }

91 assert(xferOp.isBroadcastDim(0) &&

92 "Expected AffineDimExpr or AffineConstantExpr");

93 return std::nullopt;

94 }

99 template

101

102 assert(xferOp.getTransferRank() > 0 && "unexpected 0-d transfer");

103 auto map = xferOp.getPermutationMap();

104 return AffineMap::get(map.getNumDims(), 0, map.getResults().drop_front(),

106 }

107

108

109

110

111

112

113

114 template

117 typename OpTy::Adaptor adaptor(xferOp);

118

119 auto dim = unpackedDim(xferOp);

120 auto prevIndices = adaptor.getIndices();

121 indices.append(prevIndices.begin(), prevIndices.end());

122

123 Location loc = xferOp.getLoc();

124 bool isBroadcast = !dim.has_value();

125 if (!isBroadcast) {

127 bindDims(xferOp.getContext(), d0, d1);

128 Value offset = adaptor.getIndices()[*dim];

129 indices[*dim] =

131 }

132 }

133

134 static void maybeYieldValue(OpBuilder &b, Location loc, bool hasRetVal,

136 if (hasRetVal) {

137 assert(value && "Expected non-empty value");

138 b.create scf::YieldOp(loc, value);

139 } else {

140 b.create scf::YieldOp(loc);

141 }

142 }

143

144

145

146

147

148

149

150 template

152 if (!xferOp.getMask())

154 if (xferOp.getMaskType().getRank() != 1)

156 if (xferOp.isBroadcastDim(0))

158

159 Location loc = xferOp.getLoc();

160 return b.create vector::ExtractElementOp(loc, xferOp.getMask(), iv);

161 }

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187 template

188 static Value generateInBoundsCheck(

189 OpBuilder &b, OpTy xferOp, Value iv, std::optional<int64_t> dim,

193 bool hasRetVal = !resultTypes.empty();

194 Value cond;

195

196

197 bool isBroadcast = !dim;

198 Location loc = xferOp.getLoc();

200 if (!xferOp.isDimInBounds(0) && !isBroadcast) {

203 bindDims(xferOp.getContext(), d0, d1);

204 Value base = xferOp.getIndices()[*dim];

205 Value memrefIdx =

207 cond = lb.createarith::CmpIOp(arith::CmpIPredicate::sgt, memrefDim,

208 memrefIdx);

209 }

210

211

212 if (auto maskCond = generateMaskCheck(b, xferOp, iv)) {

213 if (cond)

214 cond = lb.createarith::AndIOp(cond, maskCond);

215 else

216 cond = maskCond;

217 }

218

219

220 if (cond) {

221 auto check = lb.createscf::IfOp(

222 cond,

223

225 maybeYieldValue(b, loc, hasRetVal, inBoundsCase(b, loc));

226 },

227

229 if (outOfBoundsCase) {

230 maybeYieldValue(b, loc, hasRetVal, outOfBoundsCase(b, loc));

231 } else {

232 b.create scf::YieldOp(loc);

233 }

234 });

235

236 return hasRetVal ? check.getResult(0) : Value();

237 }

238

239

240 return inBoundsCase(b, loc);

241 }

242

243

244

245 template

246 static void generateInBoundsCheck(

247 OpBuilder &b, OpTy xferOp, Value iv, std::optional<int64_t> dim,

250 generateInBoundsCheck(

251 b, xferOp, iv, dim, TypeRange(),

252

254 inBoundsCase(b, loc);

256 },

257

259 if (outOfBoundsCase)

260 outOfBoundsCase(b, loc);

262 });

263 }

264

265

266 static ArrayAttr dropFirstElem(OpBuilder &b, ArrayAttr attr) {

267 if (!attr)

268 return attr;

270 }

271

272

273

274 template

275 static void maybeApplyPassLabel(OpBuilder &b, OpTy newXferOp,

276 unsigned targetRank) {

277 if (newXferOp.getVectorType().getRank() > targetRank)

278 newXferOp->setAttr(kPassLabel, b.getUnitAttr());

279 }

280

282

283

284 struct BufferAllocs {

285 Value dataBuffer;

286 Value maskBuffer;

287 };

288

289

293 assert(scope && "Expected op to be inside automatic allocation scope");

294 return scope;

295 }

296

297

298 template

299 static BufferAllocs allocBuffers(OpBuilder &b, OpTy xferOp) {

300 Location loc = xferOp.getLoc();

304 "AutomaticAllocationScope with >1 regions");

306

307 BufferAllocs result;

308 auto bufferType = MemRefType::get({}, xferOp.getVectorType());

309 result.dataBuffer = b.create memref::AllocaOp(loc, bufferType);

310

311 if (xferOp.getMask()) {

313 auto maskBuffer = b.create memref::AllocaOp(loc, maskType);

315 b.create memref::StoreOp(loc, xferOp.getMask(), maskBuffer);

316 result.maskBuffer = b.create memref::LoadOp(loc, maskBuffer, ValueRange());

317 }

318

319 return result;

320 }

321

322

323

324

325

326 static FailureOr unpackOneDim(MemRefType type) {

327 auto vectorType = dyn_cast(type.getElementType());

328

329

330 if (vectorType.getScalableDims().front())

331 return failure();

332 auto memrefShape = type.getShape();

334 newMemrefShape.append(memrefShape.begin(), memrefShape.end());

335 newMemrefShape.push_back(vectorType.getDimSize(0));

338 }

339

340

341

342 template

343 static Value getMaskBuffer(OpTy xferOp) {

344 assert(xferOp.getMask() && "Expected that transfer op has mask");

345 auto loadOp = xferOp.getMask().template getDefiningOpmemref::LoadOp();

346 assert(loadOp && "Expected transfer op mask produced by LoadOp");

347 return loadOp.getMemRef();

348 }

349

350

351 template

352 struct Strategy;

353

354

355 template <>

356 struct Strategy {

357

358

359 static memref::StoreOp getStoreOp(TransferReadOp xferOp) {

360 assert(xferOp->hasOneUse() && "Expected exactly one use of TransferReadOp");

361 auto storeOp = dyn_castmemref::StoreOp((*xferOp->use_begin()).getOwner());

362 assert(storeOp && "Expected TransferReadOp result used by StoreOp");

363 return storeOp;

364 }

365

366

367

368

369

370

371

372

374 return getStoreOp(xferOp).getMemRef();

375 }

376

377

378 static void getBufferIndices(TransferReadOp xferOp,

380 auto storeOp = getStoreOp(xferOp);

381 auto prevIndices = memref::StoreOpAdaptor(storeOp).getIndices();

382 indices.append(prevIndices.begin(), prevIndices.end());

383 }

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412 static TransferReadOp rewriteOp(OpBuilder &b,

414 TransferReadOp xferOp, Value buffer, Value iv,

417 getBufferIndices(xferOp, storeIndices);

418 storeIndices.push_back(iv);

419

422

423 Location loc = xferOp.getLoc();

424 auto bufferType = dyn_cast(buffer.getType());

425 auto vecType = dyn_cast(bufferType.getElementType());

426 auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());

427 auto newXferOp = b.create vector::TransferReadOp(

428 loc, vecType, xferOp.getBase(), xferIndices,

430 xferOp.getPadding(), Value(), inBoundsAttr);

431

432 maybeApplyPassLabel(b, newXferOp, options.targetRank);

433

434 b.create memref::StoreOp(loc, newXferOp.getVector(), buffer, storeIndices);

435 return newXferOp;

436 }

437

438

439

440 static Value handleOutOfBoundsDim(OpBuilder &b, TransferReadOp xferOp,

444 getBufferIndices(xferOp, storeIndices);

445 storeIndices.push_back(iv);

446

447 Location loc = xferOp.getLoc();

448 auto bufferType = dyn_cast(buffer.getType());

449 auto vecType = dyn_cast(bufferType.getElementType());

450 auto vec = b.create vector::SplatOp(loc, vecType, xferOp.getPadding());

451 b.create memref::StoreOp(loc, vec, buffer, storeIndices);

452

454 }

455

456

457 static void cleanup(PatternRewriter &rewriter, TransferReadOp xferOp,

458 scf::ForOp ) {

459 rewriter.eraseOp(getStoreOp(xferOp));

460 rewriter.eraseOp(xferOp);

461 }

462

463

464 static Value initialLoopState(TransferReadOp xferOp) { return Value(); }

465 };

466

467

468 template <>

469 struct Strategy {

470

471

472

473

474

475

476

478 auto loadOp = xferOp.getVector().getDefiningOpmemref::LoadOp();

479 assert(loadOp && "Expected transfer op vector produced by LoadOp");

480 return loadOp.getMemRef();

481 }

482

483

484 static void getBufferIndices(TransferWriteOp xferOp,

486 auto loadOp = xferOp.getVector().getDefiningOpmemref::LoadOp();

487 auto prevIndices = memref::LoadOpAdaptor(loadOp).getIndices();

488 indices.append(prevIndices.begin(), prevIndices.end());

489 }

490

491

492

493

494

495

496

497

498

499

500 static TransferWriteOp rewriteOp(OpBuilder &b,

502 TransferWriteOp xferOp, Value buffer,

505 getBufferIndices(xferOp, loadIndices);

506 loadIndices.push_back(iv);

507

510

511 Location loc = xferOp.getLoc();

512 auto vec = b.create memref::LoadOp(loc, buffer, loadIndices);

513 auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());

514 auto source = loopState.empty() ? xferOp.getBase() : loopState[0];

515 Type type = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();

516 auto newXferOp = b.create vector::TransferWriteOp(

517 loc, type, vec, source, xferIndices,

519 inBoundsAttr);

520

521 maybeApplyPassLabel(b, newXferOp, options.targetRank);

522

523 return newXferOp;

524 }

525

526

527 static Value handleOutOfBoundsDim(OpBuilder &b, TransferWriteOp xferOp,

530 return isTensorOp(xferOp) ? loopState[0] : Value();

531 }

532

533

534 static void cleanup(PatternRewriter &rewriter, TransferWriteOp xferOp,

535 scf::ForOp forOp) {

536 if (isTensorOp(xferOp)) {

537 assert(forOp->getNumResults() == 1 && "Expected one for loop result");

538 rewriter.replaceOp(xferOp, forOp->getResult(0));

539 } else {

540 rewriter.eraseOp(xferOp);

541 }

542 }

543

544

545 static Value initialLoopState(TransferWriteOp xferOp) {

546 return isTensorOp(xferOp) ? xferOp.getBase() : Value();

547 }

548 };

549

550 template

551 static LogicalResult checkPrepareXferOp(OpTy xferOp, PatternRewriter &rewriter,

553 if (xferOp->hasAttr(kPassLabel))

555 xferOp, "kPassLabel is present (vector-to-scf lowering in progress)");

556 if (xferOp.getVectorType().getRank() <= options.targetRank)

558 xferOp, "xferOp vector rank <= transformation target rank");

559 if (xferOp.getVectorType().getScalableDims().front())

561 xferOp, "Unpacking of the leading dimension into the memref is not yet "

562 "supported for scalable dims");

563 if (isTensorOp(xferOp) && options .lowerTensors)

565 xferOp, "Unpacking for tensors has been disabled.");

566 if (xferOp.getVectorType().getElementType() !=

567 xferOp.getShapedType().getElementType())

569 xferOp, "Mismatching source and destination element types.");

570

571 return success();

572 }

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597 struct PrepareTransferReadConversion

598 : public VectorToSCFPattern {

599 using VectorToSCFPattern::VectorToSCFPattern;

600

601 LogicalResult matchAndRewrite(TransferReadOp xferOp,

603 if (checkPrepareXferOp(xferOp, rewriter, options).failed())

605 xferOp, "checkPrepareXferOp conditions not met!");

606

607 auto buffers = allocBuffers(rewriter, xferOp);

608 auto *newXfer = rewriter.clone(*xferOp.getOperation());

610 if (xferOp.getMask()) {

611 dyn_cast(newXfer).getMaskMutable().assign(

612 buffers.maskBuffer);

613 }

614

615 Location loc = xferOp.getLoc();

616 rewriter.create memref::StoreOp(loc, newXfer->getResult(0),

617 buffers.dataBuffer);

618 rewriter.replaceOpWithNewOp memref::LoadOp(xferOp, buffers.dataBuffer);

619

620 return success();

621 }

622 };

623

624

625

626

627

628

629

630

631

632

633

634

635

636

637

638

639

640

641

642

643

644

645

646

647 struct PrepareTransferWriteConversion

648 : public VectorToSCFPattern {

649 using VectorToSCFPattern::VectorToSCFPattern;

650

651 LogicalResult matchAndRewrite(TransferWriteOp xferOp,

653 if (checkPrepareXferOp(xferOp, rewriter, options).failed())

655 xferOp, "checkPrepareXferOp conditions not met!");

656

657 Location loc = xferOp.getLoc();

658 auto buffers = allocBuffers(rewriter, xferOp);

659 rewriter.create memref::StoreOp(loc, xferOp.getVector(),

660 buffers.dataBuffer);

661 auto loadedVec = rewriter.create memref::LoadOp(loc, buffers.dataBuffer);

663 xferOp.getValueToStoreMutable().assign(loadedVec);

664 xferOp->setAttr(kPassLabel, rewriter.getUnitAttr());

665 });

666

667 if (xferOp.getMask()) {

669 xferOp.getMaskMutable().assign(buffers.maskBuffer);

670 });

671 }

672

673 return success();

674 }

675 };

676

677

678

679

680

681

682

683

684

685

686

687

688

689

690

691

692

693

694

695

696

697

698

699

700

701

702

703

704 struct DecomposePrintOpConversion : public VectorToSCFPatternvector::PrintOp {

705 using VectorToSCFPatternvector::PrintOp::VectorToSCFPattern;

706 LogicalResult matchAndRewrite(vector::PrintOp printOp,

708 if ( printOp .getSource())

709 return failure();

710

711 VectorType vectorType = dyn_cast(printOp.getPrintType());

712 if (!vectorType)

713 return failure();

714

715

716

717

718

719

720

721 if (vectorType.getRank() > 1 && vectorType.isScalable())

722 return failure();

723

724 auto loc = printOp.getLoc();

725 auto value = printOp.getSource();

726

727 if (auto intTy = dyn_cast(vectorType.getElementType())) {

728

729

730

731 auto width = intTy.getWidth();

732 auto legalWidth = llvm::NextPowerOf2(std::max(8u, width) - 1);

734 intTy.getSignedness());

735

736 auto signlessSourceVectorType =

737 vectorType.cloneWith({}, getIntTypeWithSignlessSemantics(intTy));

738 auto signlessTargetVectorType =

739 vectorType.cloneWith({}, getIntTypeWithSignlessSemantics(legalIntTy));

740 auto targetVectorType = vectorType.cloneWith({}, legalIntTy);

741 value = rewriter.create vector::BitCastOp(loc, signlessSourceVectorType,

742 value);

743 if (value.getType() != signlessTargetVectorType) {

744 if (width == 1 || intTy.isUnsigned())

745 value = rewriter.create arith::ExtUIOp(loc, signlessTargetVectorType,

746 value);

747 else

748 value = rewriter.create arith::ExtSIOp(loc, signlessTargetVectorType,

749 value);

750 }

751 value = rewriter.create vector::BitCastOp(loc, targetVectorType, value);

752 vectorType = targetVectorType;

753 }

754

755 auto scalableDimensions = vectorType.getScalableDims();

756 auto shape = vectorType.getShape();

757 constexpr int64_t singletonShape[] = {1};

758 if (vectorType.getRank() == 0)

759 shape = singletonShape;

760

761 if (vectorType.getRank() != 1) {

762

763

764

765 auto flatLength = std::accumulate(shape.begin(), shape.end(), 1,

766 std::multiplies<int64_t>());

767 auto flatVectorType =

768 VectorType::get({flatLength}, vectorType.getElementType());

769 value = rewriter.create vector::ShapeCastOp(loc, flatVectorType, value);

770 }

771

772 vector::PrintOp firstClose;

774 for (unsigned d = 0; d < shape.size(); d++) {

775

776 Value lowerBound = rewriter.create arith::ConstantIndexOp(loc, 0);

777 Value upperBound = rewriter.create arith::ConstantIndexOp(loc, shape[d]);

778 Value step = rewriter.create arith::ConstantIndexOp(loc, 1);

779 if (!scalableDimensions.empty() && scalableDimensions[d]) {

780 auto vscale = rewriter.create vector::VectorScaleOp(

782 upperBound = rewriter.create arith::MulIOp(loc, upperBound, vscale);

783 }

784 auto lastIndex = rewriter.create arith::SubIOp(loc, upperBound, step);

785

786

787 rewriter.create vector::PrintOp(loc, vector::PrintPunctuation::Open);

788 auto loop =

789 rewriter.create scf::ForOp(loc, lowerBound, upperBound, step);

791 loc, vector::PrintPunctuation::Close);

792 if (!firstClose)

794

795 auto loopIdx = loop.getInductionVar();

796 loopIndices.push_back(loopIdx);

797

798

800 auto notLastIndex = rewriter.create arith::CmpIOp(

801 loc, arith::CmpIPredicate::ult, loopIdx, lastIndex);

802 rewriter.create scf::IfOp(loc, notLastIndex,

804 builder.createvector::PrintOp(

805 loc, vector::PrintPunctuation::Comma);

806 builder.createscf::YieldOp(loc);

807 });

808

810 }

811

812

813

815 auto currentStride = 1;

816 for (int d = shape.size() - 1; d >= 0; d--) {

817 auto stride = rewriter.create arith::ConstantIndexOp(loc, currentStride);

818 auto index = rewriter.create arith::MulIOp(loc, stride, loopIndices[d]);

819 if (flatIndex)

820 flatIndex = rewriter.create arith::AddIOp(loc, flatIndex, index);

821 else

822 flatIndex = index;

823 currentStride *= shape[d];

824 }

825

826

827 auto element =

828 rewriter.create vector::ExtractElementOp(loc, value, flatIndex);

829 rewriter.create vector::PrintOp(loc, element,

830 vector::PrintPunctuation::NoPunctuation);

831

833 rewriter.create vector::PrintOp(loc, printOp.getPunctuation());

835 return success();

836 }

837

838 static IntegerType getIntTypeWithSignlessSemantics(IntegerType intTy) {

839 return IntegerType::get(intTy.getContext(), intTy.getWidth(),

840 IntegerType::Signless);

841 };

842 };

843

844

845

846

847

848

849

850

851

852

853

854

855

856

857

858

859

860

861

862

863

864

865

866

867

868

869

870

871

872

873 template

874 struct TransferOpConversion : public VectorToSCFPattern {

875 using VectorToSCFPattern::VectorToSCFPattern;

876

877 void initialize() {

878

879

880 this->setHasBoundedRewriteRecursion();

881 }

882

883 static void getMaskBufferLoadIndices(OpTy xferOp, Value castedMaskBuffer,

886 assert(xferOp.getMask() && "Expected transfer op to have mask");

887

888

889

890

891

892 Value maskBuffer = getMaskBuffer(xferOp);

894

895 if (auto loadOp = dyn_castmemref::LoadOp(user)) {

897 loadIndices.append(prevIndices.begin(), prevIndices.end());

898 break;

899 }

900 }

901

902

903

904 if (!xferOp.isBroadcastDim(0))

905 loadIndices.push_back(iv);

906 }

907

908 LogicalResult matchAndRewrite(OpTy xferOp,

910 if (!xferOp->hasAttr(kPassLabel))

912 xferOp, "kPassLabel is present (progressing lowering in progress)");

913

914

917 auto dataBufferType = dyn_cast(dataBuffer.getType());

918 FailureOr castedDataType = unpackOneDim(dataBufferType);

919 if (failed(castedDataType))

921 "Failed to unpack one vector dim.");

922

923 auto castedDataBuffer =

924 locB.create vector::TypeCastOp(*castedDataType, dataBuffer);

925

926

927 Value castedMaskBuffer;

928 if (xferOp.getMask()) {

929 Value maskBuffer = getMaskBuffer(xferOp);

930 if (xferOp.isBroadcastDim(0) || xferOp.getMaskType().getRank() == 1) {

931

932

933

934

935

936 castedMaskBuffer = maskBuffer;

937 } else {

938

939

940 auto maskBufferType = cast(maskBuffer.getType());

941 MemRefType castedMaskType = *unpackOneDim(maskBufferType);

942 castedMaskBuffer =

943 locB.create vector::TypeCastOp(castedMaskType, maskBuffer);

944 }

945 }

946

947

948 auto lb = locB.create arith::ConstantIndexOp(0);

949 auto ub = locB.create arith::ConstantIndexOp(

950 castedDataType->getDimSize(castedDataType->getRank() - 1));

951 auto step = locB.create arith::ConstantIndexOp(1);

952

953

954 auto loopState = Strategy::initialLoopState(xferOp);

955

956

957 auto result = locB.create scf::ForOp(

960 Type stateType = loopState.empty() ? Type() : loopState[0].getType();

961

962 auto result = generateInBoundsCheck(

963 b, xferOp, iv, unpackedDim(xferOp),

965

967

968 OpTy newXfer = Strategy::rewriteOp(

969 b, this->options, xferOp, castedDataBuffer, iv, loopState);

970

971

972

973

974

975 if (xferOp.getMask() && (xferOp.isBroadcastDim(0) ||

976 xferOp.getMaskType().getRank() > 1)) {

979

981 getMaskBufferLoadIndices(xferOp, castedMaskBuffer,

982 loadIndices, iv);

983 auto mask = b.create memref::LoadOp(loc, castedMaskBuffer,

984 loadIndices);

986 newXfer.getMaskMutable().assign(mask);

987 });

988 }

989

990 return loopState.empty() ? Value() : newXfer->getResult(0);

991 },

992

994 return Strategy::handleOutOfBoundsDim(

995 b, xferOp, castedDataBuffer, iv, loopState);

996 });

997

998 maybeYieldValue(b, loc, !loopState.empty(), result);

999 });

1000

1001 Strategy::cleanup(rewriter, xferOp, result);

1002 return success();

1003 }

1004 };

1005

1006

1007

1008 template

1009 static FailureOr<SmallVector>

1010 getMaskDimSizes(Value mask, VscaleConstantBuilder &createVscaleMultiple) {

1011 if (!mask)

1013 if (auto createMaskOp = mask.getDefiningOp vector::CreateMaskOp()) {

1014 return llvm::map_to_vector(createMaskOp.getOperands(), [](Value dimSize) {

1015 return OpFoldResult(dimSize);

1016 });

1017 }

1018 if (auto constantMask = mask.getDefiningOp vector::ConstantMaskOp()) {

1019 int dimIdx = 0;

1020 VectorType maskType = constantMask.getVectorType();

1022 return llvm::map_to_vector(

1023 constantMask.getMaskDimSizes(), [&](int64_t dimSize) {

1024

1025 if (maskType.getScalableDims()[dimIdx++])

1026 return OpFoldResult(createVscaleMultiple(dimSize));

1027 return OpFoldResult(IntegerAttr::get(indexType, dimSize));

1028 });

1029 }

1030 return failure();

1031 }

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

1047

1048

1049

1050

1051

1052

1053

1054

1055

1056

1057

1058

1059

1060

1061

1062

1063

1064

1065

1066

1067

1068

1069

1070 struct ScalableTransposeTransferWriteConversion

1071 : VectorToSCFPatternvector::TransferWriteOp {

1072 using VectorToSCFPattern::VectorToSCFPattern;

1073

1074 LogicalResult matchAndRewrite(TransferWriteOp writeOp,

1076 if (failed(checkLowerTensors(writeOp, rewriter)))

1077 return failure();

1078

1079 VectorType vectorType = writeOp.getVectorType();

1080

1081

1082

1083 ArrayRef scalableFlags = vectorType.getScalableDims();

1086 writeOp, "expected vector of the form vector<[N]xMxty>");

1087 }

1088

1089 auto permutationMap = writeOp.getPermutationMap();

1090 if (!permutationMap.isIdentity()) {

1092 writeOp, "non-identity permutations are unsupported (lower first)");

1093 }

1094

1095

1096

1097

1098 if (!writeOp.isDimInBounds(0)) {

1100 writeOp, "out-of-bounds dims are unsupported (use masking)");

1101 }

1102

1103 Value vector = writeOp.getVector();

1104 auto transposeOp = vector.getDefiningOp vector::TransposeOp();

1105 if (!transposeOp ||

1107 return rewriter.notifyMatchFailure(writeOp, "source not transpose");

1108 }

1109

1110 auto loc = writeOp.getLoc();

1111 auto createVscaleMultiple =

1113

1114 auto maskDims = getMaskDimSizes(writeOp.getMask(), createVscaleMultiple);

1115 if (failed(maskDims)) {

1117 "failed to resolve mask dims");

1118 }

1119

1120 int64_t fixedDimSize = vectorType.getDimSize(1);

1121 auto fixedDimOffsets = llvm::seq(fixedDimSize);

1122

1123

1124 auto transposeSource = transposeOp.getVector();

1126 llvm::map_to_vector(fixedDimOffsets, [&](int64_t idx) -> Value {

1127 return rewriter.create vector::ExtractOp(loc, transposeSource, idx);

1128 });

1129

1130

1131 auto lb = rewriter.create arith::ConstantIndexOp(loc, 0);

1132 auto ub =

1133 maskDims->empty()

1134 ? Value(createVscaleMultiple(vectorType.getDimSize(0)))

1136 auto step = rewriter.create arith::ConstantIndexOp(loc, 1);

1137

1138

1140 Value sliceMask = nullptr;

1141 if (!maskDims->empty()) {

1142 sliceMask = rewriter.create vector::CreateMaskOp(

1143 loc, sliceType.clone(rewriter.getI1Type()),

1145 }

1146

1147 Value initDest = isTensorOp(writeOp) ? writeOp.getBase() : Value{};

1149 auto result = rewriter.create scf::ForOp(

1150 loc, lb, ub, step, initLoopArgs,

1152

1155

1156

1158 llvm::map_to_vector(fixedDimOffsets, [&](int64_t idx) -> Value {

1159 return b.createvector::ExtractOp(

1160 loc, transposeSourceSlices[idx], iv);

1161 });

1162 auto sliceVec = b.createvector::FromElementsOp(loc, sliceType,

1163 transposeElements);

1164

1165

1167 loopIterArgs.empty() ? writeOp.getBase() : loopIterArgs.front();

1168 auto newWriteOp = b.createvector::TransferWriteOp(

1169 loc, sliceVec, dest, xferIndices,

1170 ArrayRef(writeOp.getInBoundsValues()).drop_front());

1171 if (sliceMask)

1172 newWriteOp.getMaskMutable().assign(sliceMask);

1173

1174

1175 b.createscf::YieldOp(loc, loopIterArgs.empty()

1177 : newWriteOp.getResult());

1178 });

1179

1180 if (isTensorOp(writeOp))

1181 rewriter.replaceOp(writeOp, result);

1182 else

1183 rewriter.eraseOp(writeOp);

1184

1185 return success();

1186 }

1187 };

1188

1189 }

1190

1192

1193

1194

1195 template

1196 static void maybeAssignMask(OpBuilder &b, OpTy xferOp, OpTy newXferOp,

1197 int64_t i) {

1198 if (!xferOp.getMask())

1199 return;

1200

1201 if (xferOp.isBroadcastDim(0)) {

1202

1203

1204 newXferOp.getMaskMutable().assign(xferOp.getMask());

1205 return;

1206 }

1207

1208 if (xferOp.getMaskType().getRank() > 1) {

1209

1212

1214 Location loc = xferOp.getLoc();

1215 auto newMask = b.create vector::ExtractOp(loc, xferOp.getMask(), indices);

1216 newXferOp.getMaskMutable().assign(newMask);

1217 }

1218

1219

1220

1221

1222 }

1223

1224

1225

1226

1227

1228

1229

1230

1231

1232

1233

1234

1235

1236

1237

1238

1239

1240

1241

1242

1243

1244

1245

1246

1247

1248

1249

1250

1251

1252 struct UnrollTransferReadConversion

1253 : public VectorToSCFPattern {

1254 using VectorToSCFPattern::VectorToSCFPattern;

1255

1256 void initialize() {

1257

1258

1259 setHasBoundedRewriteRecursion();

1260 }

1261

1262

1263

1265 TransferReadOp xferOp) const {

1266 if (auto insertOp = getInsertOp(xferOp))

1267 return insertOp.getDest();

1268 Location loc = xferOp.getLoc();

1269 return rewriter.create vector::SplatOp(loc, xferOp.getVectorType(),

1270 xferOp.getPadding());

1271 }

1272

1273

1274

1275 vector::InsertOp getInsertOp(TransferReadOp xferOp) const {

1276 if (xferOp->hasOneUse()) {

1278 if (auto insertOp = dyn_castvector::InsertOp(xferOpUser))

1279 return insertOp;

1280 }

1281

1282 return vector::InsertOp();

1283 }

1284

1285

1286

1287 void getInsertionIndices(TransferReadOp xferOp,

1289 if (auto insertOp = getInsertOp(xferOp)) {

1290 auto pos = insertOp.getMixedPosition();

1291 indices.append(pos.begin(), pos.end());

1292 }

1293 }

1294

1295

1296

1297 LogicalResult matchAndRewrite(TransferReadOp xferOp,

1299 if (xferOp.getVectorType().getRank() <= options.targetRank)

1301 xferOp, "vector rank is less or equal to target rank");

1302 if (failed(checkLowerTensors(xferOp, rewriter)))

1303 return failure();

1304 if (xferOp.getVectorType().getElementType() !=

1305 xferOp.getShapedType().getElementType())

1307 xferOp, "not yet supported: element type mismatch");

1308 auto xferVecType = xferOp.getVectorType();

1309 if (xferVecType.getScalableDims()[0]) {

1311 xferOp, "scalable dimensions cannot be unrolled at compile time");

1312 }

1313

1314 auto insertOp = getInsertOp(xferOp);

1315 auto vec = buildResultVector(rewriter, xferOp);

1316 auto vecType = dyn_cast(vec.getType());

1317

1319

1320 int64_t dimSize = xferVecType.getShape()[0];

1321

1322

1323 Location loc = xferOp.getLoc();

1324 for (int64_t i = 0; i < dimSize; ++i) {

1325 Value iv = rewriter.create arith::ConstantIndexOp(loc, i);

1326

1327 vec = generateInBoundsCheck(

1328 rewriter, xferOp, iv, unpackedDim(xferOp), TypeRange(vecType),

1329

1331

1334

1335

1337 getInsertionIndices(xferOp, insertionIndices);

1338 insertionIndices.push_back(rewriter.getIndexAttr(i));

1339

1340 auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());

1341 auto newXferOp = b.create vector::TransferReadOp(

1342 loc, newXferVecType, xferOp.getBase(), xferIndices,

1344 xferOp.getPadding(), Value(), inBoundsAttr);

1345 maybeAssignMask(b, xferOp, newXferOp, i);

1346 return b.create vector::InsertOp(loc, newXferOp, vec,

1347 insertionIndices);

1348 },

1349

1351

1352 return vec;

1353 });

1354 }

1355

1356 if (insertOp) {

1357

1358 rewriter.replaceOp(insertOp, vec);

1359 rewriter.eraseOp(xferOp);

1360 } else {

1361 rewriter.replaceOp(xferOp, vec);

1362 }

1363

1364 return success();

1365 }

1366 };

1367

1368

1369

1370

1371

1372

1373

1374

1375

1376

1377

1378

1379

1380

1381

1382

1383

1384

1385

1386

1387

1388

1389

1390

1391

1392

1393

1394 struct UnrollTransferWriteConversion

1395 : public VectorToSCFPattern {

1396 using VectorToSCFPattern::VectorToSCFPattern;

1397

1398 void initialize() {

1399

1400

1401 setHasBoundedRewriteRecursion();

1402 }

1403

1404

1405 Value getDataVector(TransferWriteOp xferOp) const {

1406 if (auto extractOp = getExtractOp(xferOp))

1407 return extractOp.getVector();

1408 return xferOp.getVector();

1409 }

1410

1411

1412 vector::ExtractOp getExtractOp(TransferWriteOp xferOp) const {

1413 if (auto *op = xferOp.getVector().getDefiningOp())

1414 return dyn_castvector::ExtractOp(op);

1415 return vector::ExtractOp();

1416 }

1417

1418

1419

1420 void getExtractionIndices(TransferWriteOp xferOp,

1422 if (auto extractOp = getExtractOp(xferOp)) {

1423 auto pos = extractOp.getMixedPosition();

1424 indices.append(pos.begin(), pos.end());

1425 }

1426 }

1427

1428

1429

1430 LogicalResult matchAndRewrite(TransferWriteOp xferOp,

1432 VectorType inputVectorTy = xferOp.getVectorType();

1433

1434 if (inputVectorTy.getRank() <= options.targetRank)

1435 return failure();

1436

1437 if (failed(checkLowerTensors(xferOp, rewriter)))

1438 return failure();

1439

1440 if (inputVectorTy.getElementType() !=

1441 xferOp.getShapedType().getElementType())

1442 return failure();

1443

1444 auto vec = getDataVector(xferOp);

1445 if (inputVectorTy.getScalableDims()[0]) {

1446

1447 return failure();

1448 }

1449

1450 int64_t dimSize = inputVectorTy.getShape()[0];

1451 Value source = xferOp.getBase();

1452 auto sourceType = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();

1453

1454

1455 Location loc = xferOp.getLoc();

1456 for (int64_t i = 0; i < dimSize; ++i) {

1457 Value iv = rewriter.create arith::ConstantIndexOp(loc, i);

1458

1459 auto updatedSource = generateInBoundsCheck(

1460 rewriter, xferOp, iv, unpackedDim(xferOp),

1462

1464

1467

1468

1470 getExtractionIndices(xferOp, extractionIndices);

1472

1473 auto extracted =

1474 b.create vector::ExtractOp(loc, vec, extractionIndices);

1475 auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());

1477 if (inputVectorTy.getRank() == 1) {

1478

1479

1480

1481 xferVec = b.create vector::BroadcastOp(

1482 loc, VectorType::get({}, extracted.getType()), extracted);

1483 } else {

1484 xferVec = extracted;

1485 }

1486 auto newXferOp = b.create vector::TransferWriteOp(

1487 loc, sourceType, xferVec, source, xferIndices,

1489 inBoundsAttr);

1490

1491 maybeAssignMask(b, xferOp, newXferOp, i);

1492

1493 return isTensorOp(xferOp) ? newXferOp->getResult(0) : Value();

1494 },

1495

1497 return isTensorOp(xferOp) ? source : Value();

1498 });

1499

1500 if (isTensorOp(xferOp))

1501 source = updatedSource;

1502 }

1503

1504 if (isTensorOp(xferOp))

1505 rewriter.replaceOp(xferOp, source);

1506 else

1507 rewriter.eraseOp(xferOp);

1508

1509 return success();

1510 }

1511 };

1512

1513 }

1514

1516

1517

1518

1519

1520

1521 template

1522 static std::optional<int64_t>

1523 get1dMemrefIndices(OpBuilder &b, OpTy xferOp, Value iv,

1525 auto indices = xferOp.getIndices();

1526 auto map = xferOp.getPermutationMap();

1527 assert(xferOp.getTransferRank() > 0 && "unexpected 0-d transfer");

1528

1529 memrefIndices.append(indices.begin(), indices.end());

1530 assert(map.getNumResults() == 1 &&

1531 "Expected 1 permutation map result for 1D transfer");

1532 if (auto expr = dyn_cast(map.getResult(0))) {

1533 Location loc = xferOp.getLoc();

1534 auto dim = expr.getPosition();

1536 bindDims(xferOp.getContext(), d0, d1);

1537 Value offset = memrefIndices[dim];

1538 memrefIndices[dim] =

1540 return dim;

1541 }

1542

1543 assert(xferOp.isBroadcastDim(0) &&

1544 "Expected AffineDimExpr or AffineConstantExpr");

1545 return std::nullopt;

1546 }

1547

1548

1549

1550 template

1551 struct Strategy1d;

1552

1553

1554 template <>

1555 struct Strategy1d {

1557 TransferReadOp xferOp, Value iv,

1560 auto dim = get1dMemrefIndices(b, xferOp, iv, indices);

1561 auto vec = loopState[0];

1562

1563

1564

1565 auto nextVec = generateInBoundsCheck(

1566 b, xferOp, iv, dim, TypeRange(xferOp.getVectorType()),

1567

1569 Value val = b.creatememref::LoadOp(loc, xferOp.getBase(), indices);

1570 return b.createvector::InsertElementOp(loc, val, vec, iv);

1571 },

1572

1574 b.create scf::YieldOp(loc, nextVec);

1575 }

1576

1577 static Value initialLoopState(OpBuilder &b, TransferReadOp xferOp) {

1578

1579 Location loc = xferOp.getLoc();

1580 return b.create vector::SplatOp(loc, xferOp.getVectorType(),

1581 xferOp.getPadding());

1582 }

1583 };

1584

1585

1586 template <>

1587 struct Strategy1d {

1589 TransferWriteOp xferOp, Value iv,

1592 auto dim = get1dMemrefIndices(b, xferOp, iv, indices);

1593

1594

1595 generateInBoundsCheck(

1596 b, xferOp, iv, dim,

1598 auto val =

1599 b.create vector::ExtractElementOp(loc, xferOp.getVector(), iv);

1600 b.create memref::StoreOp(loc, val, xferOp.getBase(), indices);

1601 });

1602 b.create scf::YieldOp(loc);

1603 }

1604

1605 static Value initialLoopState(OpBuilder &b, TransferWriteOp xferOp) {

1607 }

1608 };

1609

1610

1611

1612

1613

1614

1615

1616

1617

1618

1619

1620

1621

1622

1623

1624

1625

1626

1627

1628

1629

1630

1631

1632

1633

1634

1635

1636

1637

1638

1639

1640

1641 template

1642 struct TransferOp1dConversion : public VectorToSCFPattern {

1643 using VectorToSCFPattern::VectorToSCFPattern;

1644

1645 LogicalResult matchAndRewrite(OpTy xferOp,

1647

1648 if (xferOp.getTransferRank() == 0)

1649 return failure();

1650 auto map = xferOp.getPermutationMap();

1651 auto memRefType = dyn_cast(xferOp.getShapedType());

1652

1653 if (!memRefType)

1654 return failure();

1655 if (xferOp.getVectorType().getRank() != 1)

1656 return failure();

1657 if (map.isMinorIdentity() && memRefType.isLastDimUnitStride())

1658 return failure();

1659

1660

1661 Location loc = xferOp.getLoc();

1662 auto vecType = xferOp.getVectorType();

1663 auto lb = rewriter.create arith::ConstantIndexOp(loc, 0);

1665 rewriter.create arith::ConstantIndexOp(loc, vecType.getDimSize(0));

1666 if (vecType.isScalable()) {

1668 rewriter.create vector::VectorScaleOp(loc, rewriter.getIndexType());

1669 ub = rewriter.create arith::MulIOp(loc, ub, vscale);

1670 }

1671 auto step = rewriter.create arith::ConstantIndexOp(loc, 1);

1672 auto loopState = Strategy1d::initialLoopState(rewriter, xferOp);

1673

1674

1678 Strategy1d::generateForLoopBody(b, loc, xferOp, iv, loopState);

1679 });

1680

1681 return success();

1682 }

1683 };

1684

1685 }

1686 }

1687

1691 patterns.add<lowering_n_d_unrolled::UnrollTransferReadConversion,

1692 lowering_n_d_unrolled::UnrollTransferWriteConversion>(

1694 } else {

1695 patterns.add<lowering_n_d::PrepareTransferReadConversion,

1696 lowering_n_d::PrepareTransferWriteConversion,

1697 lowering_n_d::TransferOpConversion,

1698 lowering_n_d::TransferOpConversion>(

1700 }

1701 if (options.lowerScalable) {

1702 patterns.add<lowering_n_d::ScalableTransposeTransferWriteConversion>(

1704 }

1705 if (options.targetRank == 1) {

1706 patterns.add<lowering_1_d::TransferOp1dConversion,

1707 lowering_1_d::TransferOp1dConversion>(

1709 }

1710 patterns.add<lowering_n_d::DecomposePrintOpConversion>(patterns.getContext(),

1712 }

1713

1714 namespace {

1715

1716 struct ConvertVectorToSCFPass

1717 : public impl::ConvertVectorToSCFBase {

1718 ConvertVectorToSCFPass() = default;

1720 this->fullUnroll = options.unroll;

1721 this->targetRank = options.targetRank;

1722 this->lowerTensors = options.lowerTensors;

1723 this->lowerScalable = options.lowerScalable;

1724 }

1725

1726 void runOnOperation() override {

1728 options.unroll = fullUnroll;

1729 options.targetRank = targetRank;

1730 options.lowerTensors = lowerTensors;

1731 options.lowerScalable = lowerScalable;

1732

1733

1736 lowerTransferPatterns);

1738 std::move(lowerTransferPatterns));

1739

1743 }

1744 };

1745

1746 }

1747

1748 std::unique_ptr

1750 return std::make_unique(options);

1751 }

MLIR_CRUNNERUTILS_EXPORT void printClose()

static MLIRContext * getContext(OpFoldResult val)

static llvm::ManagedStatic< PassManagerOptions > options

static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)

static void printOp(llvm::raw_ostream &os, Operation *op, OpPrintingFlags &flags)

static void getXferIndices(RewriterBase &rewriter, TransferOpType xferOp, AffineMap offsetMap, ArrayRef< Value > dimValues, SmallVector< Value, 4 > &indices)

For a vector TransferOpType xferOp, an empty indices vector, and an AffineMap representing offsets to...

static Operation * getAutomaticAllocationScope(Operation *op)

Base type for affine expression.

A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.

static AffineMap get(MLIRContext *context)

Returns a zero result affine map with no dimensions or symbols: () -> ().

IntegerAttr getIndexAttr(int64_t value)

IntegerAttr getI64IntegerAttr(int64_t value)

MLIRContext * getContext() const

ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...

OpTy create(Args &&...args)

Create an operation of specific op type at the current insertion point and location.

This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...

MLIRContext is the top-level object for a collection of MLIR operations.

RAII guard to reset the insertion point of the builder when destroyed.

This class helps build Operations.

Operation * clone(Operation &op, IRMapping &mapper)

Creates a deep copy of the specified operation, remapping any operands that use values outside of the...

void setInsertionPointToStart(Block *block)

Sets the insertion point to the start of the specified block.

void setInsertionPoint(Block *block, Block::iterator insertPoint)

Set the insertion point to the specified location.

Operation * create(const OperationState &state)

Creates an operation given the fields represented as an OperationState.

void setInsertionPointAfter(Operation *op)

Sets the insertion point to the node after the specified operation, which will cause subsequent inser...

A trait of region holding operations that define a new scope for automatic allocations,...

This class implements the operand iterators for the Operation class.

Operation is the basic unit of execution within MLIR.

unsigned getNumRegions()

Returns the number of regions held by this operation.

Region & getRegion(unsigned index)

Returns the region held by this operation at position 'index'.

void setAttr(StringAttr name, Attribute value)

If the an attribute exists with the specified name, change it to the new value.

Operation * getParentWithTrait()

Returns the closest surrounding parent operation with trait Trait.

user_range getUsers()

Returns a range of all users.

A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...

std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)

Used to notify the listener that the IR failed to be rewritten because of a match failure,...

virtual void replaceOp(Operation *op, ValueRange newValues)

Replace the results of the given (original) operation with the specified list of values (replacements...

virtual void eraseOp(Operation *op)

This method erases an operation that is known to have no uses.

void modifyOpInPlace(Operation *root, CallableT &&callable)

This method is a utility wrapper around an in-place modification of an operation.

OpTy replaceOpWithNewOp(Operation *op, Args &&...args)

Replace the results of the given (original) op with a new op that is created without verification (re...

This class provides an abstraction over the various different ranges of value types.

Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...

This class provides an abstraction over the different types of ranges over Values.

This class represents an instance of an SSA value in the MLIR system, representing a computable value...

MLIRContext * getContext() const

Utility to get the associated MLIRContext that this value is defined in.

Type getType() const

Return the type of this value.

user_range getUsers() const

Operation * getDefiningOp() const

If this value is the result of an operation, return the operation that defines it.

This is a builder type that keeps local references to arguments.

Builder & dropDim(unsigned pos)

Erase a dim from shape @pos.

AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)

Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...

FailureOr< Value > getBuffer(RewriterBase &rewriter, Value value, const BufferizationOptions &options, const BufferizationState &state)

Lookup the buffer for the given value.

void populateVectorTransferPermutationMapLoweringPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)

Collect a set of transfer read/write lowering patterns that simplify the permutation map (e....

Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim)

Helper function that creates a memref::DimOp or tensor::DimOp depending on the type of source.

SmallVector< Value > getAsValues(OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > foldResults)

Convert foldResults into Values.

auto makeVscaleConstantBuilder(PatternRewriter &rewriter, Location loc)

Returns a functor (int64_t -> Value) which returns a constant vscale multiple.

Include the generated interface declarations.

Type getType(OpFoldResult ofr)

Returns the int type of the integer in ofr.

void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)

Bind a list of AffineExpr references to DimExpr at positions: [0 .

LogicalResult applyPatternsGreedily(Region &region, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)

Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...

const FrozenRewritePatternSet & patterns

void populateVectorToSCFConversionPatterns(RewritePatternSet &patterns, const VectorTransferToSCFOptions &options=VectorTransferToSCFOptions())

Collect a set of patterns to convert from the Vector dialect to SCF + func.

auto get(MLIRContext *context, Ts &&...params)

Helper method that injects context only if needed, this helps unify some of the attribute constructio...

std::unique_ptr< Pass > createConvertVectorToSCFPass(const VectorTransferToSCFOptions &options=VectorTransferToSCFOptions())

Create a pass to convert a subset of vector ops to SCF.

OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...

When lowering an N-d vector transfer op to an (N-1)-d vector transfer op, a temporary buffer is creat...