1 /* Copyright 2022 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_GML_ST_TRANSFORMS_PASSES_H
17 #define MLIR_HLO_DIALECT_GML_ST_TRANSFORMS_PASSES_H
18 
19 #include <memory>
20 #include <string>
21 
22 #include "mlir/Dialect/Func/IR/FuncOps.h"
23 #include "mlir/Pass/Pass.h"
24 
25 namespace mlir {
26 namespace gml_st {
27 
28 /// Pass to fuse producers into `gml_st.materialize` ops.
29 std::unique_ptr<OperationPass<func::FuncOp>> createFusionPass();
30 
31 /// Pass to tile operations.
32 std::unique_ptr<OperationPass<func::FuncOp>> createTilingPass();
33 std::unique_ptr<OperationPass<func::FuncOp>> createTilingPass(
34     const SmallVector<SmallVector<int64_t>>& tileSizes);
35 std::unique_ptr<OperationPass<func::FuncOp>> createTilingPass(
36     const std::string& tileSizes);
37 
38 /// Pass to compose set operations.
39 std::unique_ptr<OperationPass<func::FuncOp>> createComposeSetOpsPass();
40 
41 /// Pass to collapse (or uncollapse) materialize operations.
42 std::unique_ptr<OperationPass<func::FuncOp>> createCollapseMaterializeOpsPass(
43     bool reverse = false);
44 
45 /// Create a pass to convert `gml_st.loop` to `scf.for` and `scf.parallel`
46 /// loops and memref.load/memref.store accesses.
47 std::unique_ptr<OperationPass<func::FuncOp>> createGmlStToScfPass();
48 
49 // Pass to bufferize `linalg.tiled_loop` including the operations contained in
50 // its body.
51 std::unique_ptr<OperationPass<func::FuncOp>> CreateTiledLoopBufferizePass();
52 
53 /// Pass to vectorize linalg.generic ops tiled to gml_st.parallel and gml_st.for
54 /// loops.
55 std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeGmlStLoopsPass();
56 
57 #define GEN_PASS_REGISTRATION
58 #include "mlir-hlo/Dialect/gml_st/transforms/passes.h.inc"
59 
60 }  // namespace gml_st
61 }  // namespace mlir
62 
63 #endif  // MLIR_HLO_DIALECT_GML_ST_TRANSFORMS_PASSES_H
64