1 //===- Utils.cpp - Utilities to support the Tensor dialect ----------------===//
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 utilities for the Tensor dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Tensor/Utils/Utils.h"
14 
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
17 
18 using namespace mlir;
19 using namespace mlir::tensor;
20 
21 PadOp mlir::tensor::createPadScalarOp(Type type, Value source, Value pad,
22                                       ArrayRef<OpFoldResult> low,
23                                       ArrayRef<OpFoldResult> high, bool nofold,
24                                       Location loc, OpBuilder &builder) {
25   auto padTensorOp =
26       builder.create<PadOp>(loc, type, source, low, high, nofold);
27   int rank = padTensorOp.getResultType().getRank();
28   SmallVector<Type> blockArgTypes(rank, builder.getIndexType());
29   SmallVector<Location> blockArgLocs(rank, loc);
30   auto &region = padTensorOp.region();
31   // `builder.createBlock` changes the insertion point within the block. Create
32   // a guard to reset the insertion point of the builder after it is destroyed.
33   OpBuilder::InsertionGuard guard(builder);
34   builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
35   builder.create<YieldOp>(loc, pad);
36   return padTensorOp;
37 }
38 
39 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
40                                     Value pad, bool nofold, Location loc,
41                                     OpBuilder &b) {
42   auto zero = b.createOrFold<arith::ConstantIndexOp>(loc, 0);
43   SmallVector<OpFoldResult> low(type.getRank(), zero);
44   SmallVector<OpFoldResult> high(type.getRank(), zero);
45   for (const auto &en : enumerate(type.getShape())) {
46     // Pad only the static dimensions of the result tensor type.
47     if (ShapedType::isDynamic(en.value()))
48       continue;
49     // Compute the padding width.
50     AffineExpr d0;
51     bindDims(b.getContext(), d0);
52     auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
53     high[en.index()] =
54         makeComposedAffineApply(b, loc, en.value() - d0, {dimOp}).getResult();
55   }
56   return createPadScalarOp(type, source, pad, low, high, nofold, loc, b);
57 }
58 
59 SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
60                                                         Location loc,
61                                                         Value rankedTensor) {
62   auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
63   SmallVector<Value> dynamicDims;
64   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
65     if (en.value() == ShapedType::kDynamicSize)
66       dynamicDims.push_back(
67           b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
68   }
69   return dynamicDims;
70 }
71