MLIR  17.0.0git
VectorDistribution.h
Go to the documentation of this file.
1 //===- VectorDistribution.h - Vector distribution patterns --*- C++------*-===//
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 #ifndef MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORDISTRIBUTION_H_
10 #define MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORDISTRIBUTION_H_
11 
13 
14 namespace mlir {
15 class RewritePatternSet;
16 namespace vector {
17 
19  /// Lamdba function to let users allocate memory needed for the lowering of
20  /// WarpExecuteOnLane0Op.
21  /// The function needs to return an allocation that the lowering can use as
22  /// temporary memory. The allocation needs to match the shape of the type (the
23  /// type may be VectorType or a scalar) and be availble for the current warp.
24  /// If there are several warps running in parallel the allocation needs to be
25  /// split so that each warp has its own allocation.
27  std::function<Value(Location, OpBuilder &, WarpExecuteOnLane0Op, Type)>;
29 
30  /// Lamdba function to let user emit operation to syncronize all the thread
31  /// within a warp. After this operation all the threads can see any memory
32  /// written before the operation.
34  std::function<void(Location, OpBuilder &, WarpExecuteOnLane0Op)>;
36 };
37 
39  RewritePatternSet &patterns,
41  PatternBenefit benefit = 1);
42 
43 using DistributionMapFn = std::function<AffineMap(Value)>;
44 
45 /// Distribute transfer_write ops based on the affine map returned by
46 /// `distributionMapFn`.
47 /// Example:
48 /// ```
49 /// %0 = vector.warp_execute_on_lane_0(%id){
50 /// ...
51 /// vector.transfer_write %v, %A[%c0] : vector<32xf32>, memref<128xf32>
52 /// vector.yield
53 /// }
54 /// ```
55 /// To
56 /// ```
57 /// %r:3 = vector.warp_execute_on_lane_0(%id) -> (vector<1xf32>) {
58 /// ...
59 /// vector.yield %v : vector<32xf32>
60 /// }
61 /// vector.transfer_write %v, %A[%id] : vector<1xf32>, memref<128xf32>
62 void populateDistributeTransferWriteOpPatterns(
63  RewritePatternSet &patterns, const DistributionMapFn &distributionMapFn,
64  PatternBenefit benefit = 1);
65 
66 /// Move scalar operations with no dependency on the warp op outside of the
67 /// region.
68 void moveScalarUniformCode(WarpExecuteOnLane0Op op);
69 
70 /// Lambda signature to compute a warp shuffle of a given value of a given lane
71 /// within a given warp size.
72 using WarpShuffleFromIdxFn =
73  std::function<Value(Location, OpBuilder &b, Value, Value, int64_t)>;
74 
75 /// Collect patterns to propagate warp distribution. `distributionMapFn` is used
76 /// to decide how a value should be distributed when this cannot be inferred
77 /// from its uses.
78 void populatePropagateWarpVectorDistributionPatterns(
79  RewritePatternSet &pattern, const DistributionMapFn &distributionMapFn,
80  const WarpShuffleFromIdxFn &warpShuffleFromIdxFn,
81  PatternBenefit benefit = 1);
82 
83 /// Lambda signature to compute a reduction of a distributed value for the given
84 /// reduction kind and size.
85 using DistributedReductionFn =
86  std::function<Value(Location, OpBuilder &, Value, CombiningKind, uint32_t)>;
87 
88 /// Collect patterns to distribute vector reduction ops using given lamdba to
89 /// distribute reduction op.
90 void populateDistributeReduction(
91  RewritePatternSet &pattern,
92  const DistributedReductionFn &distributedReductionFn,
93  PatternBenefit benefit = 1);
94 
95 } // namespace vector
96 } // namespace mlir
97 #endif // MLIR_DIALECT_VECTOR_TRANSFORMS_VECTORDISTRIBUTION_H_
static llvm::ManagedStatic< PassManagerOptions > options
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:44
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
This class helps build Operations.
Definition: Builders.h:202
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:33
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:93
void populateWarpExecuteOnLane0OpToScfForPattern(RewritePatternSet &patterns, const WarpExecuteOnLane0LoweringOptions &options, PatternBenefit benefit=1)
std::function< AffineMap(Value)> DistributionMapFn
This header declares functions that assit transformations in the MemRef dialect.
std::function< Value(Location, OpBuilder &, WarpExecuteOnLane0Op, Type)> WarpAllocationFn
Lamdba function to let users allocate memory needed for the lowering of WarpExecuteOnLane0Op.
std::function< void(Location, OpBuilder &, WarpExecuteOnLane0Op)> WarpSyncronizationFn
Lamdba function to let user emit operation to syncronize all the thread within a warp.