1 //===- Utils.cpp - Utilities to support the Linalg 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 Linalg dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
14 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
15 
16 using namespace mlir;
17 
18 /// Matches a ConstantIndexOp.
19 /// TODO: This should probably just be a general matcher that uses matchConstant
20 /// and checks the operation for an index type.
21 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
22   return detail::op_matcher<arith::ConstantIndexOp>();
23 }
24 
25 /// Detects the `values` produced by a ConstantIndexOp and places the new
26 /// constant in place of the corresponding sentinel value.
27 void mlir::canonicalizeSubViewPart(
28     SmallVectorImpl<OpFoldResult> &values,
29     llvm::function_ref<bool(int64_t)> isDynamic) {
30   for (OpFoldResult &ofr : values) {
31     if (ofr.is<Attribute>())
32       continue;
33     // Newly static, move from Value to constant.
34     if (auto cstOp =
35             ofr.dyn_cast<Value>().getDefiningOp<arith::ConstantIndexOp>())
36       ofr = OpBuilder(cstOp).getIndexAttr(cstOp.value());
37   }
38 }
39 
40 void mlir::getPositionsOfShapeOne(
41     unsigned rank, ArrayRef<int64_t> shape,
42     llvm::SmallDenseSet<unsigned> &dimsToProject) {
43   dimsToProject.reserve(rank);
44   for (unsigned pos = 0, e = shape.size(); pos < e && rank > 0; ++pos) {
45     if (shape[pos] == 1) {
46       dimsToProject.insert(pos);
47       --rank;
48     }
49   }
50 }
51 
52 Value mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
53                                             OpFoldResult ofr) {
54   if (auto value = ofr.dyn_cast<Value>())
55     return value;
56   auto attr = ofr.dyn_cast<Attribute>().dyn_cast<IntegerAttr>();
57   assert(attr && "expect the op fold result casts to an integer attribute");
58   return b.create<arith::ConstantIndexOp>(loc, attr.getValue().getSExtValue());
59 }
60 
61 SmallVector<Value>
62 mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
63                                       ArrayRef<OpFoldResult> valueOrAttrVec) {
64   return llvm::to_vector<4>(
65       llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
66         return getValueOrCreateConstantIndexOp(b, loc, value);
67       }));
68 }
69 
70 Value ArithBuilder::_and(Value lhs, Value rhs) {
71   return b.create<arith::AndIOp>(loc, lhs, rhs);
72 }
73 Value ArithBuilder::add(Value lhs, Value rhs) {
74   if (lhs.getType().isa<IntegerType>())
75     return b.create<arith::AddIOp>(loc, lhs, rhs);
76   return b.create<arith::AddFOp>(loc, lhs, rhs);
77 }
78 Value ArithBuilder::mul(Value lhs, Value rhs) {
79   if (lhs.getType().isa<IntegerType>())
80     return b.create<arith::MulIOp>(loc, lhs, rhs);
81   return b.create<arith::MulFOp>(loc, lhs, rhs);
82 }
83 Value ArithBuilder::sgt(Value lhs, Value rhs) {
84   if (lhs.getType().isa<IndexType, IntegerType>())
85     return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, lhs, rhs);
86   return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGT, lhs, rhs);
87 }
88 Value ArithBuilder::slt(Value lhs, Value rhs) {
89   if (lhs.getType().isa<IndexType, IntegerType>())
90     return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, lhs, rhs);
91   return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, lhs, rhs);
92 }
93 Value ArithBuilder::select(Value cmp, Value lhs, Value rhs) {
94   return b.create<arith::SelectOp>(loc, cmp, lhs, rhs);
95 }
96