forked from OSchip/llvm-project
52 lines
1.8 KiB
C++
52 lines
1.8 KiB
C++
//===- DataLayoutAnalysis.cpp ---------------------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Analysis/DataLayoutAnalysis.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/Interfaces/DataLayoutInterfaces.h"
|
|
|
|
using namespace mlir;
|
|
|
|
DataLayoutAnalysis::DataLayoutAnalysis(Operation *root)
|
|
: defaultLayout(std::make_unique<DataLayout>(DataLayoutOpInterface())) {
|
|
// Construct a DataLayout if possible from the op.
|
|
auto computeLayout = [this](Operation *op) {
|
|
if (auto iface = dyn_cast<DataLayoutOpInterface>(op))
|
|
layouts[op] = std::make_unique<DataLayout>(iface);
|
|
if (auto module = dyn_cast<ModuleOp>(op))
|
|
layouts[op] = std::make_unique<DataLayout>(module);
|
|
};
|
|
|
|
// Compute layouts for both ancestors and descendants.
|
|
root->walk(computeLayout);
|
|
for (Operation *ancestor = root->getParentOp(); ancestor != nullptr;
|
|
ancestor = ancestor->getParentOp()) {
|
|
computeLayout(ancestor);
|
|
}
|
|
}
|
|
|
|
const DataLayout &DataLayoutAnalysis::getAbove(Operation *operation) const {
|
|
for (Operation *ancestor = operation->getParentOp(); ancestor != nullptr;
|
|
ancestor = ancestor->getParentOp()) {
|
|
auto it = layouts.find(ancestor);
|
|
if (it != layouts.end())
|
|
return *it->getSecond();
|
|
}
|
|
|
|
// Fallback to the default layout.
|
|
return *defaultLayout;
|
|
}
|
|
|
|
const DataLayout &DataLayoutAnalysis::getAtOrAbove(Operation *operation) const {
|
|
auto it = layouts.find(operation);
|
|
if (it != layouts.end())
|
|
return *it->getSecond();
|
|
return getAbove(operation);
|
|
}
|