1 //===- TensorTilingOpInterfaceImpl.h - ------------------------------------===//
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 // This file implements Tiling interface for TensorOps with ExternalModel.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
14 #define MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
15 
16 #include "mlir/IR/Dialect.h"
17 
18 namespace mlir {
19 namespace tensor {
20 
21 class PadOp;
22 
23 /// Bubbles up a slice of this pad by taking the slice first and then performing
24 /// the padding. `offsets` and `strides` specifies each dimension's start offset
25 /// and size for the slice. The slice has unit strides along all dimensions.
26 ///
27 /// Specifically, this function converts:
28 /// ```
29 /// %0 = tensor.pad %source low[...] high[...] { linalg.yield %cst }
30 /// %1 = <extract-slice> %0 offsets=[...], sizes[...]
31 /// ```
32 /// into
33 /// ```
34 /// %0 = tensor.extract_slice %source ...
35 /// %0 = tensor.pad %0 low[...] high[...] { linalg.yield %cst }
36 /// ```
37 ///
38 /// If `generateZeroSliceGuard` is true, the generated IR will contain logic
39 /// to guard against the case that we might take a zero-sized slice from the
40 /// original source. For such cases, we `tensor.generate` to generate the
41 /// full tensor.
42 Operation *bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
43                             ArrayRef<OpFoldResult> offsets,
44                             ArrayRef<OpFoldResult> sizes,
45                             bool generateZeroSliceGuard = true);
46 
47 /// Registers external models for Tiling interface for tensor ops.
48 /// Currently, it registers:
49 ///
50 /// * TilingInterface for `tensor.pad`.
51 ///
52 /// Unfortunately, a "normal" internal registration is not possible at the
53 /// moment, because of the dependency of the interface implementation for these
54 /// ops on `affine.apply` and Affine dialect already depends on TensorOps. In
55 /// order to break the cyclic dependency (TensorOps->AffineOps->TensorOps) the
56 /// implementation is moved to a separate library.
57 void registerTilingOpInterfaceExternalModels(mlir::DialectRegistry &registry);
58 
59 } // namespace tensor
60 } // namespace mlir
61 
62 #endif // MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
63