1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef MLIR_HLO_DIALECT_LHLO_TRANSFORMS_LHLO_ELEMENTAL_UTILS_H
17 #define MLIR_HLO_DIALECT_LHLO_TRANSFORMS_LHLO_ELEMENTAL_UTILS_H
18 
19 #include "mlir/IR/Builders.h"
20 
21 namespace mlir {
22 namespace func {
23 class FuncOp;
24 }  // namespace func
25 class Value;
26 class Location;
27 class Operation;
28 class ValueRange;
29 class Region;
30 enum class AtomicRMWKind : uint64_t;
31 
32 namespace scf {
33 class ForOp;
34 class ParallelOp;
35 }  // namespace scf
36 
37 namespace memref {
38 class LoadOp;
39 }  // namespace memref
40 
41 namespace lmhlo {
42 
43 Value createLoadOrUseCachedValue(Location loc, OpBuilder* b, Value memref,
44                                  ValueRange indices,
45                                  OpBuilder::InsertPoint insertPoint);
46 
47 DenseSet<Operation*> noLoaderUser(SmallVectorImpl<Operation*>& ops);
48 void cleanUnusedLhloOps(Block* parent);
49 
50 template <typename LHLO_OpTy>
51 Value elementalLower(OpBuilder* b, Location loc, LHLO_OpTy op,
52                      ValueRange outputIndex, bool checkCache = false);
53 
54 scf::ForOp createLoopAndSetInsPt(OpBuilder& b, Location loc, Value& var,
55                                  Value lb, Value ub, Value step,
56                                  ArrayRef<Value> initValues = {});
57 
58 scf::ParallelOp createParallelAndSetInsPt(OpBuilder& b, Location loc,
59                                           SmallVectorImpl<Value>& vars,
60                                           ArrayRef<Value> lbs,
61                                           ArrayRef<Value> ubs,
62                                           ArrayRef<Value> steps,
63                                           ArrayRef<Value> initValues);
64 
65 void createOffsetStore(OpBuilder& b, Location loc, Value res, Value memref,
66                        Value offset);
67 
68 memref::LoadOp createOffsetLoad(OpBuilder& b, Location loc, Value memref,
69                                 Value offset);
70 
71 }  // namespace lmhlo
72 }  // namespace mlir
73 
74 #endif  // MLIR_HLO_DIALECT_LHLO_TRANSFORMS_LHLO_ELEMENTAL_UTILS_H
75