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