xref: /llvm-project/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp (revision c02d07fdf007afc6b928cda0342751889cc2604b) 
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 TestVectorChainedReductionFoldingPatterns
424     : public PassWrapper<TestVectorChainedReductionFoldingPatterns,
425                          OperationPass<func::FuncOp>> {
426   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
427       TestVectorChainedReductionFoldingPatterns)
428 
429   StringRef getArgument() const final {
430     return "test-vector-chained-reduction-folding-patterns";
431   }
432   StringRef getDescription() const final {
433     return "Test patterns to fold chained vector reductions";
434   }
435   void runOnOperation() override {
436     RewritePatternSet patterns(&getContext());
437     populateChainedVectorReductionFoldingPatterns(patterns);
438     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
439   }
440 };
441 
442 struct TestFlattenVectorTransferPatterns
443     : public PassWrapper<TestFlattenVectorTransferPatterns,
444                          OperationPass<func::FuncOp>> {
445   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
446       TestFlattenVectorTransferPatterns)
447 
448   StringRef getArgument() const final {
449     return "test-vector-transfer-flatten-patterns";
450   }
451   StringRef getDescription() const final {
452     return "Test patterns to rewrite contiguous row-major N-dimensional "
453            "vector.transfer_{read,write} ops into 1D transfers";
454   }
455   void getDependentDialects(DialectRegistry &registry) const override {
456     registry.insert<memref::MemRefDialect>();
457     registry.insert<affine::AffineDialect>();
458     registry.insert<vector::VectorDialect>();
459   }
460   void runOnOperation() override {
461     RewritePatternSet patterns(&getContext());
462     populateFlattenVectorTransferPatterns(patterns);
463     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
464   }
465 };
466 
467 struct TestVectorScanLowering
468     : public PassWrapper<TestVectorScanLowering, OperationPass<func::FuncOp>> {
469   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorScanLowering)
470 
471   StringRef getArgument() const final { return "test-vector-scan-lowering"; }
472   StringRef getDescription() const final {
473     return "Test lowering patterns that lower the scan op in the vector "
474            "dialect";
475   }
476   void runOnOperation() override {
477     RewritePatternSet patterns(&getContext());
478     populateVectorScanLoweringPatterns(patterns);
479     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
480   }
481 };
482 
483 /// Allocate shared memory for a single warp to test lowering of
484 /// WarpExecuteOnLane0Op.
485 static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
486                                         WarpExecuteOnLane0Op warpOp,
487                                         Type type) {
488   static constexpr int64_t kSharedMemorySpace = 3;
489   // Compute type of shared memory buffer.
490   MemRefType memrefType;
491   if (auto vectorType = dyn_cast<VectorType>(type)) {
492     memrefType =
493         MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},
494                         kSharedMemorySpace);
495   } else {
496     memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);
497   }
498 
499   // Get symbol table holding all shared memory globals.
500   ModuleOp moduleOp = warpOp->getParentOfType<ModuleOp>();
501   SymbolTable symbolTable(moduleOp);
502 
503   // Create a pretty name.
504   SmallString<64> buf;
505   llvm::raw_svector_ostream os(buf);
506   interleave(memrefType.getShape(), os, "x");
507   os << "x" << memrefType.getElementType();
508   std::string symbolName = (Twine("__shared_") + os.str()).str();
509 
510   auto ip = builder.saveInsertionPoint();
511   builder.setInsertionPoint(moduleOp);
512   auto global = builder.create<memref::GlobalOp>(
513       loc,
514       /*sym_name=*/symbolName,
515       /*sym_visibility=*/builder.getStringAttr("private"),
516       /*type=*/memrefType,
517       /*initial_value=*/Attribute(),
518       /*constant=*/false,
519       /*alignment=*/IntegerAttr());
520   symbolTable.insert(global);
521   // The symbol table inserts at the end of the module, but globals are a bit
522   // nicer if they are at the beginning.
523   global->moveBefore(&moduleOp.front());
524 
525   builder.restoreInsertionPoint(ip);
526   return builder.create<memref::GetGlobalOp>(loc, memrefType, symbolName);
527 }
528 
529 static Value warpReduction(Location loc, OpBuilder &builder, Value input,
530                            CombiningKind kind, uint32_t size) {
531   // First reduce on a single thread to get per lane reduction value.
532   Value laneVal = builder.create<vector::ReductionOp>(loc, kind, input);
533   // Parallel reduction using butterfly shuffles.
534   for (uint64_t i = 1; i < size; i <<= 1) {
535     Value shuffled = builder
536                          .create<gpu::ShuffleOp>(loc, laneVal, i,
537                                                  /*width=*/size,
538                                                  /*mode=*/gpu::ShuffleMode::XOR)
539                          .getShuffleResult();
540     laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);
541   }
542   return laneVal;
543 }
544 
545 struct TestVectorDistribution
546     : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {
547   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)
548 
549   void getDependentDialects(DialectRegistry &registry) const override {
550     registry.insert<scf::SCFDialect, memref::MemRefDialect, gpu::GPUDialect,
551                     affine::AffineDialect>();
552   }
553 
554   StringRef getArgument() const final { return "test-vector-warp-distribute"; }
555   StringRef getDescription() const final {
556     return "Test vector warp distribute transformation and lowering patterns";
557   }
558   TestVectorDistribution() = default;
559   TestVectorDistribution(const TestVectorDistribution &pass)
560       : PassWrapper(pass) {}
561 
562   Option<bool> warpOpToSCF{
563       *this, "rewrite-warp-ops-to-scf-if",
564       llvm::cl::desc("Lower vector.warp_execute_on_lane0 to scf.if op"),
565       llvm::cl::init(false)};
566 
567   Option<bool> distributeTransferWriteOps{
568       *this, "distribute-transfer-write",
569       llvm::cl::desc("Test distribution of transfer write"),
570       llvm::cl::init(false)};
571 
572   Option<unsigned> maxTransferWriteElements{
573       *this, "max-transfer-write-elements",
574       llvm::cl::desc("Maximum number of transfer write elements to distribute"),
575       llvm::cl::init(1)};
576 
577   Option<bool> hoistUniform{*this, "hoist-uniform",
578                             llvm::cl::desc("Test hoist uniform"),
579                             llvm::cl::init(false)};
580 
581   Option<bool> propagateDistribution{
582       *this, "propagate-distribution",
583       llvm::cl::desc("Test distribution propgation"), llvm::cl::init(false)};
584 
585   void runOnOperation() override {
586     RewritePatternSet patterns(&getContext());
587 
588     getOperation().walk([&](Operation *op) {
589       if (auto warpOp = dyn_cast<WarpExecuteOnLane0Op>(op)) {
590         if (hoistUniform) {
591           moveScalarUniformCode(warpOp);
592         }
593         WalkResult::interrupt();
594       }
595     });
596     MLIRContext *ctx = &getContext();
597     auto distributionFn = [](Value val) {
598       // Create a map (d0, d1) -> (d1) to distribute along the inner
599       // dimension. Once we support n-d distribution we can add more
600       // complex cases.
601       VectorType vecType = dyn_cast<VectorType>(val.getType());
602       int64_t vecRank = vecType ? vecType.getRank() : 0;
603       OpBuilder builder(val.getContext());
604       if (vecRank == 0)
605         return AffineMap::get(val.getContext());
606       return AffineMap::get(vecRank, 0, builder.getAffineDimExpr(vecRank - 1));
607     };
608     auto shuffleFn = [](Location loc, OpBuilder &builder, Value val,
609                         Value srcIdx, int64_t warpSz) {
610       assert((val.getType().isF32() || val.getType().isInteger(32)) &&
611              "unsupported shuffle type");
612       Type i32Type = builder.getIntegerType(32);
613       Value srcIdxI32 =
614           builder.create<arith::IndexCastOp>(loc, i32Type, srcIdx);
615       Value warpSzI32 = builder.create<arith::ConstantOp>(
616           loc, builder.getIntegerAttr(i32Type, warpSz));
617       Value result = builder
618                          .create<gpu::ShuffleOp>(loc, val, srcIdxI32, warpSzI32,
619                                                  gpu::ShuffleMode::IDX)
620                          .getResult(0);
621       return result;
622     };
623     if (distributeTransferWriteOps && propagateDistribution) {
624       RewritePatternSet patterns(ctx);
625       vector::populatePropagateWarpVectorDistributionPatterns(
626           patterns, distributionFn, shuffleFn, /*benefit=*/1,
627           /*readBenefit=*/0);
628       vector::populateDistributeReduction(patterns, warpReduction, 1);
629       populateDistributeTransferWriteOpPatterns(patterns, distributionFn, 2);
630       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
631     } else if (distributeTransferWriteOps) {
632       RewritePatternSet patterns(ctx);
633       populateDistributeTransferWriteOpPatterns(patterns, distributionFn,
634                                                 maxTransferWriteElements);
635       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
636     } else if (propagateDistribution) {
637       RewritePatternSet patterns(ctx);
638       vector::populatePropagateWarpVectorDistributionPatterns(
639           patterns, distributionFn, shuffleFn);
640       vector::populateDistributeReduction(patterns, warpReduction);
641       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
642     }
643     WarpExecuteOnLane0LoweringOptions options;
644     options.warpAllocationFn = allocateGlobalSharedMemory;
645     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
646                                       WarpExecuteOnLane0Op warpOp) {
647       builder.create<gpu::BarrierOp>(loc);
648     };
649     // Test on one pattern in isolation.
650     if (warpOpToSCF) {
651       populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);
652       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
653       return;
654     }
655   }
656 };
657 
658 struct TestVectorExtractStridedSliceLowering
659     : public PassWrapper<TestVectorExtractStridedSliceLowering,
660                          OperationPass<func::FuncOp>> {
661   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
662       TestVectorExtractStridedSliceLowering)
663 
664   StringRef getArgument() const final {
665     return "test-vector-extract-strided-slice-lowering";
666   }
667   StringRef getDescription() const final {
668     return "Test lowering patterns that converts vector.extract_strided_slice "
669            "into a chain of vector.extract and vector.insert ops";
670   }
671   void runOnOperation() override {
672     RewritePatternSet patterns(&getContext());
673     populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);
674     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
675   }
676 };
677 
678 struct TestVectorBreakDownBitCast
679     : public PassWrapper<TestVectorBreakDownBitCast,
680                          OperationPass<func::FuncOp>> {
681   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBreakDownBitCast)
682 
683   StringRef getArgument() const final {
684     return "test-vector-break-down-bitcast";
685   }
686   StringRef getDescription() const final {
687     return "Test pattern that breaks down vector.bitcast ops ";
688   }
689   void runOnOperation() override {
690     RewritePatternSet patterns(&getContext());
691     populateBreakDownVectorBitCastOpPatterns(patterns, [](BitCastOp op) {
692       return op.getSourceVectorType().getShape().back() > 4;
693     });
694     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
695   }
696 };
697 
698 struct TestCreateVectorBroadcast
699     : public PassWrapper<TestCreateVectorBroadcast,
700                          OperationPass<func::FuncOp>> {
701   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)
702 
703   StringRef getArgument() const final { return "test-create-vector-broadcast"; }
704   StringRef getDescription() const final {
705     return "Test optimization transformations for transfer ops";
706   }
707   void getDependentDialects(DialectRegistry &registry) const override {
708     registry.insert<vector::VectorDialect>();
709   }
710 
711   void runOnOperation() override {
712     getOperation()->walk([](Operation *op) {
713       if (op->getName().getStringRef() != "test_create_broadcast")
714         return;
715       auto targetShape =
716           cast<VectorType>(op->getResult(0).getType()).getShape();
717       auto arrayAttr =
718           cast<DenseI64ArrayAttr>(op->getDiscardableAttr("broadcast_dims"))
719               .asArrayRef();
720       llvm::SetVector<int64_t> broadcastedDims;
721       broadcastedDims.insert(arrayAttr.begin(), arrayAttr.end());
722       OpBuilder b(op);
723       Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(
724           b, op->getOperand(0), targetShape, broadcastedDims);
725       op->getResult(0).replaceAllUsesWith(bcast);
726       op->erase();
727     });
728   }
729 };
730 
731 struct TestVectorGatherLowering
732     : public PassWrapper<TestVectorGatherLowering,
733                          OperationPass<func::FuncOp>> {
734   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)
735 
736   StringRef getArgument() const final { return "test-vector-gather-lowering"; }
737   StringRef getDescription() const final {
738     return "Test patterns that lower the gather op in the vector conditional "
739            "loads";
740   }
741   void getDependentDialects(DialectRegistry &registry) const override {
742     registry.insert<arith::ArithDialect, func::FuncDialect,
743                     memref::MemRefDialect, scf::SCFDialect,
744                     tensor::TensorDialect, vector::VectorDialect>();
745   }
746 
747   void runOnOperation() override {
748     RewritePatternSet patterns(&getContext());
749     populateVectorGatherLoweringPatterns(patterns);
750     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
751   }
752 };
753 
754 struct TestFoldArithExtensionIntoVectorContractPatterns
755     : public PassWrapper<TestFoldArithExtensionIntoVectorContractPatterns,
756                          OperationPass<func::FuncOp>> {
757   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
758       TestFoldArithExtensionIntoVectorContractPatterns)
759 
760   StringRef getArgument() const final {
761     return "test-fold-arith-extf-into-vector-contract-patterns";
762   }
763   StringRef getDescription() const final {
764     return "Test patterns that fold arithmetic extension ops into vector "
765            "contract ops";
766   }
767 
768   void getDependentDialects(DialectRegistry &registry) const override {
769     registry.insert<arith::ArithDialect, func::FuncDialect, nvgpu::NVGPUDialect,
770                     memref::MemRefDialect, scf::SCFDialect,
771                     tensor::TensorDialect, vector::VectorDialect>();
772   }
773 
774   void runOnOperation() override {
775     RewritePatternSet patterns(&getContext());
776     populateFoldArithExtensionPatterns(patterns);
777     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
778   }
779 };
780 } // namespace
781 
782 namespace mlir {
783 namespace test {
784 void registerTestVectorLowerings() {
785   PassRegistration<TestVectorToVectorLowering>();
786 
787   PassRegistration<TestVectorContractionPrepareForMMTLowering>();
788 
789   PassRegistration<TestVectorUnrollingPatterns>();
790 
791   PassRegistration<TestVectorTransferUnrollingPatterns>();
792 
793   PassRegistration<TestScalarVectorTransferLoweringPatterns>();
794 
795   PassRegistration<TestVectorTransferOpt>();
796 
797   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
798 
799   PassRegistration<TestSinkVectorBroadcast>();
800 
801   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
802 
803   PassRegistration<TestVectorChainedReductionFoldingPatterns>();
804 
805   PassRegistration<TestFlattenVectorTransferPatterns>();
806 
807   PassRegistration<TestVectorScanLowering>();
808 
809   PassRegistration<TestVectorDistribution>();
810 
811   PassRegistration<TestVectorExtractStridedSliceLowering>();
812 
813   PassRegistration<TestVectorBreakDownBitCast>();
814 
815   PassRegistration<TestCreateVectorBroadcast>();
816 
817   PassRegistration<TestVectorGatherLowering>();
818 
819   PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();
820 }
821 } // namespace test
822 } // namespace mlir
823