NiccoloN
53 lines
1.6 KiB
C++
53 lines
1.6 KiB
C++
#pragma once
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Transforms/FoldUtils.h"
|
|
|
|
#include "llvm/ADT/DenseMap.h"
|
|
|
|
#include <optional>
|
|
|
|
namespace onnx_mlir {
|
|
|
|
class ConstantPool {
|
|
public:
|
|
ConstantPool(mlir::Operation *constantAnchor, mlir::OpBuilder &builder);
|
|
|
|
mlir::Value getIndex(int64_t value);
|
|
mlir::Value get(mlir::Type type, mlir::Attribute value);
|
|
|
|
private:
|
|
mlir::Operation *anchor;
|
|
mlir::Block *block;
|
|
mlir::OpBuilder &builder;
|
|
llvm::DenseMap<std::pair<mlir::Type, mlir::Attribute>, mlir::Value> cache;
|
|
};
|
|
|
|
mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);
|
|
|
|
mlir::Value
|
|
getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
|
|
|
|
mlir::Value
|
|
getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
|
|
|
|
mlir::Value
|
|
createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value);
|
|
|
|
mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);
|
|
|
|
mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);
|
|
|
|
mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value);
|
|
|
|
void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);
|
|
|
|
std::optional<int64_t> matchConstantIndexValue(mlir::Value value);
|
|
|
|
std::optional<int64_t> matchConstantIndexValue(mlir::OpFoldResult value);
|
|
|
|
} // namespace onnx_mlir
|