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
18namespace mlir {
19
20struct TilingResult;
21
22namespace tensor {
23
24class 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.
45FailureOr<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.
60void registerTilingInterfaceExternalModels(mlir::DialectRegistry &registry);
61
62/// Similar to the above registeration, but it is only for `tensor.pack` and
63/// `tensor.unpack` ops.
64void registerTilingInterfaceExternalModelsForPackUnPackOps(
65 DialectRegistry &registry);
66
67} // namespace tensor
68} // namespace mlir
69
70#endif // MLIR_DIALECT_TENSOR_IR_TENSORTILINGINTERFACEIMPL_H_
71

source code of mlir/include/mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h