1 /* Copyright 2019 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 TENSORFLOW_COMPILER_MLIR_XLA_EXPERIMENTAL_CONV_EMITTER_CONV_EMITTER_TRANSFORMS_H_ 17 #define TENSORFLOW_COMPILER_MLIR_XLA_EXPERIMENTAL_CONV_EMITTER_CONV_EMITTER_TRANSFORMS_H_ 18 19 #include "absl/types/span.h" 20 #include "mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project 21 #include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project 22 #include "mlir/IR/Operation.h" // from @llvm-project 23 #include "tensorflow/core/platform/types.h" 24 25 namespace xla { 26 namespace experimental { 27 28 struct BoundAffineMap { 29 mlir::AffineMap affine_map; 30 std::vector<mlir::Value> operands; 31 }; 32 33 BoundAffineMap GetBoundAffineMapFrom(mlir::Operation* op); 34 mlir::Operation* CloneWithNewAffineMap(mlir::Operation* op, 35 BoundAffineMap new_affine, 36 mlir::OpBuilder builder); 37 38 bool IsSimpleLoop(mlir::AffineForOp loop); 39 std::vector<mlir::AffineForOp> CreateNestedSimpleLoops( 40 absl::Span<const int64_t> upper_bounds, mlir::OpBuilder builder); 41 void SetBoundForSimpleLoop(mlir::AffineForOp loop, mlir::AffineExpr new_bound, 42 mlir::OpBuilder builder); 43 44 // Tile a loop with trip count N by `size`. For now, N has to be a multiple of 45 // size, but later this constraint will be removed. 46 // 47 // The major loop (with trip count N / size) stays as-is, while the minor loop 48 // (with trip count `size`) will take over the body of `target`, and be placed 49 // as the new body of `target`. 50 // 51 // `target` has to be within the same "perfectly nested loop group" as `loop`. 52 // See the documentation for mlir::getPerfectlyNestedLoops. 53 // 54 // Example: 55 // Before tiling `loop` with tile size X: 56 // for (loop in N) 57 // for (unrelated_loop in ...) 58 // for (target in ...) 59 // // pass loop into affine maps 60 // After: 61 // for (loop in N / X) 62 // for (unrelated_loop in ...) 63 // for (target in ...) 64 // for (tiled_loop in X) 65 // // rewrite all affine exprs from loop to `loop * X + tiled_loop`. 66 // 67 // Design note: 68 // TileLoop is different from mlir::tile. At the moment, mlir::tile is not well 69 // documented about the exact tiling semantics, but the observed behavior is: 70 // for (i from 0 to N) 71 // for (unrelated_loop in ...) 72 // for (target in ...) 73 // // pass i into affine maps 74 // => 75 // for (i from 0 to N, step = X) 76 // for (unrelated_loop in ...) 77 // for (target in ...) 78 // for (j from i to min(i + X, N), step = 1) 79 // // pass j into affine maps 80 // 81 // There are two differences between mlir::tile and TileLoop: 82 // * TileLoop always puts the tiling logic "stepping" logic into AffineExprs. 83 // With that all index calculation is done in AffineExprs and easier to 84 // analyze in a single place. 85 // * TileLoop doesn't plan to use max() and min() to resolve the issue when 86 // N % X != 0. max() and min() are not representable in AffineExprs. 87 // TODO(timshen): support the case where N % X != 0. 88 // 89 // TODO(timshen): consider the possibility to reuse mlir::tile's logic to 90 // achieve the same goal. 91 mlir::AffineForOp TileLoop(mlir::AffineForOp loop, int64_t size, 92 mlir::AffineForOp target); 93 94 // Sinks a segment of perfectly nested loops to the bottom. It implements this 95 // by rotating the loop nest by rotate_amount. 96 void SinkPerfectlyNestedLoops(llvm::MutableArrayRef<mlir::AffineForOp> loops, 97 int rotate_amount); 98 99 } // namespace experimental 100 } // namespace xla 101 102 #endif // TENSORFLOW_COMPILER_MLIR_XLA_EXPERIMENTAL_CONV_EMITTER_CONV_EMITTER_TRANSFORMS_H_ 103