xref: /llvm-project/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp (revision 09dfc5713d7e2342bea4c8447d1ed76c85eb8225) 
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)applyPatternsGreedily(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());
100     }
101     if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
102       auto insert = writeOp.getVector().getDefiningOp<InsertStridedSliceOp>();
103       if (!insert)
104         return std::nullopt;
105       ArrayRef<int64_t> shape = insert.getSourceVectorType().getShape();
106       return SmallVector<int64_t>(shape);
107     }
108     return std::nullopt;
109   }
110 
111   static LogicalResult filter(Operation *op) {
112     return success(isa<arith::AddFOp, arith::SelectOp, arith::CmpFOp,
113                        ContractionOp, TransferReadOp, TransferWriteOp>(op));
114   }
115 };
116 
117 struct TestVectorContractionPrepareForMMTLowering
118     : public PassWrapper<TestVectorContractionPrepareForMMTLowering,
119                          OperationPass<func::FuncOp>> {
120   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
121       TestVectorContractionPrepareForMMTLowering)
122 
123   StringRef getArgument() const final {
124     return "test-vector-contraction-prepare-for-mmt-lowering";
125   }
126   StringRef getDescription() const final {
127     return "Test vector.contraction matmul canonicalization for MMT lowering.";
128   }
129   TestVectorContractionPrepareForMMTLowering() = default;
130 
131   void getDependentDialects(DialectRegistry &registry) const override {
132     registry.insert<affine::AffineDialect, arith::ArithDialect,
133                     vector::VectorDialect>();
134   }
135 
136   void runOnOperation() override {
137     MLIRContext *ctx = &getContext();
138     RewritePatternSet patterns(ctx);
139     vector::populateVectorContractCanonicalizeMatmulToMMT(patterns);
140     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
141   }
142 };
143 
144 struct TestVectorUnrollingPatterns
145     : public PassWrapper<TestVectorUnrollingPatterns,
146                          OperationPass<func::FuncOp>> {
147   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorUnrollingPatterns)
148 
149   StringRef getArgument() const final {
150     return "test-vector-unrolling-patterns";
151   }
152   StringRef getDescription() const final {
153     return "Test lowering patterns to unroll contract ops in the vector "
154            "dialect";
155   }
156   TestVectorUnrollingPatterns() = default;
157   TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass)
158       : PassWrapper(pass) {}
159   void runOnOperation() override {
160     MLIRContext *ctx = &getContext();
161     RewritePatternSet patterns(ctx);
162     populateVectorUnrollPatterns(
163         patterns, UnrollVectorOptions()
164                       .setNativeShape(ArrayRef<int64_t>{2, 2})
165                       .setFilterConstraint([](Operation *op) {
166                         return success(isa<arith::AddFOp, vector::FMAOp,
167                                            vector::MultiDimReductionOp>(op));
168                       }));
169     populateVectorUnrollPatterns(
170         patterns, UnrollVectorOptions()
171                       .setNativeShape(ArrayRef<int64_t>{2})
172                       .setFilterConstraint([](Operation *op) {
173                         return success(isa<vector::ReductionOp>(op));
174                       }));
175     populateVectorUnrollPatterns(
176         patterns, UnrollVectorOptions()
177                       .setNativeShape(ArrayRef<int64_t>{1, 3, 4, 2})
178                       .setFilterConstraint([](Operation *op) {
179                         return success(isa<vector::TransposeOp>(op));
180                       }));
181 
182     if (unrollBasedOnType) {
183       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
184           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
185         vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
186         SmallVector<int64_t> nativeShape(contractOp.getIteratorTypes().size(),
187                                          4);
188         Type lhsType = contractOp.getLhsType().getElementType();
189         nativeShape[nativeShape.size() - 1] = lhsType.isF16() ? 4 : 2;
190         return nativeShape;
191       };
192 
193       UnrollVectorOptions opts;
194       opts.setNativeShapeFn(nativeShapeFn)
195           .setFilterConstraint(
196               [](Operation *op) { return success(isa<ContractionOp>(op)); });
197 
198       if (!unrollOrder.empty()) {
199         opts.setUnrollTraversalOrderFn(
200             [this](Operation *op) -> std::optional<SmallVector<int64_t>> {
201               vector::ContractionOp contractOp =
202                   cast<vector::ContractionOp>(op);
203               if (contractOp.getIteratorTypes().size() == unrollOrder.size())
204                 return SmallVector<int64_t>(unrollOrder.begin(),
205                                             unrollOrder.end());
206               return std::nullopt;
207             });
208       }
209       populateVectorUnrollPatterns(patterns, opts);
210     } else {
211       auto nativeShapeFn =
212           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
213         auto contractOp = dyn_cast<ContractionOp>(op);
214         if (!contractOp)
215           return std::nullopt;
216         return SmallVector<int64_t>(contractOp.getIteratorTypes().size(), 2);
217       };
218       populateVectorUnrollPatterns(patterns,
219                                    UnrollVectorOptions()
220                                        .setNativeShapeFn(nativeShapeFn)
221                                        .setFilterConstraint([](Operation *op) {
222                                          return success(isa<ContractionOp>(op));
223                                        }));
224     }
225     populateVectorToVectorCanonicalizationPatterns(patterns);
226     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
227   }
228 
229   ListOption<int64_t> unrollOrder{*this, "unroll-order",
230                                   llvm::cl::desc("set the unroll order")};
231 
232   Option<bool> unrollBasedOnType{
233       *this, "unroll-based-on-type",
234       llvm::cl::desc("Set the unroll factor based on type of the operation"),
235       llvm::cl::init(false)};
236 };
237 
238 struct TestVectorTransferUnrollingPatterns
239     : public PassWrapper<TestVectorTransferUnrollingPatterns,
240                          OperationPass<func::FuncOp>> {
241   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
242       TestVectorTransferUnrollingPatterns)
243 
244   TestVectorTransferUnrollingPatterns() = default;
245   TestVectorTransferUnrollingPatterns(
246       const TestVectorTransferUnrollingPatterns &pass)
247       : PassWrapper(pass) {}
248 
249   void getDependentDialects(DialectRegistry &registry) const override {
250     registry.insert<affine::AffineDialect>();
251   }
252   StringRef getArgument() const final {
253     return "test-vector-transfer-unrolling-patterns";
254   }
255   StringRef getDescription() const final {
256     return "Test lowering patterns to unroll transfer ops in the vector "
257            "dialect";
258   }
259   void runOnOperation() override {
260     MLIRContext *ctx = &getContext();
261     RewritePatternSet patterns(ctx);
262     UnrollVectorOptions opts;
263     opts.setNativeShape(ArrayRef<int64_t>{2, 2})
264         .setFilterConstraint([](Operation *op) {
265           return success(isa<vector::TransferReadOp, vector::TransferWriteOp,
266                              vector::GatherOp>(op));
267         });
268     if (reverseUnrollOrder.getValue()) {
269       opts.setUnrollTraversalOrderFn(
270           [](Operation *op) -> std::optional<SmallVector<int64_t>> {
271             int64_t numLoops = 0;
272             if (auto readOp = dyn_cast<vector::TransferReadOp>(op))
273               numLoops = readOp.getVectorType().getRank();
274             else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op))
275               numLoops = writeOp.getVectorType().getRank();
276             else if (auto gatherOp = dyn_cast<vector::GatherOp>(op))
277               numLoops = gatherOp.getVectorType().getRank();
278             else
279               return std::nullopt;
280             auto order = llvm::reverse(llvm::seq<int64_t>(0, numLoops));
281             return llvm::to_vector(order);
282           });
283     }
284     populateVectorUnrollPatterns(patterns, opts);
285     populateVectorToVectorCanonicalizationPatterns(patterns);
286     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
287   }
288 
289   Option<bool> reverseUnrollOrder{
290       *this, "reverse-unroll-order",
291       llvm::cl::desc(
292           "reverse the order of unrolling of vector transfer operations"),
293       llvm::cl::init(false)};
294 };
295 
296 struct TestScalarVectorTransferLoweringPatterns
297     : public PassWrapper<TestScalarVectorTransferLoweringPatterns,
298                          OperationPass<func::FuncOp>> {
299   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
300       TestScalarVectorTransferLoweringPatterns)
301 
302   TestScalarVectorTransferLoweringPatterns() = default;
303   TestScalarVectorTransferLoweringPatterns(
304       const TestScalarVectorTransferLoweringPatterns &pass)
305       : PassWrapper(pass) {}
306 
307   StringRef getArgument() const final {
308     return "test-scalar-vector-transfer-lowering";
309   }
310   StringRef getDescription() const final {
311     return "Test lowering of scalar vector transfers to memref loads/stores.";
312   }
313 
314   void getDependentDialects(DialectRegistry &registry) const override {
315     registry.insert<affine::AffineDialect, memref::MemRefDialect,
316                     tensor::TensorDialect, vector::VectorDialect>();
317   }
318 
319   Option<bool> allowMultipleUses{
320       *this, "allow-multiple-uses",
321       llvm::cl::desc("Fold transfer operations with multiple uses"),
322       llvm::cl::init(false)};
323 
324   void runOnOperation() override {
325     MLIRContext *ctx = &getContext();
326     RewritePatternSet patterns(ctx);
327     vector::populateScalarVectorTransferLoweringPatterns(
328         patterns, /*benefit=*/1, allowMultipleUses.getValue());
329     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
330   }
331 };
332 
333 struct TestVectorTransferOpt
334     : public PassWrapper<TestVectorTransferOpt, OperationPass<func::FuncOp>> {
335   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorTransferOpt)
336 
337   StringRef getArgument() const final { return "test-vector-transferop-opt"; }
338   StringRef getDescription() const final {
339     return "Test optimization transformations for transfer ops";
340   }
341   void runOnOperation() override {
342     IRRewriter rewriter(&getContext());
343     transferOpflowOpt(rewriter, getOperation());
344   }
345 };
346 
347 struct TestVectorTransferCollapseInnerMostContiguousDims
348     : public PassWrapper<TestVectorTransferCollapseInnerMostContiguousDims,
349                          OperationPass<func::FuncOp>> {
350   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
351       TestVectorTransferCollapseInnerMostContiguousDims)
352 
353   TestVectorTransferCollapseInnerMostContiguousDims() = default;
354   TestVectorTransferCollapseInnerMostContiguousDims(
355       const TestVectorTransferCollapseInnerMostContiguousDims &pass) = default;
356 
357   void getDependentDialects(DialectRegistry &registry) const override {
358     registry.insert<memref::MemRefDialect, affine::AffineDialect>();
359   }
360 
361   StringRef getArgument() const final {
362     return "test-vector-transfer-collapse-inner-most-dims";
363   }
364 
365   StringRef getDescription() const final {
366     return "Test lowering patterns that reduces the rank of the vector "
367            "transfer memory and vector operands.";
368   }
369 
370   void runOnOperation() override {
371     RewritePatternSet patterns(&getContext());
372     populateVectorTransferCollapseInnerMostContiguousDimsPatterns(patterns);
373     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
374   }
375 };
376 
377 struct TestVectorSinkPatterns
378     : public PassWrapper<TestVectorSinkPatterns, OperationPass<func::FuncOp>> {
379   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorSinkPatterns)
380 
381   TestVectorSinkPatterns() = default;
382   TestVectorSinkPatterns(const TestVectorSinkPatterns &pass) = default;
383 
384   void getDependentDialects(DialectRegistry &registry) const override {
385     registry.insert<memref::MemRefDialect, affine::AffineDialect>();
386   }
387 
388   StringRef getArgument() const final { return "test-vector-sink-patterns"; }
389 
390   StringRef getDescription() const final {
391     return "Test lowering patterns that eliminate redundant broadcast "
392            "and transpose operations.";
393   }
394 
395   void runOnOperation() override {
396     RewritePatternSet patterns(&getContext());
397     populateSinkVectorOpsPatterns(patterns);
398     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
399   }
400 };
401 
402 struct TestVectorReduceToContractPatternsPatterns
403     : public PassWrapper<TestVectorReduceToContractPatternsPatterns,
404                          OperationPass<func::FuncOp>> {
405   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
406       TestVectorReduceToContractPatternsPatterns)
407 
408   StringRef getArgument() const final {
409     return "test-vector-reduction-to-contract-patterns";
410   }
411   StringRef getDescription() const final {
412     return "Test patterns to convert multireduce op to contract and combine "
413            "broadcast/transpose to contract";
414   }
415   void runOnOperation() override {
416     RewritePatternSet patterns(&getContext());
417     populateVectorReductionToContractPatterns(patterns);
418     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
419   }
420 };
421 
422 struct TestVectorChainedReductionFoldingPatterns
423     : public PassWrapper<TestVectorChainedReductionFoldingPatterns,
424                          OperationPass<func::FuncOp>> {
425   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
426       TestVectorChainedReductionFoldingPatterns)
427 
428   StringRef getArgument() const final {
429     return "test-vector-chained-reduction-folding-patterns";
430   }
431   StringRef getDescription() const final {
432     return "Test patterns to fold chained vector reductions";
433   }
434   void runOnOperation() override {
435     RewritePatternSet patterns(&getContext());
436     populateChainedVectorReductionFoldingPatterns(patterns);
437     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
438   }
439 };
440 
441 struct TestVectorBreakDownReductionPatterns
442     : public PassWrapper<TestVectorBreakDownReductionPatterns,
443                          OperationPass<func::FuncOp>> {
444   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
445       TestVectorBreakDownReductionPatterns)
446 
447   StringRef getArgument() const final {
448     return "test-vector-break-down-reduction-patterns";
449   }
450   StringRef getDescription() const final {
451     return "Test patterns to break down vector reductions into arith "
452            "reductions";
453   }
454   void runOnOperation() override {
455     RewritePatternSet patterns(&getContext());
456     populateBreakDownVectorReductionPatterns(patterns,
457                                              /*maxNumElementsToExtract=*/2);
458     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
459   }
460 };
461 
462 struct TestFlattenVectorTransferPatterns
463     : public PassWrapper<TestFlattenVectorTransferPatterns,
464                          OperationPass<func::FuncOp>> {
465   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
466       TestFlattenVectorTransferPatterns)
467 
468   TestFlattenVectorTransferPatterns() = default;
469   TestFlattenVectorTransferPatterns(
470       const TestFlattenVectorTransferPatterns &pass)
471       : PassWrapper(pass) {}
472 
473   StringRef getArgument() const final {
474     return "test-vector-transfer-flatten-patterns";
475   }
476 
477   StringRef getDescription() const final {
478     return "Test patterns to rewrite contiguous row-major N-dimensional "
479            "vector.transfer_{read,write} ops into 1D transfers";
480   }
481 
482   void getDependentDialects(DialectRegistry &registry) const override {
483     registry.insert<memref::MemRefDialect>();
484     registry.insert<affine::AffineDialect>();
485     registry.insert<vector::VectorDialect>();
486   }
487 
488   Option<unsigned> targetVectorBitwidth{
489       *this, "target-vector-bitwidth",
490       llvm::cl::desc(
491           "Minimum vector bitwidth to enable the flattening transformation. "
492           "For scalable vectors this is the base size, i.e. the size "
493           "corresponding to vscale=1."),
494       llvm::cl::init(std::numeric_limits<unsigned>::max())};
495 
496   void runOnOperation() override {
497     RewritePatternSet patterns(&getContext());
498     populateFlattenVectorTransferPatterns(patterns, targetVectorBitwidth);
499     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
500   }
501 };
502 
503 struct TestVectorScanLowering
504     : public PassWrapper<TestVectorScanLowering, OperationPass<func::FuncOp>> {
505   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorScanLowering)
506 
507   StringRef getArgument() const final { return "test-vector-scan-lowering"; }
508   StringRef getDescription() const final {
509     return "Test lowering patterns that lower the scan op in the vector "
510            "dialect";
511   }
512   void runOnOperation() override {
513     RewritePatternSet patterns(&getContext());
514     populateVectorScanLoweringPatterns(patterns);
515     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
516   }
517 };
518 
519 /// Allocate shared memory for a single warp to test lowering of
520 /// WarpExecuteOnLane0Op.
521 static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
522                                         gpu::WarpExecuteOnLane0Op warpOp,
523                                         Type type) {
524   static constexpr int64_t kSharedMemorySpace = 3;
525   // Compute type of shared memory buffer.
526   MemRefType memrefType;
527   if (auto vectorType = dyn_cast<VectorType>(type)) {
528     memrefType =
529         MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},
530                         kSharedMemorySpace);
531   } else {
532     memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);
533   }
534 
535   // Get symbol table holding all shared memory globals.
536   ModuleOp moduleOp = warpOp->getParentOfType<ModuleOp>();
537   SymbolTable symbolTable(moduleOp);
538 
539   // Create a pretty name.
540   SmallString<64> buf;
541   llvm::raw_svector_ostream os(buf);
542   interleave(memrefType.getShape(), os, "x");
543   os << "x" << memrefType.getElementType();
544   std::string symbolName = (Twine("__shared_") + os.str()).str();
545 
546   auto ip = builder.saveInsertionPoint();
547   builder.setInsertionPoint(moduleOp);
548   auto global = builder.create<memref::GlobalOp>(
549       loc,
550       /*sym_name=*/symbolName,
551       /*sym_visibility=*/builder.getStringAttr("private"),
552       /*type=*/memrefType,
553       /*initial_value=*/Attribute(),
554       /*constant=*/false,
555       /*alignment=*/IntegerAttr());
556   symbolTable.insert(global);
557   // The symbol table inserts at the end of the module, but globals are a bit
558   // nicer if they are at the beginning.
559   global->moveBefore(&moduleOp.front());
560 
561   builder.restoreInsertionPoint(ip);
562   return builder.create<memref::GetGlobalOp>(loc, memrefType, symbolName);
563 }
564 
565 static Value warpReduction(Location loc, OpBuilder &builder, Value input,
566                            CombiningKind kind, uint32_t size) {
567   // First reduce on a single thread to get per lane reduction value.
568   Value laneVal = builder.create<vector::ReductionOp>(loc, kind, input);
569   // Parallel reduction using butterfly shuffles.
570   for (uint64_t i = 1; i < size; i <<= 1) {
571     Value shuffled = builder
572                          .create<gpu::ShuffleOp>(loc, laneVal, i,
573                                                  /*width=*/size,
574                                                  /*mode=*/gpu::ShuffleMode::XOR)
575                          .getShuffleResult();
576     laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);
577   }
578   return laneVal;
579 }
580 
581 struct TestVectorDistribution
582     : public PassWrapper<TestVectorDistribution, OperationPass<func::FuncOp>> {
583   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorDistribution)
584 
585   void getDependentDialects(DialectRegistry &registry) const override {
586     registry
587         .insert<vector::VectorDialect, scf::SCFDialect, memref::MemRefDialect,
588                 gpu::GPUDialect, affine::AffineDialect>();
589   }
590 
591   StringRef getArgument() const final { return "test-vector-warp-distribute"; }
592   StringRef getDescription() const final {
593     return "Test vector warp distribute transformation and lowering patterns";
594   }
595   TestVectorDistribution() = default;
596   TestVectorDistribution(const TestVectorDistribution &pass)
597       : PassWrapper(pass) {}
598 
599   Option<bool> warpOpToSCF{
600       *this, "rewrite-warp-ops-to-scf-if",
601       llvm::cl::desc("Lower vector.warp_execute_on_lane0 to scf.if op"),
602       llvm::cl::init(false)};
603 
604   Option<bool> distributeTransferWriteOps{
605       *this, "distribute-transfer-write",
606       llvm::cl::desc("Test distribution of transfer write"),
607       llvm::cl::init(false)};
608 
609   Option<unsigned> maxTransferWriteElements{
610       *this, "max-transfer-write-elements",
611       llvm::cl::desc("Maximum number of transfer write elements to distribute"),
612       llvm::cl::init(1)};
613 
614   Option<bool> hoistUniform{*this, "hoist-uniform",
615                             llvm::cl::desc("Test hoist uniform"),
616                             llvm::cl::init(false)};
617 
618   Option<bool> propagateDistribution{
619       *this, "propagate-distribution",
620       llvm::cl::desc("Test distribution propagation"), llvm::cl::init(false)};
621 
622   void runOnOperation() override {
623     RewritePatternSet patterns(&getContext());
624 
625     getOperation().walk([&](Operation *op) {
626       if (auto warpOp = dyn_cast<gpu::WarpExecuteOnLane0Op>(op)) {
627         if (hoistUniform) {
628           moveScalarUniformCode(warpOp);
629         }
630         WalkResult::interrupt();
631       }
632     });
633     MLIRContext *ctx = &getContext();
634     auto distributionFn = [](Value val) {
635       // Create an identity dim map of the same rank as the vector.
636       VectorType vecType = dyn_cast<VectorType>(val.getType());
637       int64_t vecRank = vecType ? vecType.getRank() : 0;
638       OpBuilder builder(val.getContext());
639       if (vecRank == 0)
640         return AffineMap::get(val.getContext());
641       return AffineMap::getMultiDimIdentityMap(vecRank, val.getContext());
642     };
643     auto shuffleFn = [](Location loc, OpBuilder &builder, Value val,
644                         Value srcIdx, int64_t warpSz) {
645       assert((val.getType().isF32() || val.getType().isInteger(32)) &&
646              "unsupported shuffle type");
647       Type i32Type = builder.getIntegerType(32);
648       Value srcIdxI32 =
649           builder.create<arith::IndexCastOp>(loc, i32Type, srcIdx);
650       Value warpSzI32 = builder.create<arith::ConstantOp>(
651           loc, builder.getIntegerAttr(i32Type, warpSz));
652       Value result = builder
653                          .create<gpu::ShuffleOp>(loc, val, srcIdxI32, warpSzI32,
654                                                  gpu::ShuffleMode::IDX)
655                          .getResult(0);
656       return result;
657     };
658     if (distributeTransferWriteOps && propagateDistribution) {
659       RewritePatternSet patterns(ctx);
660       vector::populatePropagateWarpVectorDistributionPatterns(
661           patterns, distributionFn, shuffleFn, /*benefit=*/1,
662           /*readBenefit=*/0);
663       vector::populateDistributeReduction(patterns, warpReduction, 1);
664       populateDistributeTransferWriteOpPatterns(patterns, distributionFn, 2);
665       (void)applyPatternsGreedily(getOperation(), std::move(patterns));
666     } else if (distributeTransferWriteOps) {
667       RewritePatternSet patterns(ctx);
668       populateDistributeTransferWriteOpPatterns(patterns, distributionFn,
669                                                 maxTransferWriteElements);
670       (void)applyPatternsGreedily(getOperation(), std::move(patterns));
671     } else if (propagateDistribution) {
672       RewritePatternSet patterns(ctx);
673       vector::populatePropagateWarpVectorDistributionPatterns(
674           patterns, distributionFn, shuffleFn);
675       vector::populateDistributeReduction(patterns, warpReduction);
676       (void)applyPatternsGreedily(getOperation(), std::move(patterns));
677     }
678     WarpExecuteOnLane0LoweringOptions options;
679     options.warpAllocationFn = allocateGlobalSharedMemory;
680     options.warpSyncronizationFn = [](Location loc, OpBuilder &builder,
681                                       gpu::WarpExecuteOnLane0Op warpOp) {
682       builder.create<gpu::BarrierOp>(loc);
683     };
684     // Test on one pattern in isolation.
685     if (warpOpToSCF) {
686       populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);
687       (void)applyPatternsGreedily(getOperation(), std::move(patterns));
688       return;
689     }
690   }
691 };
692 
693 struct TestVectorExtractStridedSliceLowering
694     : public PassWrapper<TestVectorExtractStridedSliceLowering,
695                          OperationPass<func::FuncOp>> {
696   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
697       TestVectorExtractStridedSliceLowering)
698 
699   StringRef getArgument() const final {
700     return "test-vector-extract-strided-slice-lowering";
701   }
702   StringRef getDescription() const final {
703     return "Test lowering patterns that converts vector.extract_strided_slice "
704            "into a chain of vector.extract and vector.insert ops";
705   }
706   void runOnOperation() override {
707     RewritePatternSet patterns(&getContext());
708     populateVectorExtractStridedSliceToExtractInsertChainPatterns(patterns);
709     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
710   }
711 };
712 
713 struct TestVectorBreakDownBitCast
714     : public PassWrapper<TestVectorBreakDownBitCast,
715                          OperationPass<func::FuncOp>> {
716   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBreakDownBitCast)
717 
718   StringRef getArgument() const final {
719     return "test-vector-break-down-bitcast";
720   }
721   StringRef getDescription() const final {
722     return "Test pattern that breaks down vector.bitcast ops ";
723   }
724   void runOnOperation() override {
725     RewritePatternSet patterns(&getContext());
726     populateBreakDownVectorBitCastOpPatterns(patterns, [](BitCastOp op) {
727       return op.getSourceVectorType().getShape().back() > 4;
728     });
729     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
730   }
731 };
732 
733 struct TestCreateVectorBroadcast
734     : public PassWrapper<TestCreateVectorBroadcast,
735                          OperationPass<func::FuncOp>> {
736   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestCreateVectorBroadcast)
737 
738   StringRef getArgument() const final { return "test-create-vector-broadcast"; }
739   StringRef getDescription() const final {
740     return "Test optimization transformations for transfer ops";
741   }
742   void getDependentDialects(DialectRegistry &registry) const override {
743     registry.insert<vector::VectorDialect>();
744   }
745 
746   void runOnOperation() override {
747     getOperation()->walk([](Operation *op) {
748       if (op->getName().getStringRef() != "test_create_broadcast")
749         return;
750       auto targetShape =
751           cast<VectorType>(op->getResult(0).getType()).getShape();
752       auto arrayAttr =
753           cast<DenseI64ArrayAttr>(op->getDiscardableAttr("broadcast_dims"))
754               .asArrayRef();
755       llvm::SetVector<int64_t> broadcastedDims;
756       broadcastedDims.insert(arrayAttr.begin(), arrayAttr.end());
757       OpBuilder b(op);
758       Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(
759           b, op->getOperand(0), targetShape, broadcastedDims);
760       op->getResult(0).replaceAllUsesWith(bcast);
761       op->erase();
762     });
763   }
764 };
765 
766 struct TestVectorGatherLowering
767     : public PassWrapper<TestVectorGatherLowering,
768                          OperationPass<func::FuncOp>> {
769   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)
770 
771   StringRef getArgument() const final { return "test-vector-gather-lowering"; }
772   StringRef getDescription() const final {
773     return "Test patterns that lower the gather op in the vector conditional "
774            "loads";
775   }
776   void getDependentDialects(DialectRegistry &registry) const override {
777     registry.insert<arith::ArithDialect, func::FuncDialect,
778                     memref::MemRefDialect, scf::SCFDialect,
779                     tensor::TensorDialect, vector::VectorDialect>();
780   }
781 
782   void runOnOperation() override {
783     RewritePatternSet patterns(&getContext());
784     populateVectorGatherLoweringPatterns(patterns);
785     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
786   }
787 };
788 
789 struct TestFoldArithExtensionIntoVectorContractPatterns
790     : public PassWrapper<TestFoldArithExtensionIntoVectorContractPatterns,
791                          OperationPass<func::FuncOp>> {
792   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
793       TestFoldArithExtensionIntoVectorContractPatterns)
794 
795   StringRef getArgument() const final {
796     return "test-fold-arith-extf-into-vector-contract-patterns";
797   }
798   StringRef getDescription() const final {
799     return "Test patterns that fold arithmetic extension ops into vector "
800            "contract ops";
801   }
802 
803   void getDependentDialects(DialectRegistry &registry) const override {
804     registry.insert<arith::ArithDialect, func::FuncDialect, nvgpu::NVGPUDialect,
805                     memref::MemRefDialect, scf::SCFDialect,
806                     tensor::TensorDialect, vector::VectorDialect>();
807   }
808 
809   void runOnOperation() override {
810     RewritePatternSet patterns(&getContext());
811     populateFoldArithExtensionPatterns(patterns);
812     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
813   }
814 };
815 
816 struct TestVectorEmulateMaskedLoadStore final
817     : public PassWrapper<TestVectorEmulateMaskedLoadStore,
818                          OperationPass<func::FuncOp>> {
819   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorEmulateMaskedLoadStore)
820 
821   StringRef getArgument() const override {
822     return "test-vector-emulate-masked-load-store";
823   }
824   StringRef getDescription() const override {
825     return "Test patterns that emulate the maskedload/maskedstore op by "
826            " memref.load/store and scf.if";
827   }
828   void getDependentDialects(DialectRegistry &registry) const override {
829     registry
830         .insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
831                 scf::SCFDialect, vector::VectorDialect>();
832   }
833 
834   void runOnOperation() override {
835     RewritePatternSet patterns(&getContext());
836     populateVectorMaskedLoadStoreEmulationPatterns(patterns);
837     (void)applyPatternsGreedily(getOperation(), std::move(patterns));
838   }
839 };
840 
841 struct TestVectorLinearize final
842     : public PassWrapper<TestVectorLinearize, OperationPass<>> {
843   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorLinearize)
844 
845   TestVectorLinearize() = default;
846   TestVectorLinearize(const TestVectorLinearize &pass) : PassWrapper(pass) {}
847 
848   StringRef getArgument() const override { return "test-vector-linearize"; }
849   StringRef getDescription() const override {
850     return "Linearizes ND vectors for N >= 2 into 1D vectors";
851   }
852   void getDependentDialects(DialectRegistry &registry) const override {
853     registry.insert<vector::VectorDialect>();
854   }
855 
856   Option<unsigned> targetVectorBitwidth{
857       *this, "target-vector-bitwidth",
858       llvm::cl::desc(
859           "Minimum vector bitwidth to enable the flattening transformation"),
860       llvm::cl::init(std::numeric_limits<unsigned>::max())};
861   void runOnOperation() override {
862     auto *context = &getContext();
863 
864     TypeConverter typeConverter;
865     RewritePatternSet patterns(context);
866     ConversionTarget target(*context);
867 
868     vector::populateVectorLinearizeTypeConversionsAndLegality(
869         typeConverter, patterns, target, targetVectorBitwidth);
870     vector::populateVectorLinearizeShuffleLikeOpsPatterns(
871         typeConverter, patterns, target, targetVectorBitwidth);
872     if (failed(applyPartialConversion(getOperation(), target,
873                                       std::move(patterns))))
874       return signalPassFailure();
875   }
876 };
877 
878 struct TestEliminateVectorMasks
879     : public PassWrapper<TestEliminateVectorMasks,
880                          OperationPass<func::FuncOp>> {
881   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestEliminateVectorMasks)
882 
883   TestEliminateVectorMasks() = default;
884   TestEliminateVectorMasks(const TestEliminateVectorMasks &pass)
885       : PassWrapper(pass) {}
886 
887   Option<unsigned> vscaleMin{
888       *this, "vscale-min", llvm::cl::desc("Minimum possible value of vscale."),
889       llvm::cl::init(1)};
890   Option<unsigned> vscaleMax{
891       *this, "vscale-max", llvm::cl::desc("Maximum possible value of vscale."),
892       llvm::cl::init(16)};
893 
894   StringRef getArgument() const final { return "test-eliminate-vector-masks"; }
895   StringRef getDescription() const final {
896     return "Test eliminating vector masks";
897   }
898   void runOnOperation() override {
899     IRRewriter rewriter(&getContext());
900     eliminateVectorMasks(rewriter, getOperation(),
901                          VscaleRange{vscaleMin, vscaleMax});
902   }
903 };
904 } // namespace
905 
906 namespace mlir {
907 namespace test {
908 void registerTestVectorLowerings() {
909   PassRegistration<TestVectorToVectorLowering>();
910 
911   PassRegistration<TestVectorContractionPrepareForMMTLowering>();
912 
913   PassRegistration<TestVectorUnrollingPatterns>();
914 
915   PassRegistration<TestVectorTransferUnrollingPatterns>();
916 
917   PassRegistration<TestScalarVectorTransferLoweringPatterns>();
918 
919   PassRegistration<TestVectorTransferOpt>();
920 
921   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
922 
923   PassRegistration<TestVectorSinkPatterns>();
924 
925   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
926 
927   PassRegistration<TestVectorChainedReductionFoldingPatterns>();
928 
929   PassRegistration<TestVectorBreakDownReductionPatterns>();
930 
931   PassRegistration<TestFlattenVectorTransferPatterns>();
932 
933   PassRegistration<TestVectorScanLowering>();
934 
935   PassRegistration<TestVectorDistribution>();
936 
937   PassRegistration<TestVectorExtractStridedSliceLowering>();
938 
939   PassRegistration<TestVectorBreakDownBitCast>();
940 
941   PassRegistration<TestCreateVectorBroadcast>();
942 
943   PassRegistration<TestVectorGatherLowering>();
944 
945   PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();
946 
947   PassRegistration<TestVectorEmulateMaskedLoadStore>();
948 
949   PassRegistration<TestVectorLinearize>();
950 
951   PassRegistration<TestEliminateVectorMasks>();
952 }
953 } // namespace test
954 } // namespace mlir
955