10d6b83d022ff9ce630c123931f134fe6d02c2b4
[lldb.git] / mlir / include / mlir / Transforms / LoopFusionUtils.h
1 //===- LoopFusionUtils.h - Loop fusion utilities ----------------*- 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 // This header file defines prototypes for various loop fusion utility
10 // methods: these are not passes by themselves but are used either by passes,
11 // optimization sequences, or in turn by other transformation utilities.
12 //
13 //===----------------------------------------------------------------------===//
14
15 #ifndef MLIR_TRANSFORMS_LOOP_FUSION_UTILS_H
16 #define MLIR_TRANSFORMS_LOOP_FUSION_UTILS_H
17
18 #include "mlir/IR/Value.h"
19 #include "mlir/Support/LLVM.h"
20 #include "llvm/ADT/DenseMap.h"
21 #include "llvm/ADT/SmallVector.h"
22
23 namespace mlir {
24 class AffineForOp;
25 struct ComputationSliceState;
26 class Operation;
27
28 // TODO: Extend this module to include utility functions for querying fusion
29 // cost/storage reduction, and for performing the loop fusion transformation.
30
31 struct FusionResult {
32   enum ResultEnum {
33     Success,
34     FailPrecondition,     // Failed precondition for fusion. (e.g. same block).
35     FailBlockDependence,  // Fusion would violate another dependence in block.
36     FailFusionDependence, // Fusion would reverse dependences between loops.
37     FailComputationSlice, // Unable to compute src loop computation slice.
38   } value;
39   FusionResult(ResultEnum v) : value(v) {}
40 };
41
42 /// Describes the fusion strategy to be used in the Affine loop fusion
43 /// utilities. Currently, it is used to specialized the loop fusion utilities
44 /// with the assumptions made in the AffineLoopFusion pass for producer-consumer
45 /// and sibling fusion, while sharing a single implementation. The latter
46 /// strategies are also limited to scenarios where a single memref is involved
47 /// in the producer-consume or sibling relationship between the candidate
48 /// loops. We use 'memref' to keep track of such a memref.
49 // TODO: Remove 'memref' when we support more generic scenarios.
50 // TODO: Generalize utilities so that producer-consumer and sibling fusion
51 // strategies can be used without the assumptions made in the AffineLoopFusion
52 // pass.
53 class FusionStrategy {
54 public:
55   enum StrategyEnum {
56     // Generic loop fusion: Arbitrary loops are considered for fusion. No
57     // assumptions about a specific fusion strategy from AffineLoopFusion pass
58     // are made.
59     // TODO: Generic fusion is not fully implemented by fusion utilities yet.
60     // It should only be used for testing.
61     Generic,
62     // Producer-consumer fusion: Only loops with a producer-consumer
63     // memref dependence are considered for fusion. Currently, assumptions from
64     // the producer-consumer fusion implementation in AffineLoopFusion pass are
65     // made. See pass for specific details.
66     ProducerConsumer,
67     // Sibling fusion: Only sibling loops with no producer-consumer memref
68     // dependences are considered for fusion. Memref reuse is taken into account
69     // for profitability. Currently, assumptions from the sibling fusion
70     // implementation in AffineLoopFusion pass are made. See pass for specific
71     // details.
72     Sibling
73   };
74
75   /// Construct a generic or producer-consumer fusion strategy.
76   FusionStrategy(StrategyEnum strategy) : strategy(strategy) {
77     assert(strategy != Sibling &&
78            "Sibling fusion strategy requires a specific memref");
79   }
80
81   /// Construct a sibling fusion strategy targeting 'memref'. This construct
82   /// should only be used for sibling fusion.
83   FusionStrategy(Value memref) : strategy(Sibling), memref(memref) {}
84
85   /// Returns the fusion strategy.
86   StrategyEnum getStrategy() const { return strategy; };
87
88   /// Returns the memref attached to this sibling fusion strategy.
89   Value getSiblingFusionMemRef() const {
90     assert(strategy == Sibling && "Memref is only valid for sibling fusion");
91     return memref;
92   }
93
94 private:
95   /// Fusion strategy.
96   StrategyEnum strategy;
97
98   /// Target memref for this fusion transformation. Only used for sibling
99   /// fusion.
100   Value memref;
101 };
102
103 /// Checks the feasibility of fusing the loop nest rooted at 'srcForOp' into the
104 /// loop nest rooted at 'dstForOp' at 'dstLoopDepth'. Returns FusionResult
105 /// 'Success' if fusion of the src/dst loop nests is feasible (i.e. they are
106 /// in the same block and dependences would not be violated). Otherwise
107 /// returns a FusionResult explaining why fusion is not feasible.
108 /// NOTE: This function is not feature complete and should only be used in
109 /// testing.
110 /// TODO: Update comments when this function is fully implemented.
111 FusionResult
112 canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, unsigned dstLoopDepth,
113              ComputationSliceState *srcSlice,
114              FusionStrategy fusionStrategy = FusionStrategy::Generic);
115
116 /// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
117 /// and source slice loop bounds specified in 'srcSlice'.
118 void fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
119                const ComputationSliceState &srcSlice);
120
121 /// LoopNestStats aggregates various per-loop statistics (eg. loop trip count
122 /// and operation count) for a loop nest up until (and including) the innermost
123 /// loop body.
124 struct LoopNestStats {
125   /// Map from AffineForOp to immediate child AffineForOps in its loop body.
126   DenseMap<Operation *, SmallVector<AffineForOp, 2>> loopMap;
127   /// Map from AffineForOp to count of operations in its loop body.
128   DenseMap<Operation *, uint64_t> opCountMap;
129   /// Map from AffineForOp to its constant trip count.
130   DenseMap<Operation *, uint64_t> tripCountMap;
131 };
132
133 /// Collect loop nest statistics (eg. loop trip count and operation count)
134 /// in 'stats' for loop nest rooted at 'forOp'. Returns true on success,
135 /// returns false otherwise.
136 // TODO: Consider moving this to LoopUtils.
137 bool getLoopNestStats(AffineForOp forOp, LoopNestStats *stats);
138
139 /// Computes the total cost of the loop nest rooted at 'forOp' using 'stats'.
140 /// Currently, the total cost is computed by counting the total operation
141 /// instance count (i.e. total number of operations in the loop body * loop
142 /// trip count) for the entire loop nest.
143 // TODO: Improve this cost model.
144 int64_t getComputeCost(AffineForOp forOp, LoopNestStats &stats);
145
146 /// Computes and returns in 'computeCost', the total compute cost of fusing the
147 /// 'slice' of the loop nest rooted at 'srcForOp' into 'dstForOp'. Currently,
148 /// the total cost is computed by counting the total operation instance count
149 /// (i.e. total number of operations in the loop body * loop trip count) for
150 /// the entire loop nest.
151 /// Returns true on success, failure otherwise (e.g. non-constant trip counts).
152 // TODO: Improve this cost model.
153 bool getFusionComputeCost(AffineForOp srcForOp, LoopNestStats &srcStats,
154                           AffineForOp dstForOp, LoopNestStats &dstStats,
155                           const ComputationSliceState &slice,
156                           int64_t *computeCost);
157
158 /// Returns in 'producerConsumerMemrefs' the memrefs involved in a
159 /// producer-consumer dependence between write ops in 'srcOps' and read ops in
160 /// 'dstOps'.
161 void gatherProducerConsumerMemrefs(ArrayRef<Operation *> srcOps,
162                                    ArrayRef<Operation *> dstOps,
163                                    DenseSet<Value> &producerConsumerMemrefs);
164 } // end namespace mlir
165
166 #endif // MLIR_TRANSFORMS_LOOP_FUSION_UTILS_H