1//===- DataLayoutAnalysis.cpp ---------------------------------------------===//
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/Analysis/DataLayoutAnalysis.h"
10#include "mlir/IR/BuiltinOps.h"
11#include "mlir/IR/Operation.h"
12#include "mlir/Interfaces/DataLayoutInterfaces.h"
13#include "mlir/Support/LLVM.h"
14#include <memory>
15
16using namespace mlir;
17
18DataLayoutAnalysis::DataLayoutAnalysis(Operation *root)
19 : defaultLayout(std::make_unique<DataLayout>(args: DataLayoutOpInterface())) {
20 // Construct a DataLayout if possible from the op.
21 auto computeLayout = [this](Operation *op) {
22 if (auto iface = dyn_cast<DataLayoutOpInterface>(op))
23 layouts[op] = std::make_unique<DataLayout>(iface);
24 if (auto module = dyn_cast<ModuleOp>(op))
25 layouts[op] = std::make_unique<DataLayout>(module);
26 };
27
28 // Compute layouts for both ancestors and descendants.
29 root->walk(callback&: computeLayout);
30 for (Operation *ancestor = root->getParentOp(); ancestor != nullptr;
31 ancestor = ancestor->getParentOp()) {
32 computeLayout(ancestor);
33 }
34}
35
36const DataLayout &DataLayoutAnalysis::getAbove(Operation *operation) const {
37 for (Operation *ancestor = operation->getParentOp(); ancestor != nullptr;
38 ancestor = ancestor->getParentOp()) {
39 auto it = layouts.find(Val: ancestor);
40 if (it != layouts.end())
41 return *it->getSecond();
42 }
43
44 // Fallback to the default layout.
45 return *defaultLayout;
46}
47
48const DataLayout &DataLayoutAnalysis::getAtOrAbove(Operation *operation) const {
49 auto it = layouts.find(Val: operation);
50 if (it != layouts.end())
51 return *it->getSecond();
52 return getAbove(operation);
53}
54

source code of mlir/lib/Analysis/DataLayoutAnalysis.cpp