MLIR  20.0.0git
TensorTilingInterfaceImpl.h
Go to the documentation of this file.
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 
20 struct TilingResult;
21 
22 namespace tensor {
23 
24 class PadOp;
25 
26 /// Bubbles up a slice of this pad by taking the slice first and then performing
27 /// the padding. `offsets` and `strides` specifies each dimension's start offset
28 /// and size for the slice. The slice has unit strides along all dimensions.
29 ///
30 /// Specifically, this function converts:
31 /// ```
32 /// %0 = tensor.pad %source low[...] high[...] { linalg.yield %cst }
33 /// %1 = <extract-slice> %0 offsets=[...], sizes[...]
34 /// ```
35 /// into
36 /// ```
37 /// %0 = tensor.extract_slice %source ...
38 /// %0 = tensor.pad %0 low[...] high[...] { linalg.yield %cst }
39 /// ```
40 ///
41 /// If `generateZeroSliceGuard` is true, the generated IR will contain logic
42 /// to guard against the case that we might take a zero-sized slice from the
43 /// original source. For such cases, we `tensor.generate` to generate the
44 /// full tensor.
45 FailureOr<TilingResult> bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
46  ArrayRef<OpFoldResult> offsets,
47  ArrayRef<OpFoldResult> sizes,
48  bool generateZeroSliceGuard = true);
49 
50 /// Registers external models for Tiling interface for tensor ops.
51 /// Currently, it registers:
52 ///
53 /// * TilingInterface for `tensor.pad`, `tensor.pack`, and `tensor.unpack`.
54 ///
55 /// Unfortunately, a "normal" internal registration is not possible at the
56 /// moment, because of the dependency of the interface implementation for these
57 /// ops on `affine.apply` and Affine dialect already depends on TensorOps. In
58 /// order to break the cyclic dependency (TensorOps->AffineOps->TensorOps) the
59 /// implementation is moved to a separate library.
61 
62 /// Similar to the above registeration, but it is only for `tensor.pack` and
63 /// `tensor.unpack` ops.
65  DialectRegistry &registry);
66 
67 } // namespace tensor
68 } // namespace mlir
69 
70 #endif // MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
FailureOr< TilingResult > bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, bool generateZeroSliceGuard=true)
Bubbles up a slice of this pad by taking the slice first and then performing the padding.
void registerTilingInterfaceExternalModels(mlir::DialectRegistry &registry)
Registers external models for Tiling interface for tensor ops.
void registerTilingInterfaceExternalModelsForPackUnPackOps(DialectRegistry &registry)
Similar to the above registeration, but it is only for tensor.pack and tensor.unpack ops.
Include the generated interface declarations.