xref: /llvm-project/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp (revision 8063622721d0b2b70e44f6e747eec54cdaec2e76) 
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<unsigned> maxTransferWriteElements{
572       *this, "max-transfer-write-elements",
573       llvm::cl::desc("Maximum number of transfer write elements to distribute"),
574       llvm::cl::init(1)};
575 
576   Option<bool> hoistUniform{*this, "hoist-uniform",
577                             llvm::cl::desc("Test hoist uniform"),
578                             llvm::cl::init(false)};
579 
580   Option<bool> propagateDistribution{
581       *this, "propagate-distribution",
582       llvm::cl::desc("Test distribution propgation"), llvm::cl::init(false)};
583 
584   void runOnOperation() override {
585     RewritePatternSet patterns(&getContext());
586 
587     getOperation().walk([&](Operation *op) {
588       if (auto warpOp = dyn_cast<WarpExecuteOnLane0Op>(op)) {
589         if (hoistUniform) {
590           moveScalarUniformCode(warpOp);
591         }
592         WalkResult::interrupt();
593       }
594     });
595     MLIRContext *ctx = &getContext();
596     auto distributionFn = [](Value val) {
597       // Create a map (d0, d1) -> (d1) to distribute along the inner
598       // dimension. Once we support n-d distribution we can add more
599       // complex cases.
600       VectorType vecType = dyn_cast<VectorType>(val.getType());
601       int64_t vecRank = vecType ? vecType.getRank() : 0;
602       OpBuilder builder(val.getContext());
603       if (vecRank == 0)
604         return AffineMap::get(val.getContext());
605       return AffineMap::get(vecRank, 0, builder.getAffineDimExpr(vecRank - 1));
606     };
607     auto shuffleFn = [](Location loc, OpBuilder &builder, Value val,
608                         Value srcIdx, int64_t warpSz) {
609       assert((val.getType().isF32() || val.getType().isInteger(32)) &&
610              "unsupported shuffle type");
611       Type i32Type = builder.getIntegerType(32);
612       Value srcIdxI32 =
613           builder.create<arith::IndexCastOp>(loc, i32Type, srcIdx);
614       Value warpSzI32 = builder.create<arith::ConstantOp>(
615           loc, builder.getIntegerAttr(i32Type, warpSz));
616       Value result = builder
617                          .create<gpu::ShuffleOp>(loc, val, srcIdxI32, warpSzI32,
618                                                  gpu::ShuffleMode::IDX)
619                          .getResult(0);
620       return result;
621     };
622     if (distributeTransferWriteOps && propagateDistribution) {
623       RewritePatternSet patterns(ctx);
624       vector::populatePropagateWarpVectorDistributionPatterns(
625           patterns, distributionFn, shuffleFn, /*benefit=*/1,
626           /*readBenefit=*/0);
627       vector::populateDistributeReduction(patterns, warpReduction, 1);
628       populateDistributeTransferWriteOpPatterns(patterns, distributionFn, 2);
629       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
630     } else if (distributeTransferWriteOps) {
631       RewritePatternSet patterns(ctx);
632       populateDistributeTransferWriteOpPatterns(patterns, distributionFn,
633                                                 maxTransferWriteElements);
634       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
635     } else if (propagateDistribution) {
636       RewritePatternSet patterns(ctx);
637       vector::populatePropagateWarpVectorDistributionPatterns(
638           patterns, distributionFn, shuffleFn);
639       vector::populateDistributeReduction(patterns, warpReduction);
640       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
641     }
642     WarpExecuteOnLane0LoweringOptions options;
643     options.warpAllocationFn = allocateGlobalSharedMemory;
644     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
645                                       WarpExecuteOnLane0Op warpOp) {
646       builder.create<gpu::BarrierOp>(loc);
647     };
648     // Test on one pattern in isolation.
649     if (warpOpToSCF) {
650       populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);
651       (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
652       return;
653     }
654   }
655 };
656 
657 struct TestVectorExtractStridedSliceLowering
658     : public PassWrapper<TestVectorExtractStridedSliceLowering,
659                          OperationPass<func::FuncOp>> {
660   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
661       TestVectorExtractStridedSliceLowering)
662 
663   StringRef getArgument() const final {
664     return "test-vector-extract-strided-slice-lowering";
665   }
666   StringRef getDescription() const final {
667     return "Test lowering patterns that converts vector.extract_strided_slice "
668            "into a chain of vector.extract and vector.insert ops";
669   }
670   void runOnOperation() override {
671     RewritePatternSet patterns(&getContext());
672     populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);
673     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
674   }
675 };
676 
677 struct TestVectorBreakDownBitCast
678     : public PassWrapper<TestVectorBreakDownBitCast,
679                          OperationPass<func::FuncOp>> {
680   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBreakDownBitCast)
681 
682   StringRef getArgument() const final {
683     return "test-vector-break-down-bitcast";
684   }
685   StringRef getDescription() const final {
686     return "Test pattern that breaks down vector.bitcast ops ";
687   }
688   void runOnOperation() override {
689     RewritePatternSet patterns(&getContext());
690     populateBreakDownVectorBitCastOpPatterns(patterns, [](BitCastOp op) {
691       return op.getSourceVectorType().getShape().back() > 4;
692     });
693     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
694   }
695 };
696 
697 struct TestCreateVectorBroadcast
698     : public PassWrapper<TestCreateVectorBroadcast,
699                          OperationPass<func::FuncOp>> {
700   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)
701 
702   StringRef getArgument() const final { return "test-create-vector-broadcast"; }
703   StringRef getDescription() const final {
704     return "Test optimization transformations for transfer ops";
705   }
706   void getDependentDialects(DialectRegistry &registry) const override {
707     registry.insert<vector::VectorDialect>();
708   }
709 
710   void runOnOperation() override {
711     getOperation()->walk([](Operation *op) {
712       if (op->getName().getStringRef() != "test_create_broadcast")
713         return;
714       auto targetShape =
715           cast<VectorType>(op->getResult(0).getType()).getShape();
716       auto arrayAttr =
717           cast<DenseI64ArrayAttr>(op->getDiscardableAttr("broadcast_dims"))
718               .asArrayRef();
719       llvm::SetVector<int64_t> broadcastedDims;
720       broadcastedDims.insert(arrayAttr.begin(), arrayAttr.end());
721       OpBuilder b(op);
722       Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(
723           b, op->getOperand(0), targetShape, broadcastedDims);
724       op->getResult(0).replaceAllUsesWith(bcast);
725       op->erase();
726     });
727   }
728 };
729 
730 struct TestVectorGatherLowering
731     : public PassWrapper<TestVectorGatherLowering,
732                          OperationPass<func::FuncOp>> {
733   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)
734 
735   StringRef getArgument() const final { return "test-vector-gather-lowering"; }
736   StringRef getDescription() const final {
737     return "Test patterns that lower the gather op in the vector conditional "
738            "loads";
739   }
740   void getDependentDialects(DialectRegistry &registry) const override {
741     registry.insert<arith::ArithDialect, func::FuncDialect,
742                     memref::MemRefDialect, scf::SCFDialect,
743                     tensor::TensorDialect, vector::VectorDialect>();
744   }
745 
746   void runOnOperation() override {
747     RewritePatternSet patterns(&getContext());
748     populateVectorGatherLoweringPatterns(patterns);
749     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
750   }
751 };
752 
753 struct TestFoldArithExtensionIntoVectorContractPatterns
754     : public PassWrapper<TestFoldArithExtensionIntoVectorContractPatterns,
755                          OperationPass<func::FuncOp>> {
756   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
757       TestFoldArithExtensionIntoVectorContractPatterns)
758 
759   StringRef getArgument() const final {
760     return "test-fold-arith-extf-into-vector-contract-patterns";
761   }
762   StringRef getDescription() const final {
763     return "Test patterns that fold arithmetic extension ops into vector "
764            "contract ops";
765   }
766 
767   void getDependentDialects(DialectRegistry &registry) const override {
768     registry.insert<arith::ArithDialect, func::FuncDialect, nvgpu::NVGPUDialect,
769                     memref::MemRefDialect, scf::SCFDialect,
770                     tensor::TensorDialect, vector::VectorDialect>();
771   }
772 
773   void runOnOperation() override {
774     RewritePatternSet patterns(&getContext());
775     populateFoldArithExtensionPatterns(patterns);
776     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
777   }
778 };
779 } // namespace
780 
781 namespace mlir {
782 namespace test {
783 void registerTestVectorLowerings() {
784   PassRegistration<TestVectorToVectorLowering>();
785 
786   PassRegistration<TestVectorContractionPrepareForMMTLowering>();
787 
788   PassRegistration<TestVectorUnrollingPatterns>();
789 
790   PassRegistration<TestVectorTransferUnrollingPatterns>();
791 
792   PassRegistration<TestScalarVectorTransferLoweringPatterns>();
793 
794   PassRegistration<TestVectorTransferOpt>();
795 
796   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
797 
798   PassRegistration<TestSinkVectorBroadcast>();
799 
800   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
801 
802   PassRegistration<TestVectorChainedReductionFoldingPatterns>();
803 
804   PassRegistration<TestFlattenVectorTransferPatterns>();
805 
806   PassRegistration<TestVectorScanLowering>();
807 
808   PassRegistration<TestVectorDistribution>();
809 
810   PassRegistration<TestVectorExtractStridedSliceLowering>();
811 
812   PassRegistration<TestVectorBreakDownBitCast>();
813 
814   PassRegistration<TestCreateVectorBroadcast>();
815 
816   PassRegistration<TestVectorGatherLowering>();
817 
818   PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();
819 }
820 } // namespace test
821 } // namespace mlir
822