1 //===- HoistPadding.h - Hoisting for tensor::PadOp -*- 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 #ifndef MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H 10 #define MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H 11 12 #include "mlir/Support/LogicalResult.h" 13 14 namespace mlir { 15 class Value; 16 17 namespace tensor { 18 class PadOp; 19 } // namespace tensor 20 21 namespace linalg { 22 class GenericOp; 23 24 /// Mechanically hoist padding operations on tensors by `numLoops` into a new, 25 /// generally larger tensor. This achieves packing of multiple padding ops into 26 /// a larger tensor. On success, `opToHoist` is replaced by the cloned version 27 /// in the packing loop so the caller can continue reasoning about the padding 28 /// operation. If `transposeVector` is non-empty, hoist padding introduces a 29 /// GenericOp to transpose the padded tensor before inserting it into the packed 30 /// tensor. A `transposeVector` can change the storage order of the padded 31 /// tensor but does not change the order of the pack or compute loops. 32 /// 33 /// 34 /// Example in pseudo-mlir: 35 /// ======================= 36 /// 37 /// If hoistPaddingOnTensors is called with `nLoops` = 2 on the following IR. 38 /// ``` 39 /// scf.for (%i, %j, %k) 40 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32> 41 /// %0 = tensor.pad %st0 low[0, 0] high[...] { 42 /// ^bb0( ... ): 43 /// linalg.yield %pad 44 /// } : tensor<?x?xf32> to tensor<4x8xf32> 45 /// compute(%0) 46 /// ``` 47 /// 48 /// IR resembling the following is produced: 49 /// 50 /// ``` 51 /// scf.for (%i) { 52 /// %packed_init = linalg.init_tensor range(%j) : tensor<?x4x8xf32> 53 /// %packed = scf.for (%k) iter_args(%p : %packed_init) { 54 /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32> 55 /// %0 = tensor.pad %st0 low[0, 0] high[...] { 56 /// ^bb0( ... ): 57 /// linalg.yield %pad 58 /// } : tensor<?x?xf32> to tensor<4x8xf32> 59 /// %1 = tensor.insert_slice %0 ... 60 /// : tensor<4x8xf32> to tensor<?x4x8xf32> 61 /// scf.yield %1: tensor<?x4x8xf32> 62 /// } -> tensor<?x4x8xf32> 63 /// scf.for (%j, %k) { 64 /// %st0 = tensor.extract_slice %packed [%k, 0, 0][1, 4, 8][1, 1, 1] : 65 /// tensor<?x4x8xf32> to tensor<4x8xf32> 66 /// compute(%st0) 67 /// } 68 /// } 69 /// ``` 70 FailureOr<Value> hoistPaddingOnTensors( 71 tensor::PadOp opToHoist, int numLoops, ArrayRef<int64_t> transposeVector, 72 tensor::PadOp &hoistedOp, SmallVectorImpl<GenericOp> &transposeOps); 73 74 } // namespace linalg 75 } // namespace mlir 76 77 #endif // MLIR_DIALECT_LINALG_TRANSFORMS_HOISTPADDING_H 78