1088 lines · cpp
1//===- TestVectorTransforms.cpp - Test Vector transforms and lowerings ----===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include <optional>10 11#include "mlir/Analysis/SliceAnalysis.h"12#include "mlir/Dialect/Affine/IR/AffineOps.h"13#include "mlir/Dialect/Arith/IR/Arith.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/GPU/IR/GPUDialect.h"16#include "mlir/Dialect/Linalg/Passes.h"17#include "mlir/Dialect/MemRef/IR/MemRef.h"18#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"19#include "mlir/Dialect/SCF/IR/SCF.h"20#include "mlir/Dialect/SCF/Transforms/Patterns.h"21#include "mlir/Dialect/Tensor/IR/Tensor.h"22#include "mlir/Dialect/Vector/IR/VectorOps.h"23#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"24#include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"25#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"26#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"27#include "mlir/Pass/Pass.h"28#include "mlir/Pass/PassManager.h"29#include "mlir/Support/LLVM.h"30#include "mlir/Transforms/GreedyPatternRewriteDriver.h"31 32using namespace mlir;33using namespace mlir::linalg;34using namespace mlir::vector;35 36namespace {37 38struct TestVectorToVectorLowering39 : public PassWrapper<TestVectorToVectorLowering,40 OperationPass<func::FuncOp>> {41 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorToVectorLowering)42 43 TestVectorToVectorLowering() = default;44 TestVectorToVectorLowering(const TestVectorToVectorLowering &pass)45 : PassWrapper(pass) {}46 StringRef getArgument() const final {47 return "test-vector-to-vector-lowering";48 }49 StringRef getDescription() const final {50 return "Test lowering patterns between ops in the vector dialect";51 }52 53 void getDependentDialects(DialectRegistry ®istry) const override {54 registry.insert<affine::AffineDialect>();55 registry.insert<vector::VectorDialect>();56 }57 58 Option<bool> unroll{*this, "unroll", llvm::cl::desc("Include unrolling"),59 llvm::cl::init(false)};60 61 void runOnOperation() override {62 auto *ctx = &getContext();63 RewritePatternSet patterns(ctx);64 if (unroll) {65 populateVectorUnrollPatterns(66 patterns,67 UnrollVectorOptions().setNativeShapeFn(getShape).setFilterConstraint(68 filter));69 }70 populateVectorToVectorCanonicalizationPatterns(patterns);71 populateBubbleVectorBitCastOpPatterns(patterns);72 populateCastAwayVectorLeadingOneDimPatterns(patterns);73 (void)applyPatternsGreedily(getOperation(), std::move(patterns));74 }75 76private:77 // Return the target shape based on op type.78 static std::optional<SmallVector<int64_t>> getShape(Operation *op) {79 if (isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp>(op))80 return SmallVector<int64_t>(2, 2);81 if (isa<vector::ContractionOp>(op))82 return SmallVector<int64_t>(3, 2);83 // For transfer ops, just propagate the shape coming from84 // InsertStridedSlices/ExtractStridedSlices.85 if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {86 VectorType dstVec;87 for (Operation *users : readOp->getUsers()) {88 auto extract = dyn_cast<ExtractStridedSliceOp>(users);89 if (!extract)90 return std::nullopt;91 auto vecType = cast<VectorType>(extract.getResult().getType());92 if (dstVec && dstVec != vecType)93 return std::nullopt;94 dstVec = vecType;95 }96 return SmallVector<int64_t>(dstVec.getShape());97 }98 if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {99 auto insert = writeOp.getVector().getDefiningOp<InsertStridedSliceOp>();100 if (!insert)101 return std::nullopt;102 ArrayRef<int64_t> shape = insert.getSourceVectorType().getShape();103 return SmallVector<int64_t>(shape);104 }105 return std::nullopt;106 }107 108 static LogicalResult filter(Operation *op) {109 return success(isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp,110 ContractionOp, TransferReadOp, TransferWriteOp>(op));111 }112};113 114struct TestVectorContractionPrepareForMMTLowering115 : public PassWrapper<TestVectorContractionPrepareForMMTLowering,116 OperationPass<func::FuncOp>> {117 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(118 TestVectorContractionPrepareForMMTLowering)119 120 StringRef getArgument() const final {121 return "test-vector-contraction-prepare-for-mmt-lowering";122 }123 StringRef getDescription() const final {124 return "Test vector.contraction matmul canonicalization for MMT lowering.";125 }126 TestVectorContractionPrepareForMMTLowering() = default;127 128 void getDependentDialects(DialectRegistry ®istry) const override {129 registry.insert<affine::AffineDialect, arith::ArithDialect,130 vector::VectorDialect>();131 }132 133 void runOnOperation() override {134 MLIRContext *ctx = &getContext();135 RewritePatternSet patterns(ctx);136 vector::populateVectorContractCanonicalizeMatmulToMMT(patterns);137 (void)applyPatternsGreedily(getOperation(), std::move(patterns));138 }139};140 141struct TestVectorUnrollingPatterns142 : public PassWrapper<TestVectorUnrollingPatterns,143 OperationPass<func::FuncOp>> {144 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorUnrollingPatterns)145 146 StringRef getArgument() const final {147 return "test-vector-unrolling-patterns";148 }149 StringRef getDescription() const final {150 return "Test lowering patterns to unroll contract ops in the vector "151 "dialect";152 }153 TestVectorUnrollingPatterns() = default;154 TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass)155 : PassWrapper(pass) {}156 void runOnOperation() override {157 MLIRContext *ctx = &getContext();158 RewritePatternSet patterns(ctx);159 populateVectorUnrollPatterns(160 patterns,161 UnrollVectorOptions()162 .setNativeShape(ArrayRef<int64_t>{2, 2})163 .setFilterConstraint([](Operation *op) {164 return success(165 isa<arith::AddFOp, vector::FMAOp, vector::MultiDimReductionOp,166 vector::BroadcastOp, vector::LoadOp, vector::StoreOp>(167 op));168 }));169 populateVectorUnrollPatterns(170 patterns, UnrollVectorOptions()171 .setNativeShape(ArrayRef<int64_t>{2})172 .setFilterConstraint([](Operation *op) {173 return success(isa<vector::ReductionOp>(op));174 }));175 populateVectorUnrollPatterns(patterns,176 UnrollVectorOptions()177 .setNativeShape(ArrayRef<int64_t>{8})178 .setFilterConstraint([](Operation *op) {179 return success(isa<vector::StepOp>(op));180 }));181 populateVectorUnrollPatterns(182 patterns,183 UnrollVectorOptions()184 .setNativeShapeFn(185 [](Operation *op) -> std::optional<SmallVector<int64_t>> {186 auto shapeCast = dyn_cast<vector::ShapeCastOp>(op);187 if (!shapeCast)188 return std::nullopt;189 190 auto resultShape = shapeCast.getResultVectorType().getShape();191 // Special case with leading unit dims and different inner dim192 // for result and target shape.193 if (resultShape.size() == 2 && resultShape[0] == 1 &&194 resultShape[1] == 32) {195 return SmallVector<int64_t>{1, 16};196 }197 // Default case: [2,4] for all tests.198 return SmallVector<int64_t>{2, 4};199 })200 .setFilterConstraint([](Operation *op) {201 return success(isa<vector::ShapeCastOp>(op));202 }));203 populateVectorUnrollPatterns(204 patterns, UnrollVectorOptions()205 .setNativeShape(ArrayRef<int64_t>{1, 3, 4, 2})206 .setFilterConstraint([](Operation *op) {207 return success(isa<vector::TransposeOp>(op));208 }));209 210 if (unrollBasedOnType) {211 UnrollVectorOptions::NativeShapeFnType nativeShapeFn =212 [](Operation *op) -> std::optional<SmallVector<int64_t>> {213 vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);214 SmallVector<int64_t> nativeShape(contractOp.getIteratorTypes().size(),215 4);216 Type lhsType = contractOp.getLhsType().getElementType();217 nativeShape[nativeShape.size() - 1] = lhsType.isF16() ? 4 : 2;218 return nativeShape;219 };220 221 UnrollVectorOptions opts;222 opts.setNativeShapeFn(nativeShapeFn)223 .setFilterConstraint(224 [](Operation *op) { return success(isa<ContractionOp>(op)); });225 226 if (!unrollOrder.empty()) {227 opts.setUnrollTraversalOrderFn(228 [this](Operation *op) -> std::optional<SmallVector<int64_t>> {229 vector::ContractionOp contractOp =230 cast<vector::ContractionOp>(op);231 if (contractOp.getIteratorTypes().size() == unrollOrder.size())232 return SmallVector<int64_t>(unrollOrder.begin(),233 unrollOrder.end());234 return std::nullopt;235 });236 }237 populateVectorUnrollPatterns(patterns, opts);238 } else {239 auto nativeShapeFn =240 [](Operation *op) -> std::optional<SmallVector<int64_t>> {241 auto contractOp = dyn_cast<ContractionOp>(op);242 if (!contractOp)243 return std::nullopt;244 return SmallVector<int64_t>(contractOp.getIteratorTypes().size(), 2);245 };246 populateVectorUnrollPatterns(patterns,247 UnrollVectorOptions()248 .setNativeShapeFn(nativeShapeFn)249 .setFilterConstraint([](Operation *op) {250 return success(isa<ContractionOp>(op));251 }));252 }253 populateVectorToVectorCanonicalizationPatterns(patterns);254 (void)applyPatternsGreedily(getOperation(), std::move(patterns));255 }256 257 ListOption<int64_t> unrollOrder{*this, "unroll-order",258 llvm::cl::desc("set the unroll order")};259 260 Option<bool> unrollBasedOnType{261 *this, "unroll-based-on-type",262 llvm::cl::desc("Set the unroll factor based on type of the operation"),263 llvm::cl::init(false)};264};265 266struct TestVectorTransferUnrollingPatterns267 : public PassWrapper<TestVectorTransferUnrollingPatterns,268 OperationPass<func::FuncOp>> {269 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(270 TestVectorTransferUnrollingPatterns)271 272 TestVectorTransferUnrollingPatterns() = default;273 TestVectorTransferUnrollingPatterns(274 const TestVectorTransferUnrollingPatterns &pass)275 : PassWrapper(pass) {}276 277 void getDependentDialects(DialectRegistry ®istry) const override {278 registry.insert<affine::AffineDialect>();279 }280 StringRef getArgument() const final {281 return "test-vector-transfer-unrolling-patterns";282 }283 StringRef getDescription() const final {284 return "Test lowering patterns to unroll transfer ops in the vector "285 "dialect";286 }287 void runOnOperation() override {288 MLIRContext *ctx = &getContext();289 RewritePatternSet patterns(ctx);290 UnrollVectorOptions opts;291 opts.setNativeShape(ArrayRef<int64_t>{2, 2})292 .setFilterConstraint([](Operation *op) {293 return success(isa<vector::TransferReadOp, vector::TransferWriteOp,294 vector::GatherOp>(op));295 });296 if (reverseUnrollOrder.getValue()) {297 opts.setUnrollTraversalOrderFn(298 [](Operation *op) -> std::optional<SmallVector<int64_t>> {299 int64_t numLoops = 0;300 if (auto readOp = dyn_cast<vector::TransferReadOp>(op))301 numLoops = readOp.getVectorType().getRank();302 else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op))303 numLoops = writeOp.getVectorType().getRank();304 else if (auto gatherOp = dyn_cast<vector::GatherOp>(op))305 numLoops = gatherOp.getVectorType().getRank();306 else307 return std::nullopt;308 auto order = llvm::reverse(llvm::seq<int64_t>(0, numLoops));309 return llvm::to_vector(order);310 });311 }312 populateVectorUnrollPatterns(patterns, opts);313 populateVectorToVectorCanonicalizationPatterns(patterns);314 (void)applyPatternsGreedily(getOperation(), std::move(patterns));315 }316 317 Option<bool> reverseUnrollOrder{318 *this, "reverse-unroll-order",319 llvm::cl::desc(320 "reverse the order of unrolling of vector transfer operations"),321 llvm::cl::init(false)};322};323 324struct TestScalarVectorTransferLoweringPatterns325 : public PassWrapper<TestScalarVectorTransferLoweringPatterns,326 OperationPass<func::FuncOp>> {327 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(328 TestScalarVectorTransferLoweringPatterns)329 330 TestScalarVectorTransferLoweringPatterns() = default;331 TestScalarVectorTransferLoweringPatterns(332 const TestScalarVectorTransferLoweringPatterns &pass)333 : PassWrapper(pass) {}334 335 StringRef getArgument() const final {336 return "test-scalar-vector-transfer-lowering";337 }338 StringRef getDescription() const final {339 return "Test lowering of scalar vector transfers to memref loads/stores.";340 }341 342 void getDependentDialects(DialectRegistry ®istry) const override {343 registry.insert<affine::AffineDialect, memref::MemRefDialect,344 tensor::TensorDialect, vector::VectorDialect>();345 }346 347 Option<bool> allowMultipleUses{348 *this, "allow-multiple-uses",349 llvm::cl::desc("Fold transfer operations with multiple uses"),350 llvm::cl::init(false)};351 352 void runOnOperation() override {353 MLIRContext *ctx = &getContext();354 RewritePatternSet patterns(ctx);355 vector::populateScalarVectorTransferLoweringPatterns(356 patterns, /*benefit=*/1, allowMultipleUses.getValue());357 (void)applyPatternsGreedily(getOperation(), std::move(patterns));358 }359};360 361struct TestVectorTransferOpt362 : public PassWrapper<TestVectorTransferOpt, OperationPass<func::FuncOp>> {363 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorTransferOpt)364 365 StringRef getArgument() const final { return "test-vector-transferop-opt"; }366 StringRef getDescription() const final {367 return "Test optimization transformations for transfer ops";368 }369 void runOnOperation() override {370 IRRewriter rewriter(&getContext());371 transferOpflowOpt(rewriter, getOperation());372 }373};374 375struct TestVectorSinkPatterns376 : public PassWrapper<TestVectorSinkPatterns, OperationPass<func::FuncOp>> {377 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorSinkPatterns)378 379 TestVectorSinkPatterns() = default;380 TestVectorSinkPatterns(const TestVectorSinkPatterns &pass) = default;381 382 void getDependentDialects(DialectRegistry ®istry) const override {383 registry.insert<memref::MemRefDialect, affine::AffineDialect>();384 }385 386 StringRef getArgument() const final { return "test-vector-sink-patterns"; }387 388 StringRef getDescription() const final {389 return "Test lowering patterns that eliminate redundant broadcast "390 "and transpose operations.";391 }392 393 void runOnOperation() override {394 RewritePatternSet patterns(&getContext());395 populateSinkVectorOpsPatterns(patterns);396 populateSinkVectorMemOpsPatterns(patterns);397 (void)applyPatternsGreedily(getOperation(), std::move(patterns));398 }399};400 401struct TestVectorReduceToContractPatternsPatterns402 : public PassWrapper<TestVectorReduceToContractPatternsPatterns,403 OperationPass<func::FuncOp>> {404 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(405 TestVectorReduceToContractPatternsPatterns)406 407 StringRef getArgument() const final {408 return "test-vector-reduction-to-contract-patterns";409 }410 StringRef getDescription() const final {411 return "Test patterns to convert multireduce op to contract and combine "412 "broadcast/transpose to contract";413 }414 void runOnOperation() override {415 RewritePatternSet patterns(&getContext());416 populateVectorReductionToContractPatterns(patterns);417 (void)applyPatternsGreedily(getOperation(), std::move(patterns));418 }419};420 421struct TestVectorChainedReductionFoldingPatterns422 : public PassWrapper<TestVectorChainedReductionFoldingPatterns,423 OperationPass<func::FuncOp>> {424 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(425 TestVectorChainedReductionFoldingPatterns)426 427 StringRef getArgument() const final {428 return "test-vector-chained-reduction-folding-patterns";429 }430 StringRef getDescription() const final {431 return "Test patterns to fold chained vector reductions";432 }433 void runOnOperation() override {434 RewritePatternSet patterns(&getContext());435 populateChainedVectorReductionFoldingPatterns(patterns);436 (void)applyPatternsGreedily(getOperation(), std::move(patterns));437 }438};439 440struct TestVectorBreakDownReductionPatterns441 : public PassWrapper<TestVectorBreakDownReductionPatterns,442 OperationPass<func::FuncOp>> {443 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(444 TestVectorBreakDownReductionPatterns)445 446 StringRef getArgument() const final {447 return "test-vector-break-down-reduction-patterns";448 }449 StringRef getDescription() const final {450 return "Test patterns to break down vector reductions into arith "451 "reductions";452 }453 void runOnOperation() override {454 RewritePatternSet patterns(&getContext());455 populateBreakDownVectorReductionPatterns(patterns,456 /*maxNumElementsToExtract=*/2);457 (void)applyPatternsGreedily(getOperation(), std::move(patterns));458 }459};460 461struct TestFlattenVectorTransferPatterns462 : public PassWrapper<TestFlattenVectorTransferPatterns,463 OperationPass<func::FuncOp>> {464 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(465 TestFlattenVectorTransferPatterns)466 467 TestFlattenVectorTransferPatterns() = default;468 TestFlattenVectorTransferPatterns(469 const TestFlattenVectorTransferPatterns &pass)470 : PassWrapper(pass) {}471 472 StringRef getArgument() const final {473 return "test-vector-transfer-flatten-patterns";474 }475 476 StringRef getDescription() const final {477 return "Test patterns to rewrite contiguous row-major N-dimensional "478 "vector.transfer_{read,write} ops into 1D transfers";479 }480 481 void getDependentDialects(DialectRegistry ®istry) const override {482 registry.insert<memref::MemRefDialect>();483 registry.insert<affine::AffineDialect>();484 registry.insert<vector::VectorDialect>();485 }486 487 Option<unsigned> targetVectorBitwidth{488 *this, "target-vector-bitwidth",489 llvm::cl::desc(490 "Minimum vector bitwidth to enable the flattening transformation. "491 "For scalable vectors this is the base size, i.e. the size "492 "corresponding to vscale=1."),493 llvm::cl::init(std::numeric_limits<unsigned>::max())};494 495 void runOnOperation() override {496 RewritePatternSet patterns(&getContext());497 populateFlattenVectorTransferPatterns(patterns, targetVectorBitwidth);498 (void)applyPatternsGreedily(getOperation(), std::move(patterns));499 }500};501 502struct TestVectorScanLowering503 : public PassWrapper<TestVectorScanLowering, OperationPass<func::FuncOp>> {504 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorScanLowering)505 506 StringRef getArgument() const final { return "test-vector-scan-lowering"; }507 StringRef getDescription() const final {508 return "Test lowering patterns that lower the scan op in the vector "509 "dialect";510 }511 void runOnOperation() override {512 RewritePatternSet patterns(&getContext());513 populateVectorScanLoweringPatterns(patterns);514 (void)applyPatternsGreedily(getOperation(), std::move(patterns));515 }516};517 518/// Allocate shared memory for a single warp to test lowering of519/// WarpExecuteOnLane0Op.520static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,521 gpu::WarpExecuteOnLane0Op warpOp,522 Type type) {523 static constexpr int64_t kSharedMemorySpace = 3;524 // Compute type of shared memory buffer.525 MemRefType memrefType;526 if (auto vectorType = dyn_cast<VectorType>(type)) {527 memrefType =528 MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},529 kSharedMemorySpace);530 } else {531 memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);532 }533 534 // Get symbol table holding all shared memory globals.535 ModuleOp moduleOp = warpOp->getParentOfType<ModuleOp>();536 SymbolTable symbolTable(moduleOp);537 538 // Create a pretty name.539 SmallString<64> buf;540 llvm::raw_svector_ostream os(buf);541 interleave(memrefType.getShape(), os, "x");542 os << "x" << memrefType.getElementType();543 std::string symbolName = (Twine("__shared_") + os.str()).str();544 545 auto ip = builder.saveInsertionPoint();546 builder.setInsertionPoint(moduleOp);547 auto global = memref::GlobalOp::create(548 builder, loc,549 /*sym_name=*/symbolName,550 /*sym_visibility=*/builder.getStringAttr("private"),551 /*type=*/memrefType,552 /*initial_value=*/Attribute(),553 /*constant=*/false,554 /*alignment=*/IntegerAttr());555 symbolTable.insert(global);556 // The symbol table inserts at the end of the module, but globals are a bit557 // nicer if they are at the beginning.558 global->moveBefore(&moduleOp.front());559 560 builder.restoreInsertionPoint(ip);561 return memref::GetGlobalOp::create(builder, loc, memrefType, symbolName);562}563 564static Value warpReduction(Location loc, OpBuilder &builder, Value input,565 CombiningKind kind, uint32_t size) {566 // First reduce on a single thread to get per lane reduction value.567 Value laneVal = vector::ReductionOp::create(builder, loc, kind, input);568 // Parallel reduction using butterfly shuffles.569 for (uint64_t i = 1; i < size; i <<= 1) {570 Value shuffled = gpu::ShuffleOp::create(builder, loc, laneVal, i,571 /*width=*/size,572 /*mode=*/gpu::ShuffleMode::XOR)573 .getShuffleResult();574 laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);575 }576 return laneVal;577}578 579struct TestVectorDistribution580 : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {581 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)582 583 void getDependentDialects(DialectRegistry ®istry) const override {584 registry585 .insert<vector::VectorDialect, scf::SCFDialect, memref::MemRefDialect,586 gpu::GPUDialect, affine::AffineDialect>();587 }588 589 StringRef getArgument() const final { return "test-vector-warp-distribute"; }590 StringRef getDescription() const final {591 return "Test vector warp distribute transformation and lowering patterns";592 }593 TestVectorDistribution() = default;594 TestVectorDistribution(const TestVectorDistribution &pass)595 : PassWrapper(pass) {}596 597 Option<bool> warpOpToSCF{598 *this, "rewrite-warp-ops-to-scf-if",599 llvm::cl::desc("Lower vector.warp_execute_on_lane0 to scf.if op"),600 llvm::cl::init(false)};601 602 Option<bool> distributeTransferWriteOps{603 *this, "distribute-transfer-write",604 llvm::cl::desc("Test distribution of transfer write"),605 llvm::cl::init(false)};606 607 Option<unsigned> maxTransferWriteElements{608 *this, "max-transfer-write-elements",609 llvm::cl::desc("Maximum number of transfer write elements to distribute"),610 llvm::cl::init(1)};611 612 Option<bool> hoistUniform{*this, "hoist-uniform",613 llvm::cl::desc("Test hoist uniform"),614 llvm::cl::init(false)};615 616 Option<bool> propagateDistribution{617 *this, "propagate-distribution",618 llvm::cl::desc("Test distribution propagation"), llvm::cl::init(false)};619 620 void runOnOperation() override {621 RewritePatternSet patterns(&getContext());622 623 getOperation().walk([&](Operation *op) {624 if (auto warpOp = dyn_cast<gpu::WarpExecuteOnLane0Op>(op)) {625 if (hoistUniform) {626 moveScalarUniformCode(warpOp);627 }628 WalkResult::interrupt();629 }630 });631 MLIRContext *ctx = &getContext();632 auto distributionFn = [](Value val) {633 // Create an identity dim map of the same rank as the vector.634 VectorType vecType = dyn_cast<VectorType>(val.getType());635 int64_t vecRank = vecType ? vecType.getRank() : 0;636 OpBuilder builder(val.getContext());637 if (vecRank == 0)638 return AffineMap::get(val.getContext());639 return AffineMap::getMultiDimIdentityMap(vecRank, val.getContext());640 };641 auto shuffleFn = [](Location loc, OpBuilder &builder, Value val,642 Value srcIdx, int64_t warpSz) {643 assert((val.getType().isF32() || val.getType().isInteger(32)) &&644 "unsupported shuffle type");645 Type i32Type = builder.getIntegerType(32);646 Value srcIdxI32 =647 arith::IndexCastOp::create(builder, loc, i32Type, srcIdx);648 Value warpSzI32 = arith::ConstantOp::create(649 builder, loc, builder.getIntegerAttr(i32Type, warpSz));650 Value result = gpu::ShuffleOp::create(builder, loc, val, srcIdxI32,651 warpSzI32, gpu::ShuffleMode::IDX)652 .getResult(0);653 return result;654 };655 if (distributeTransferWriteOps && propagateDistribution) {656 RewritePatternSet patterns(ctx);657 vector::populatePropagateWarpVectorDistributionPatterns(658 patterns, distributionFn, shuffleFn, /*benefit=*/1,659 /*readBenefit=*/0);660 vector::populateDistributeReduction(patterns, warpReduction, 1);661 populateDistributeTransferWriteOpPatterns(patterns, distributionFn, 2);662 (void)applyPatternsGreedily(getOperation(), std::move(patterns));663 } else if (distributeTransferWriteOps) {664 RewritePatternSet patterns(ctx);665 populateDistributeTransferWriteOpPatterns(patterns, distributionFn,666 maxTransferWriteElements);667 (void)applyPatternsGreedily(getOperation(), std::move(patterns));668 } else if (propagateDistribution) {669 RewritePatternSet patterns(ctx);670 vector::populatePropagateWarpVectorDistributionPatterns(671 patterns, distributionFn, shuffleFn);672 vector::populateDistributeReduction(patterns, warpReduction);673 (void)applyPatternsGreedily(getOperation(), std::move(patterns));674 }675 WarpExecuteOnLane0LoweringOptions options;676 options.warpAllocationFn = allocateGlobalSharedMemory;677 options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,678 gpu::WarpExecuteOnLane0Op warpOp) {679 gpu::BarrierOp::create(builder, loc);680 };681 // Test on one pattern in isolation.682 if (warpOpToSCF) {683 populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);684 (void)applyPatternsGreedily(getOperation(), std::move(patterns));685 return;686 }687 }688};689 690struct TestVectorExtractStridedSliceLowering691 : public PassWrapper<TestVectorExtractStridedSliceLowering,692 OperationPass<func::FuncOp>> {693 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(694 TestVectorExtractStridedSliceLowering)695 696 StringRef getArgument() const final {697 return "test-vector-extract-strided-slice-lowering";698 }699 StringRef getDescription() const final {700 return "Test lowering patterns that converts vector.extract_strided_slice "701 "into a chain of vector.extract and vector.insert ops";702 }703 void runOnOperation() override {704 RewritePatternSet patterns(&getContext());705 populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);706 (void)applyPatternsGreedily(getOperation(), std::move(patterns));707 }708};709 710struct TestVectorBreakDownBitCast711 : public PassWrapper<TestVectorBreakDownBitCast,712 OperationPass<func::FuncOp>> {713 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBreakDownBitCast)714 715 StringRef getArgument() const final {716 return "test-vector-break-down-bitcast";717 }718 StringRef getDescription() const final {719 return "Test pattern that breaks down vector.bitcast ops ";720 }721 void runOnOperation() override {722 RewritePatternSet patterns(&getContext());723 populateBreakDownVectorBitCastOpPatterns(patterns, [](BitCastOp op) {724 return op.getSourceVectorType().getShape().back() > 4;725 });726 (void)applyPatternsGreedily(getOperation(), std::move(patterns));727 }728};729 730struct TestCreateVectorBroadcast731 : public PassWrapper<TestCreateVectorBroadcast,732 OperationPass<func::FuncOp>> {733 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)734 735 StringRef getArgument() const final { return "test-create-vector-broadcast"; }736 StringRef getDescription() const final {737 return "Test optimization transformations for transfer ops";738 }739 void getDependentDialects(DialectRegistry ®istry) const override {740 registry.insert<vector::VectorDialect>();741 }742 743 void runOnOperation() override {744 getOperation()->walk([](Operation *op) {745 if (op->getName().getStringRef() != "test_create_broadcast")746 return;747 auto targetShape =748 cast<VectorType>(op->getResult(0).getType()).getShape();749 auto arrayAttr =750 cast<DenseI64ArrayAttr>(op->getDiscardableAttr("broadcast_dims"))751 .asArrayRef();752 llvm::SetVector<int64_t> broadcastedDims;753 broadcastedDims.insert_range(arrayAttr);754 OpBuilder b(op);755 Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(756 b, op->getOperand(0), targetShape, broadcastedDims);757 op->getResult(0).replaceAllUsesWith(bcast);758 op->erase();759 });760 }761};762 763struct TestVectorGatherLowering764 : public PassWrapper<TestVectorGatherLowering,765 OperationPass<func::FuncOp>> {766 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)767 768 StringRef getArgument() const final { return "test-vector-gather-lowering"; }769 StringRef getDescription() const final {770 return "Test patterns that lower the gather op in the vector conditional "771 "loads";772 }773 void getDependentDialects(DialectRegistry ®istry) const override {774 registry.insert<arith::ArithDialect, func::FuncDialect,775 memref::MemRefDialect, scf::SCFDialect,776 tensor::TensorDialect, vector::VectorDialect>();777 }778 779 void runOnOperation() override {780 RewritePatternSet patterns(&getContext());781 populateVectorGatherLoweringPatterns(patterns);782 populateVectorGatherToConditionalLoadPatterns(patterns);783 (void)applyPatternsGreedily(getOperation(), std::move(patterns));784 }785};786 787struct TestFoldArithExtensionIntoVectorContractPatterns788 : public PassWrapper<TestFoldArithExtensionIntoVectorContractPatterns,789 OperationPass<func::FuncOp>> {790 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(791 TestFoldArithExtensionIntoVectorContractPatterns)792 793 StringRef getArgument() const final {794 return "test-fold-arith-extf-into-vector-contract-patterns";795 }796 StringRef getDescription() const final {797 return "Test patterns that fold arithmetic extension ops into vector "798 "contract ops";799 }800 801 void getDependentDialects(DialectRegistry ®istry) const override {802 registry.insert<arith::ArithDialect, func::FuncDialect, nvgpu::NVGPUDialect,803 memref::MemRefDialect, scf::SCFDialect,804 tensor::TensorDialect, vector::VectorDialect>();805 }806 807 void runOnOperation() override {808 RewritePatternSet patterns(&getContext());809 populateFoldArithExtensionPatterns(patterns);810 (void)applyPatternsGreedily(getOperation(), std::move(patterns));811 }812};813 814struct TestVectorEmulateMaskedLoadStore final815 : public PassWrapper<TestVectorEmulateMaskedLoadStore,816 OperationPass<func::FuncOp>> {817 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorEmulateMaskedLoadStore)818 819 StringRef getArgument() const override {820 return "test-vector-emulate-masked-load-store";821 }822 StringRef getDescription() const override {823 return "Test patterns that emulate the maskedload/maskedstore op by "824 " memref.load/store and scf.if";825 }826 void getDependentDialects(DialectRegistry ®istry) const override {827 registry828 .insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,829 scf::SCFDialect, vector::VectorDialect>();830 }831 832 void runOnOperation() override {833 RewritePatternSet patterns(&getContext());834 populateVectorMaskedLoadStoreEmulationPatterns(patterns);835 (void)applyPatternsGreedily(getOperation(), std::move(patterns));836 }837};838 839/// Get the set of operand/result types to check for sufficiently840/// small inner-most dimension size.841static SmallVector<std::pair<Type, unsigned>>842getTypeBitWidthBoundPairs(Operation *op, unsigned targetBitWidth) {843 844 if (auto insertOp = dyn_cast<vector::InsertOp>(op)) {845 unsigned w = targetBitWidth < std::numeric_limits<unsigned>::max()846 ? targetBitWidth + 1847 : targetBitWidth;848 return {{insertOp.getValueToStoreType(), w}};849 }850 851 auto resultTypes = op->getResultTypes();852 SmallVector<std::pair<Type, unsigned>> resultsWithBitWidth;853 resultsWithBitWidth.reserve(resultTypes.size());854 for (Type type : resultTypes) {855 resultsWithBitWidth.push_back({type, targetBitWidth});856 }857 return resultsWithBitWidth;858}859 860/// If `type` is VectorType with trailing dimension of (bit) size greater than861/// or equal to `targetBitWidth`, its defining op is considered legal.862static bool863isNotLinearizableBecauseLargeInnerDimension(Type type,864 unsigned targetBitWidth) {865 866 VectorType vecType = dyn_cast<VectorType>(type);867 868 // Not linearizable for reasons other than what this function checks.869 if (!vecType || vecType.getRank() == 0)870 return false;871 872 // The width of the type 'index' is unbounded (and therefore potentially above873 // the target width).874 if (vecType.getElementType().isIndex())875 return true;876 877 unsigned finalDimSize = vecType.getShape().back();878 unsigned nbBitsPerElm = vecType.getElementTypeBitWidth();879 unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm;880 return trailingVecDimBitWidth >= targetBitWidth;881}882 883static bool884isNotLinearizableBecauseLargeInnerDimension(Operation *op,885 unsigned targetBitWidth) {886 // Check on bitwidths.887 SmallVector<std::pair<Type, unsigned>> toCheck =888 getTypeBitWidthBoundPairs(op, targetBitWidth);889 return llvm::any_of(toCheck, [&](std::pair<Type, unsigned> typeWidth) {890 return isNotLinearizableBecauseLargeInnerDimension(typeWidth.first,891 typeWidth.second);892 });893}894 895void populateWithBitWidthConstraints(TypeConverter &typeConverter,896 ConversionTarget &target,897 unsigned targetBitWidth) {898 899 // The general purpose definition of what ops are legal must come first.900 populateForVectorLinearize(typeConverter, target);901 902 // Extend the set of legal ops to include those with large inner-most903 // dimensions on selected operands/results.904 target.markUnknownOpDynamicallyLegal(905 [=](Operation *op) -> std::optional<bool> {906 if (isNotLinearizableBecauseLargeInnerDimension(op, targetBitWidth)) {907 return true;908 }909 return {};910 });911}912 913struct TestVectorBitWidthLinearize final914 : public PassWrapper<TestVectorBitWidthLinearize, OperationPass<>> {915 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBitWidthLinearize)916 917 TestVectorBitWidthLinearize() = default;918 TestVectorBitWidthLinearize(const TestVectorBitWidthLinearize &pass)919 : PassWrapper(pass) {}920 921 StringRef getArgument() const override {922 return "test-bit-width-constrained-vector-linearize";923 }924 StringRef getDescription() const override {925 return "Linearizes ND vectors for N >= 2 into 1D vectors, with constraints "926 "in inner-most dimension's bit width.";927 }928 void getDependentDialects(DialectRegistry ®istry) const override {929 registry.insert<vector::VectorDialect>();930 }931 932 Option<unsigned> targetVectorBitwidth{933 *this, "target-vector-bitwidth",934 llvm::cl::desc(935 "Minimum vector bitwidth to enable the flattening transformation"),936 llvm::cl::init(std::numeric_limits<unsigned>::max())};937 void runOnOperation() override {938 auto *context = &getContext();939 940 TypeConverter typeConverter;941 RewritePatternSet patterns(context);942 ConversionTarget target(*context);943 944 populateWithBitWidthConstraints(typeConverter, target,945 targetVectorBitwidth);946 947 vector::populateVectorLinearizeBasePatterns(typeConverter, target,948 patterns);949 950 vector::populateVectorLinearizeShuffleLikeOpsPatterns(typeConverter, target,951 patterns);952 953 if (failed(applyPartialConversion(getOperation(), target,954 std::move(patterns))))955 return signalPassFailure();956 }957};958 959struct TestVectorLinearize final960 : public PassWrapper<TestVectorLinearize, OperationPass<>> {961 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorLinearize)962 963 TestVectorLinearize() = default;964 965 StringRef getArgument() const override { return "test-vector-linearize"; }966 StringRef getDescription() const override {967 return "Linearizes ND vectors for N >= 2 into 1D vectors";968 }969 void getDependentDialects(DialectRegistry ®istry) const override {970 registry.insert<vector::VectorDialect, arith::ArithDialect>();971 }972 973 void runOnOperation() override {974 MLIRContext &context = getContext();975 TypeConverter converter;976 RewritePatternSet patterns(&context);977 ConversionTarget target(context);978 979 vector::populateForVectorLinearize(converter, target);980 981 vector::populateVectorLinearizeBasePatterns(converter, target, patterns);982 vector::populateVectorLinearizeShuffleLikeOpsPatterns(converter, target,983 patterns);984 mlir::scf::populateSCFStructuralTypeConversionsAndLegality(985 converter, patterns, target);986 987 if (failed(applyPartialConversion(getOperation(), target,988 std::move(patterns))))989 return signalPassFailure();990 }991};992 993struct TestEliminateVectorMasks994 : public PassWrapper<TestEliminateVectorMasks,995 OperationPass<func::FuncOp>> {996 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestEliminateVectorMasks)997 998 TestEliminateVectorMasks() = default;999 TestEliminateVectorMasks(const TestEliminateVectorMasks &pass)1000 : PassWrapper(pass) {}1001 1002 Option<unsigned> vscaleMin{1003 *this, "vscale-min", llvm::cl::desc("Minimum possible value of vscale."),1004 llvm::cl::init(1)};1005 Option<unsigned> vscaleMax{1006 *this, "vscale-max", llvm::cl::desc("Maximum possible value of vscale."),1007 llvm::cl::init(16)};1008 1009 StringRef getArgument() const final { return "test-eliminate-vector-masks"; }1010 StringRef getDescription() const final {1011 return "Test eliminating vector masks";1012 }1013 void runOnOperation() override {1014 IRRewriter rewriter(&getContext());1015 eliminateVectorMasks(rewriter, getOperation(),1016 VscaleRange{vscaleMin, vscaleMax});1017 }1018};1019 1020struct TestVectorShuffleLowering1021 : public PassWrapper<TestVectorShuffleLowering,1022 OperationPass<func::FuncOp>> {1023 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorShuffleLowering)1024 1025 StringRef getArgument() const final { return "test-vector-shuffle-lowering"; }1026 StringRef getDescription() const final {1027 return "Test lowering patterns for vector.shuffle with mixed-size inputs";1028 }1029 void runOnOperation() override {1030 RewritePatternSet patterns(&getContext());1031 populateVectorShuffleLoweringPatterns(patterns);1032 (void)applyPatternsGreedily(getOperation(), std::move(patterns));1033 }1034};1035} // namespace1036 1037namespace mlir {1038namespace test {1039void registerTestVectorLowerings() {1040 PassRegistration<TestVectorToVectorLowering>();1041 1042 PassRegistration<TestVectorContractionPrepareForMMTLowering>();1043 1044 PassRegistration<TestVectorUnrollingPatterns>();1045 1046 PassRegistration<TestVectorTransferUnrollingPatterns>();1047 1048 PassRegistration<TestScalarVectorTransferLoweringPatterns>();1049 1050 PassRegistration<TestVectorTransferOpt>();1051 1052 PassRegistration<TestVectorSinkPatterns>();1053 1054 PassRegistration<TestVectorReduceToContractPatternsPatterns>();1055 1056 PassRegistration<TestVectorChainedReductionFoldingPatterns>();1057 1058 PassRegistration<TestVectorBreakDownReductionPatterns>();1059 1060 PassRegistration<TestFlattenVectorTransferPatterns>();1061 1062 PassRegistration<TestVectorScanLowering>();1063 1064 PassRegistration<TestVectorShuffleLowering>();1065 1066 PassRegistration<TestVectorDistribution>();1067 1068 PassRegistration<TestVectorExtractStridedSliceLowering>();1069 1070 PassRegistration<TestVectorBreakDownBitCast>();1071 1072 PassRegistration<TestCreateVectorBroadcast>();1073 1074 PassRegistration<TestVectorGatherLowering>();1075 1076 PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();1077 1078 PassRegistration<TestVectorEmulateMaskedLoadStore>();1079 1080 PassRegistration<TestVectorLinearize>();1081 1082 PassRegistration<TestVectorBitWidthLinearize>();1083 1084 PassRegistration<TestEliminateVectorMasks>();1085}1086} // namespace test1087} // namespace mlir1088