xref: /llvm-project/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp (revision 576b184d6e3b633f51b908b61ebd281d2ecbf66f) 
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include <optional>
10 #include <type_traits>
11 
12 #include "mlir/Analysis/SliceAnalysis.h"
13 #include "mlir/Dialect/Affine/IR/AffineOps.h"
14 #include "mlir/Dialect/Arith/IR/Arith.h"
15 #include "mlir/Dialect/Func/IR/FuncOps.h"
16 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
17 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
18 #include "mlir/Dialect/Linalg/IR/Linalg.h"
19 #include "mlir/Dialect/Linalg/Passes.h"
20 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
21 #include "mlir/Dialect/MemRef/IR/MemRef.h"
22 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
23 #include "mlir/Dialect/SCF/IR/SCF.h"
24 #include "mlir/Dialect/Tensor/IR/Tensor.h"
25 #include "mlir/Dialect/Vector/IR/VectorOps.h"
26 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
27 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
28 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
29 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
30 #include "mlir/Pass/Pass.h"
31 #include "mlir/Pass/PassManager.h"
32 #include "mlir/Support/LLVM.h"
33 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
34 
35 using namespace mlir;
36 using namespace mlir::linalg;
37 using namespace mlir::vector;
38 
39 namespace {
40 
41 struct TestVectorToVectorLowering
42     : public PassWrapper<TestVectorToVectorLowering,
43                          OperationPass<func::FuncOp>> {
44   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorToVectorLowering)
45 
46   TestVectorToVectorLowering() = default;
47   TestVectorToVectorLowering(const TestVectorToVectorLowering &pass)
48       : PassWrapper(pass) {}
49   StringRef getArgument() const final {
50     return "test-vector-to-vector-lowering";
51   }
52   StringRef getDescription() const final {
53     return "Test lowering patterns between ops in the vector dialect";
54   }
55 
56   void getDependentDialects(DialectRegistry &registry) const override {
57     registry.insert<affine::AffineDialect>();
58     registry.insert<vector::VectorDialect>();
59   }
60 
61   Option<bool> unroll{*this, "unroll", llvm::cl::desc("Include unrolling"),
62                       llvm::cl::init(false)};
63 
64   void runOnOperation() override {
65     auto *ctx = &getContext();
66     RewritePatternSet patterns(ctx);
67     if (unroll) {
68       populateVectorUnrollPatterns(
69           patterns,
70           UnrollVectorOptions().setNativeShapeFn(getShape).setFilterConstraint(
71               filter));
72     }
73     populateVectorToVectorCanonicalizationPatterns(patterns);
74     populateBubbleVectorBitCastOpPatterns(patterns);
75     populateCastAwayVectorLeadingOneDimPatterns(patterns);
76     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
77   }
78 
79 private:
80   // Return the target shape based on op type.
81   static std::optional<SmallVector<int64_t>> getShape(Operation *op) {
82     if (isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp>(op))
83       return SmallVector<int64_t>(2, 2);
84     if (isa<vector::ContractionOp>(op))
85       return SmallVector<int64_t>(3, 2);
86     // For transfer ops, just propagate the shape coming from
87     // InsertStridedSlices/ExtractStridedSlices.
88     if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
89       VectorType dstVec;
90       for (Operation *users : readOp->getUsers()) {
91         auto extract = dyn_cast<ExtractStridedSliceOp>(users);
92         if (!extract)
93           return std::nullopt;
94         auto vecType = cast<VectorType>(extract.getResult().getType());
95         if (dstVec && dstVec != vecType)
96           return std::nullopt;
97         dstVec = vecType;
98       }
99       return SmallVector<int64_t>(dstVec.getShape().begin(),
100                                   dstVec.getShape().end());
101     }
102     if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
103       auto insert = writeOp.getVector().getDefiningOp<InsertStridedSliceOp>();
104       if (!insert)
105         return std::nullopt;
106       ArrayRef<int64_t> shape = insert.getSourceVectorType().getShape();
107       return SmallVector<int64_t>(shape.begin(), shape.end());
108     }
109     return std::nullopt;
110   }
111 
112   static LogicalResult filter(Operation *op) {
113     return success(isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp,
114                        ContractionOp, TransferReadOp, TransferWriteOp>(op));
115   }
116 };
117 
118 struct TestVectorContractionPrepareForMMTLowering
119     : public PassWrapper<TestVectorContractionPrepareForMMTLowering,
120                          OperationPass<func::FuncOp>> {
121   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
122       TestVectorContractionPrepareForMMTLowering)
123 
124   StringRef getArgument() const final {
125     return "test-vector-contraction-prepare-for-mmt-lowering";
126   }
127   StringRef getDescription() const final {
128     return "Test vector.contraction matmul canonicalization for MMT lowering.";
129   }
130   TestVectorContractionPrepareForMMTLowering() = default;
131 
132   void getDependentDialects(DialectRegistry &registry) const override {
133     registry.insert<affine::AffineDialect, arith::ArithDialect,
134                     vector::VectorDialect>();
135   }
136 
137   void runOnOperation() override {
138     MLIRContext *ctx = &getContext();
139     RewritePatternSet patterns(ctx);
140     vector::populateVectorContractCanonicalizeMatmulToMMT(patterns);
141     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
142   }
143 };
144 
145 struct TestVectorUnrollingPatterns
146     : public PassWrapper<TestVectorUnrollingPatterns,
147                          OperationPass<func::FuncOp>> {
148   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorUnrollingPatterns)
149 
150   StringRef getArgument() const final {
151     return "test-vector-unrolling-patterns";
152   }
153   StringRef getDescription() const final {
154     return "Test lowering patterns to unroll contract ops in the vector "
155            "dialect";
156   }
157   TestVectorUnrollingPatterns() = default;
158   TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass)
159       : PassWrapper(pass) {}
160   void runOnOperation() override {
161     MLIRContext *ctx = &getContext();
162     RewritePatternSet patterns(ctx);
163     populateVectorUnrollPatterns(
164         patterns, UnrollVectorOptions()
165                       .setNativeShape(ArrayRef<int64_t>{2, 2})
166                       .setFilterConstraint([](Operation *op) {
167                         return success(isa<arith::AddFOp, vector::FMAOp,
168                                            vector::MultiDimReductionOp>(op));
169                       }));
170     populateVectorUnrollPatterns(
171         patterns, UnrollVectorOptions()
172                       .setNativeShape(ArrayRef<int64_t>{2})
173                       .setFilterConstraint([](Operation *op) {
174                         return success(isa<vector::ReductionOp>(op));
175                       }));
176     populateVectorUnrollPatterns(
177         patterns, UnrollVectorOptions()
178                       .setNativeShape(ArrayRef<int64_t>{1, 3, 4, 2})
179                       .setFilterConstraint([](Operation *op) {
180                         return success(isa<vector::TransposeOp>(op));
181                       }));
182 
183     if (unrollBasedOnType) {
184       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
185           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
186         vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
187         SmallVector<int64_t> nativeShape(contractOp.getIteratorTypes().size(),
188                                          4);
189         Type lhsType = contractOp.getLhsType().getElementType();
190         nativeShape[nativeShape.size() - 1] = lhsType.isF16() ? 4 : 2;
191         return nativeShape;
192       };
193 
194       UnrollVectorOptions opts;
195       opts.setNativeShapeFn(nativeShapeFn)
196           .setFilterConstraint(
197               [](Operation *op) { return success(isa<ContractionOp>(op)); });
198 
199       if (!unrollOrder.empty()) {
200         opts.setUnrollTraversalOrderFn(
201             [this](Operation *op) -> std::optional<SmallVector<int64_t>> {
202               vector::ContractionOp contractOp =
203                   cast<vector::ContractionOp>(op);
204               if (contractOp.getIteratorTypes().size() == unrollOrder.size())
205                 return SmallVector<int64_t>(unrollOrder.begin(),
206                                             unrollOrder.end());
207               return std::nullopt;
208             });
209       }
210       populateVectorUnrollPatterns(patterns, opts);
211     } else {
212       auto nativeShapeFn =
213           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
214         auto contractOp = dyn_cast<ContractionOp>(op);
215         if (!contractOp)
216           return std::nullopt;
217         return SmallVector<int64_t>(contractOp.getIteratorTypes().size(), 2);
218       };
219       populateVectorUnrollPatterns(patterns,
220                                    UnrollVectorOptions()
221                                        .setNativeShapeFn(nativeShapeFn)
222                                        .setFilterConstraint([](Operation *op) {
223                                          return success(isa<ContractionOp>(op));
224                                        }));
225     }
226     populateVectorToVectorCanonicalizationPatterns(patterns);
227     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
228   }
229 
230   ListOption<int64_t> unrollOrder{*this, "unroll-order",
231                                   llvm::cl::desc("set the unroll order")};
232 
233   Option<bool> unrollBasedOnType{
234       *this, "unroll-based-on-type",
235       llvm::cl::desc("Set the unroll factor based on type of the operation"),
236       llvm::cl::init(false)};
237 };
238 
239 struct TestVectorTransferUnrollingPatterns
240     : public PassWrapper<TestVectorTransferUnrollingPatterns,
241                          OperationPass<func::FuncOp>> {
242   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
243       TestVectorTransferUnrollingPatterns)
244 
245   TestVectorTransferUnrollingPatterns() = default;
246   TestVectorTransferUnrollingPatterns(
247       const TestVectorTransferUnrollingPatterns &pass)
248       : PassWrapper(pass) {}
249 
250   void getDependentDialects(DialectRegistry &registry) const override {
251     registry.insert<affine::AffineDialect>();
252   }
253   StringRef getArgument() const final {
254     return "test-vector-transfer-unrolling-patterns";
255   }
256   StringRef getDescription() const final {
257     return "Test lowering patterns to unroll transfer ops in the vector "
258            "dialect";
259   }
260   void runOnOperation() override {
261     MLIRContext *ctx = &getContext();
262     RewritePatternSet patterns(ctx);
263     UnrollVectorOptions opts;
264     opts.setNativeShape(ArrayRef<int64_t>{2, 2})
265         .setFilterConstraint([](Operation *op) {
266           return success(isa<vector::TransferReadOp, vector::TransferWriteOp,
267                              vector::GatherOp>(op));
268         });
269     if (reverseUnrollOrder.getValue()) {
270       opts.setUnrollTraversalOrderFn(
271           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
272             int64_t numLoops = 0;
273             if (auto readOp = dyn_cast<vector::TransferReadOp>(op))
274               numLoops = readOp.getVectorType().getRank();
275             else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op))
276               numLoops = writeOp.getVectorType().getRank();
277             else if (auto gatherOp = dyn_cast<vector::GatherOp>(op))
278               numLoops = gatherOp.getVectorType().getRank();
279             else
280               return std::nullopt;
281             auto order = llvm::reverse(llvm::seq<int64_t>(0, numLoops));
282             return llvm::to_vector(order);
283           });
284     }
285     populateVectorUnrollPatterns(patterns, opts);
286     populateVectorToVectorCanonicalizationPatterns(patterns);
287     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
288   }
289 
290   Option<bool> reverseUnrollOrder{
291       *this, "reverse-unroll-order",
292       llvm::cl::desc(
293           "reverse the order of unrolling of vector transfer operations"),
294       llvm::cl::init(false)};
295 };
296 
297 struct TestScalarVectorTransferLoweringPatterns
298     : public PassWrapper<TestScalarVectorTransferLoweringPatterns,
299                          OperationPass<func::FuncOp>> {
300   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
301       TestScalarVectorTransferLoweringPatterns)
302 
303   TestScalarVectorTransferLoweringPatterns() = default;
304   TestScalarVectorTransferLoweringPatterns(
305       const TestScalarVectorTransferLoweringPatterns &pass)
306       : PassWrapper(pass) {}
307 
308   StringRef getArgument() const final {
309     return "test-scalar-vector-transfer-lowering";
310   }
311   StringRef getDescription() const final {
312     return "Test lowering of scalar vector transfers to memref loads/stores.";
313   }
314 
315   void getDependentDialects(DialectRegistry &registry) const override {
316     registry.insert<affine::AffineDialect, memref::MemRefDialect,
317                     tensor::TensorDialect, vector::VectorDialect>();
318   }
319 
320   Option<bool> allowMultipleUses{
321       *this, "allow-multiple-uses",
322       llvm::cl::desc("Fold transfer operations with multiple uses"),
323       llvm::cl::init(false)};
324 
325   void runOnOperation() override {
326     MLIRContext *ctx = &getContext();
327     RewritePatternSet patterns(ctx);
328     vector::populateScalarVectorTransferLoweringPatterns(
329         patterns, /*benefit=*/1, allowMultipleUses.getValue());
330     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
331   }
332 };
333 
334 struct TestVectorTransferOpt
335     : public PassWrapper<TestVectorTransferOpt, OperationPass<func::FuncOp>> {
336   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorTransferOpt)
337 
338   StringRef getArgument() const final { return "test-vector-transferop-opt"; }
339   StringRef getDescription() const final {
340     return "Test optimization transformations for transfer ops";
341   }
342   void runOnOperation() override {
343     IRRewriter rewriter(&getContext());
344     transferOpflowOpt(rewriter, getOperation());
345   }
346 };
347 
348 struct TestVectorTransferCollapseInnerMostContiguousDims
349     : public PassWrapper<TestVectorTransferCollapseInnerMostContiguousDims,
350                          OperationPass<func::FuncOp>> {
351   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
352       TestVectorTransferCollapseInnerMostContiguousDims)
353 
354   TestVectorTransferCollapseInnerMostContiguousDims() = default;
355   TestVectorTransferCollapseInnerMostContiguousDims(
356       const TestVectorTransferCollapseInnerMostContiguousDims &pass) = default;
357 
358   void getDependentDialects(DialectRegistry &registry) const override {
359     registry.insert<memref::MemRefDialect, affine::AffineDialect>();
360   }
361 
362   StringRef getArgument() const final {
363     return "test-vector-transfer-collapse-inner-most-dims";
364   }
365 
366   StringRef getDescription() const final {
367     return "Test lowering patterns that reducedes the rank of the vector "
368            "transfer memory and vector operands.";
369   }
370 
371   void runOnOperation() override {
372     RewritePatternSet patterns(&getContext());
373     populateVectorTransferCollapseInnerMostContiguousDimsPatterns(patterns);
374     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
375   }
376 };
377 
378 struct TestSinkVectorBroadcast
379     : public PassWrapper<TestSinkVectorBroadcast, OperationPass<func::FuncOp>> {
380   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSinkVectorBroadcast)
381 
382   TestSinkVectorBroadcast() = default;
383   TestSinkVectorBroadcast(const TestSinkVectorBroadcast &pass) = default;
384 
385   void getDependentDialects(DialectRegistry &registry) const override {
386     registry.insert<memref::MemRefDialect, affine::AffineDialect>();
387   }
388 
389   StringRef getArgument() const final { return "test-sink-vector-broadcast"; }
390 
391   StringRef getDescription() const final {
392     return "Test lowering patterns that eliminate redundant brodacast "
393            "operations.";
394   }
395 
396   void runOnOperation() override {
397     RewritePatternSet patterns(&getContext());
398     populateSinkVectorBroadcastPatterns(patterns);
399     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
400   }
401 };
402 
403 struct TestVectorReduceToContractPatternsPatterns
404     : public PassWrapper<TestVectorReduceToContractPatternsPatterns,
405                          OperationPass<func::FuncOp>> {
406   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
407       TestVectorReduceToContractPatternsPatterns)
408 
409   StringRef getArgument() const final {
410     return "test-vector-reduction-to-contract-patterns";
411   }
412   StringRef getDescription() const final {
413     return "Test patterns to convert multireduce op to contract and combine "
414            "broadcast/transpose to contract";
415   }
416   void runOnOperation() override {
417     RewritePatternSet patterns(&getContext());
418     populateVectorReductionToContractPatterns(patterns);
419     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
420   }
421 };
422 
423 struct TestFlattenVectorTransferPatterns
424     : public PassWrapper<TestFlattenVectorTransferPatterns,
425                          OperationPass<func::FuncOp>> {
426   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
427       TestFlattenVectorTransferPatterns)
428 
429   StringRef getArgument() const final {
430     return "test-vector-transfer-flatten-patterns";
431   }
432   StringRef getDescription() const final {
433     return "Test patterns to rewrite contiguous row-major N-dimensional "
434            "vector.transfer_{read,write} ops into 1D transfers";
435   }
436   void getDependentDialects(DialectRegistry &registry) const override {
437     registry.insert<memref::MemRefDialect>();
438   }
439   void runOnOperation() override {
440     RewritePatternSet patterns(&getContext());
441     populateFlattenVectorTransferPatterns(patterns);
442     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
443   }
444 };
445 
446 struct TestVectorScanLowering
447     : public PassWrapper<TestVectorScanLowering, OperationPass<func::FuncOp>> {
448   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorScanLowering)
449 
450   StringRef getArgument() const final { return "test-vector-scan-lowering"; }
451   StringRef getDescription() const final {
452     return "Test lowering patterns that lower the scan op in the vector "
453            "dialect";
454   }
455   void runOnOperation() override {
456     RewritePatternSet patterns(&getContext());
457     populateVectorScanLoweringPatterns(patterns);
458     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
459   }
460 };
461 
462 /// Allocate shared memory for a single warp to test lowering of
463 /// WarpExecuteOnLane0Op.
464 static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
465                                         WarpExecuteOnLane0Op warpOp,
466                                         Type type) {
467   static constexpr int64_t kSharedMemorySpace = 3;
468   // Compute type of shared memory buffer.
469   MemRefType memrefType;
470   if (auto vectorType = dyn_cast<VectorType>(type)) {
471     memrefType =
472         MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},
473                         kSharedMemorySpace);
474   } else {
475     memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);
476   }
477 
478   // Get symbol table holding all shared memory globals.
479   ModuleOp moduleOp = warpOp->getParentOfType<ModuleOp>();
480   SymbolTable symbolTable(moduleOp);
481 
482   // Create a pretty name.
483   SmallString<64> buf;
484   llvm::raw_svector_ostream os(buf);
485   interleave(memrefType.getShape(), os, "x");
486   os << "x" << memrefType.getElementType();
487   std::string symbolName = (Twine("__shared_") + os.str()).str();
488 
489   auto ip = builder.saveInsertionPoint();
490   builder.setInsertionPoint(moduleOp);
491   auto global = builder.create<memref::GlobalOp>(
492       loc,
493       /*sym_name=*/symbolName,
494       /*sym_visibility=*/builder.getStringAttr("private"),
495       /*type=*/memrefType,
496       /*initial_value=*/Attribute(),
497       /*constant=*/false,
498       /*alignment=*/IntegerAttr());
499   symbolTable.insert(global);
500   // The symbol table inserts at the end of the module, but globals are a bit
501   // nicer if they are at the beginning.
502   global->moveBefore(&moduleOp.front());
503 
504   builder.restoreInsertionPoint(ip);
505   return builder.create<memref::GetGlobalOp>(loc, memrefType, symbolName);
506 }
507 
508 static Value warpReduction(Location loc, OpBuilder &builder, Value input,
509                            CombiningKind kind, uint32_t size) {
510   // First reduce on a single thread to get per lane reduction value.
511   Value laneVal = builder.create<vector::ReductionOp>(loc, kind, input);
512   // Parallel reduction using butterfly shuffles.
513   for (uint64_t i = 1; i < size; i <<= 1) {
514     Value shuffled = builder
515                          .create<gpu::ShuffleOp>(loc, laneVal, i,
516                                                  /*width=*/size,
517                                                  /*mode=*/gpu::ShuffleMode::XOR)
518                          .getShuffleResult();
519     laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);
520   }
521   return laneVal;
522 }
523 
524 struct TestVectorDistribution
525     : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {
526   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)
527 
528   void getDependentDialects(DialectRegistry &registry) const override {
529     registry.insert<scf::SCFDialect, memref::MemRefDialect, gpu::GPUDialect,
530                     affine::AffineDialect>();
531   }
532 
533   StringRef getArgument() const final { return "test-vector-warp-distribute"; }
534   StringRef getDescription() const final {
535     return "Test vector warp distribute transformation and lowering patterns";
536   }
537   TestVectorDistribution() = default;
538   TestVectorDistribution(const TestVectorDistribution &pass)
539       : PassWrapper(pass) {}
540 
541   Option<bool> warpOpToSCF{
542       *this, "rewrite-warp-ops-to-scf-if",
543       llvm::cl::desc("Lower vector.warp_execute_on_lane0 to scf.if op"),
544       llvm::cl::init(false)};
545 
546   Option<bool> distributeTransferWriteOps{
547       *this, "distribute-transfer-write",
548       llvm::cl::desc("Test distribution of transfer write"),
549       llvm::cl::init(false)};
550 
551   Option<bool> hoistUniform{*this, "hoist-uniform",
552                             llvm::cl::desc("Test hoist uniform"),
553                             llvm::cl::init(false)};
554 
555   Option<bool> propagateDistribution{
556       *this, "propagate-distribution",
557       llvm::cl::desc("Test distribution propgation"), llvm::cl::init(false)};
558 
559   void runOnOperation() override {
560     RewritePatternSet patterns(&getContext());
561 
562     getOperation().walk([&](Operation *op) {
563       if (auto warpOp = dyn_cast<WarpExecuteOnLane0Op>(op)) {
564         if (hoistUniform) {
565           moveScalarUniformCode(warpOp);
566         }
567         WalkResult::interrupt();
568       }
569     });
570     MLIRContext *ctx = &getContext();
571     auto distributionFn = [](Value val) {
572       // Create a map (d0, d1) -> (d1) to distribute along the inner
573       // dimension. Once we support n-d distribution we can add more
574       // complex cases.
575       VectorType vecType = dyn_cast<VectorType>(val.getType());
576       int64_t vecRank = vecType ? vecType.getRank() : 0;
577       OpBuilder builder(val.getContext());
578       if (vecRank == 0)
579         return AffineMap::get(val.getContext());
580       return AffineMap::get(vecRank, 0, builder.getAffineDimExpr(vecRank - 1));
581     };
582     auto shuffleFn = [](Location loc, OpBuilder &builder, Value val,
583                         Value srcIdx, int64_t warpSz) {
584       assert((val.getType().isF32() || val.getType().isInteger(32)) &&
585              "unsupported shuffle type");
586       Type i32Type = builder.getIntegerType(32);
587       Value srcIdxI32 =
588           builder.create<arith::IndexCastOp>(loc, i32Type, srcIdx);
589       Value warpSzI32 = builder.create<arith::ConstantOp>(
590           loc, builder.getIntegerAttr(i32Type, warpSz));
591       Value result = builder
592                          .create<gpu::ShuffleOp>(loc, val, srcIdxI32, warpSzI32,
593                                                  gpu::ShuffleMode::IDX)
594                          .getResult(0);
595       return result;
596     };
597     if (distributeTransferWriteOps) {
598       RewritePatternSet patterns(ctx);
599       populateDistributeTransferWriteOpPatterns(patterns, distributionFn);
600       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
601     }
602     if (propagateDistribution) {
603       RewritePatternSet patterns(ctx);
604       vector::populatePropagateWarpVectorDistributionPatterns(
605           patterns, distributionFn, shuffleFn);
606       vector::populateDistributeReduction(patterns, warpReduction);
607       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
608     }
609     WarpExecuteOnLane0LoweringOptions options;
610     options.warpAllocationFn = allocateGlobalSharedMemory;
611     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
612                                       WarpExecuteOnLane0Op warpOp) {
613       builder.create<gpu::BarrierOp>(loc);
614     };
615     // Test on one pattern in isolation.
616     if (warpOpToSCF) {
617       populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);
618       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
619       return;
620     }
621   }
622 };
623 
624 struct TestVectorExtractStridedSliceLowering
625     : public PassWrapper<TestVectorExtractStridedSliceLowering,
626                          OperationPass<func::FuncOp>> {
627   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
628       TestVectorExtractStridedSliceLowering)
629 
630   StringRef getArgument() const final {
631     return "test-vector-extract-strided-slice-lowering";
632   }
633   StringRef getDescription() const final {
634     return "Test lowering patterns that converts vector.extract_strided_slice "
635            "into a chain of vector.extract and vector.insert ops";
636   }
637   void runOnOperation() override {
638     RewritePatternSet patterns(&getContext());
639     populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);
640     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
641   }
642 };
643 
644 struct TestVectorBreakDownBitCast
645     : public PassWrapper<TestVectorBreakDownBitCast,
646                          OperationPass<func::FuncOp>> {
647   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBreakDownBitCast)
648 
649   StringRef getArgument() const final {
650     return "test-vector-break-down-bitcast";
651   }
652   StringRef getDescription() const final {
653     return "Test pattern that breaks down vector.bitcast ops ";
654   }
655   void runOnOperation() override {
656     RewritePatternSet patterns(&getContext());
657     populateBreakDownVectorBitCastOpPatterns(patterns, [](BitCastOp op) {
658       return op.getSourceVectorType().getShape().back() > 4;
659     });
660     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
661   }
662 };
663 
664 struct TestCreateVectorBroadcast
665     : public PassWrapper<TestCreateVectorBroadcast,
666                          OperationPass<func::FuncOp>> {
667   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)
668 
669   StringRef getArgument() const final { return "test-create-vector-broadcast"; }
670   StringRef getDescription() const final {
671     return "Test optimization transformations for transfer ops";
672   }
673   void getDependentDialects(DialectRegistry &registry) const override {
674     registry.insert<vector::VectorDialect>();
675   }
676 
677   void runOnOperation() override {
678     getOperation()->walk([](Operation *op) {
679       if (op->getName().getStringRef() != "test_create_broadcast")
680         return;
681       auto targetShape =
682           cast<VectorType>(op->getResult(0).getType()).getShape();
683       auto arrayAttr =
684           cast<DenseI64ArrayAttr>(op->getAttr("broadcast_dims")).asArrayRef();
685       llvm::SetVector<int64_t> broadcastedDims;
686       broadcastedDims.insert(arrayAttr.begin(), arrayAttr.end());
687       OpBuilder b(op);
688       Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(
689           b, op->getOperand(0), targetShape, broadcastedDims);
690       op->getResult(0).replaceAllUsesWith(bcast);
691       op->erase();
692     });
693   }
694 };
695 
696 struct TestVectorGatherLowering
697     : public PassWrapper<TestVectorGatherLowering,
698                          OperationPass<func::FuncOp>> {
699   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)
700 
701   StringRef getArgument() const final { return "test-vector-gather-lowering"; }
702   StringRef getDescription() const final {
703     return "Test patterns that lower the gather op in the vector conditional "
704            "loads";
705   }
706   void getDependentDialects(DialectRegistry &registry) const override {
707     registry.insert<arith::ArithDialect, func::FuncDialect,
708                     memref::MemRefDialect, scf::SCFDialect,
709                     tensor::TensorDialect, vector::VectorDialect>();
710   }
711 
712   void runOnOperation() override {
713     RewritePatternSet patterns(&getContext());
714     populateVectorGatherLoweringPatterns(patterns);
715     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
716   }
717 };
718 
719 struct TestFoldArithExtensionIntoVectorContractPatterns
720     : public PassWrapper<TestFoldArithExtensionIntoVectorContractPatterns,
721                          OperationPass<func::FuncOp>> {
722   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
723       TestFoldArithExtensionIntoVectorContractPatterns)
724 
725   StringRef getArgument() const final {
726     return "test-fold-arith-extf-into-vector-contract-patterns";
727   }
728   StringRef getDescription() const final {
729     return "Test patterns that fold arithmetic extension ops into vector "
730            "contract ops";
731   }
732 
733   void getDependentDialects(DialectRegistry &registry) const override {
734     registry.insert<arith::ArithDialect, func::FuncDialect, nvgpu::NVGPUDialect,
735                     memref::MemRefDialect, scf::SCFDialect,
736                     tensor::TensorDialect, vector::VectorDialect>();
737   }
738 
739   void runOnOperation() override {
740     RewritePatternSet patterns(&getContext());
741     populateFoldArithExtensionPatterns(patterns);
742     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
743   }
744 };
745 } // namespace
746 
747 namespace mlir {
748 namespace test {
749 void registerTestVectorLowerings() {
750   PassRegistration<TestVectorToVectorLowering>();
751 
752   PassRegistration<TestVectorContractionPrepareForMMTLowering>();
753 
754   PassRegistration<TestVectorUnrollingPatterns>();
755 
756   PassRegistration<TestVectorTransferUnrollingPatterns>();
757 
758   PassRegistration<TestScalarVectorTransferLoweringPatterns>();
759 
760   PassRegistration<TestVectorTransferOpt>();
761 
762   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
763 
764   PassRegistration<TestSinkVectorBroadcast>();
765 
766   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
767 
768   PassRegistration<TestFlattenVectorTransferPatterns>();
769 
770   PassRegistration<TestVectorScanLowering>();
771 
772   PassRegistration<TestVectorDistribution>();
773 
774   PassRegistration<TestVectorExtractStridedSliceLowering>();
775 
776   PassRegistration<TestVectorBreakDownBitCast>();
777 
778   PassRegistration<TestCreateVectorBroadcast>();
779 
780   PassRegistration<TestVectorGatherLowering>();
781 
782   PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();
783 }
784 } // namespace test
785 } // namespace mlir
786