1 /* Copyright 2020 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_XLA_SERVICE_MLIR_GPU_PASSES_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_MLIR_GPU_PASSES_H_ 18 19 #include <memory> 20 21 #include "mlir/Pass/Pass.h" // from @llvm-project 22 23 namespace xla { 24 namespace mlir_gpu { 25 26 // TODO(herhut, pifon): Move these passes to MLIR Core. 27 28 /// Replaces a FusionOp by the operations contained in its region. 29 std::unique_ptr<mlir::FunctionPass> createFusionOpRemoverPass(); 30 31 /// Replaces a load that immediately follows a store to the same address with 32 /// the stored value. This needs generalization. 33 std::unique_ptr<mlir::FunctionPass> createStoreForwardingPass(); 34 35 /// Removes temporary buffers that are only written to but never read from or 36 /// that are read but the read value is not used. Needs an analysis that proves 37 /// that loads and stores are side-effect free (in bounds, no aliasing, etc.). 38 std::unique_ptr<mlir::FunctionPass> createDeadTempBufferRemovalPass(); 39 40 /// Sorts the operands to the kernel for a deterministic order. First operands 41 /// that are defined by function arguments, followed by operands that are 42 /// returned from the function. This only works for simple functions without 43 /// control flow and can be used in cases where the kernel is extracted and used 44 /// independently of the host-side code. 45 std::unique_ptr<mlir::FunctionPass> createRewriteKernelSignaturePass(); 46 47 /// We need to direct fusion to the inner loops. This cannot be done with 48 /// a passmanager alone ATM, as nested pass managers require operations to 49 /// be closed from above. 50 std::unique_ptr<mlir::FunctionPass> createFuseInnerParallelLoopsPass(); 51 52 /// Greedily maps loops to GPU hardware dimensions. 53 std::unique_ptr<mlir::FunctionPass> createMapParallelLoopsPass(); 54 55 /// Collapses all loop dimension into the first one. 56 std::unique_ptr<mlir::FunctionPass> 57 createParallelLoopCollapsingToFirstDimPass(); 58 59 #define GEN_PASS_REGISTRATION 60 #include "tensorflow/compiler/xla/service/mlir_gpu/passes.h.inc" 61 62 } // namespace mlir_gpu 63 } // namespace xla 64 65 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_MLIR_GPU_PASSES_H_ 66