61 lines
2.0 KiB
C++
61 lines
2.0 KiB
C++
//===- ShapeMappingAnalysis.h - Preserve shape Info ------------*- C++ -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_
|
|
#define MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_
|
|
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "llvm/ADT/DenseMap.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/SmallVector.h"
|
|
|
|
namespace mlir {
|
|
|
|
namespace shape {
|
|
|
|
/// ShapeMappingValue works as the value of ShapeMappingAnalysis table, where
|
|
/// `funcSymbol` is the symbol of mapping function, and `inputs` are the actual
|
|
/// parameters for the function.
|
|
struct ShapeMappingValue {
|
|
ShapeMappingValue() = default;
|
|
ShapeMappingValue(FlatSymbolRefAttr symbol, llvm::SmallVector<Value> &&inps)
|
|
: funcSymbol(symbol), inputs(inps) {}
|
|
|
|
FlatSymbolRefAttr funcSymbol;
|
|
llvm::SmallVector<Value> inputs;
|
|
};
|
|
|
|
/// ShapeMappingAnalysis is used together with OutlineShapeComputationPass to
|
|
/// preserve Value and corresponding shape function / arguments mapping
|
|
/// information
|
|
struct ShapeMappingAnalysis {
|
|
ShapeMappingAnalysis(Operation *op) : operation(op) { (void)operation; }
|
|
|
|
/// Dumps the shape mapping information to the given stream.
|
|
void print(raw_ostream &os) const {
|
|
os << "// ---- Shape Mapping Information -----\n";
|
|
for (const auto &it : shapeMapping) {
|
|
const ShapeMappingValue &mappingValue = it.second;
|
|
os << "// Shape for " << it.first << " :: " << mappingValue.funcSymbol;
|
|
llvm::interleaveComma(mappingValue.inputs, os << "(");
|
|
os << ")\n";
|
|
}
|
|
}
|
|
|
|
llvm::DenseMap<Value, ShapeMappingValue> shapeMapping;
|
|
|
|
private:
|
|
Operation *operation;
|
|
};
|
|
|
|
} // namespace shape
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_DIALECT_SHAPE_ANALYSIS_SHAPEMAPPINGANALYSIS_H_
|