1 //===----------------------------------------------------------------------===//
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 #include "mlir/Dialect/MemRef/IR/MemRef.h"
10 #include "mlir/Dialect/Tensor/IR/Tensor.h"
11 #include "mlir/Transforms/InliningUtils.h"
12 
13 using namespace mlir;
14 using namespace mlir::memref;
15 
16 #include "mlir/Dialect/MemRef/IR/MemRefOpsDialect.cpp.inc"
17 
18 //===----------------------------------------------------------------------===//
19 // MemRefDialect Dialect Interfaces
20 //===----------------------------------------------------------------------===//
21 
22 namespace {
23 struct MemRefInlinerInterface : public DialectInlinerInterface {
24   using DialectInlinerInterface::DialectInlinerInterface;
25   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
26                        BlockAndValueMapping &valueMapping) const final {
27     return true;
28   }
29   bool isLegalToInline(Operation *, Region *, bool wouldBeCloned,
30                        BlockAndValueMapping &) const final {
31     return true;
32   }
33 };
34 } // end anonymous namespace
35 
36 SmallVector<Value, 4> mlir::getDynOperands(Location loc, Value val,
37                                            OpBuilder &b) {
38   SmallVector<Value, 4> dynOperands;
39   auto shapedType = val.getType().cast<ShapedType>();
40   for (auto dim : llvm::enumerate(shapedType.getShape())) {
41     if (dim.value() == ShapedType::kDynamicSize)
42       dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
43   }
44   return dynOperands;
45 }
46 
47 // Helper function that creates a memref::DimOp or tensor::DimOp depending on
48 // the type of `source`.
49 // TODO: Move helper function out of MemRef dialect.
50 Value mlir::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
51                               int64_t dim) {
52   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
53     return b.createOrFold<memref::DimOp>(loc, source, dim);
54   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
55     return b.createOrFold<tensor::DimOp>(loc, source, dim);
56   llvm_unreachable("Expected MemRefType or TensorType");
57 }
58 
59 void mlir::memref::MemRefDialect::initialize() {
60   addOperations<DmaStartOp, DmaWaitOp,
61 #define GET_OP_LIST
62 #include "mlir/Dialect/MemRef/IR/MemRefOps.cpp.inc"
63                 >();
64   addInterfaces<MemRefInlinerInterface>();
65 }
66