NiccoloN
63 lines
2.5 KiB
C++
63 lines
2.5 KiB
C++
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/Dialect.h"
|
|
|
|
#include "ConstantUtils.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
Block* getHostConstantBlock(Operation* anchorOp) {
|
|
assert(anchorOp && "expected a valid anchor operation");
|
|
|
|
for (Operation* current = anchorOp; current; current = current->getParentOp())
|
|
if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(current))
|
|
return current->getBlock();
|
|
|
|
if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
|
|
return &funcOp.getBody().front();
|
|
if (auto moduleOp = anchorOp->getParentOfType<ModuleOp>())
|
|
return moduleOp.getBody();
|
|
return anchorOp->getBlock();
|
|
}
|
|
|
|
Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, OperationFolder& folder) {
|
|
assert(anchorOp && "expected a valid anchor operation");
|
|
Block* hostBlock = getHostConstantBlock(anchorOp);
|
|
for (Operation& op : *hostBlock) {
|
|
auto constantOp = dyn_cast<arith::ConstantOp>(&op);
|
|
if (!constantOp || constantOp.getType() != type || constantOp.getValue() != value)
|
|
continue;
|
|
return constantOp.getResult();
|
|
}
|
|
|
|
auto* arithDialect = anchorOp->getContext()->getOrLoadDialect<arith::ArithDialect>();
|
|
return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
|
|
}
|
|
|
|
Value getOrCreateHostConstantLike(arith::ConstantOp constantOp, OperationFolder& folder) {
|
|
return getOrCreateHostConstant(constantOp.getOperation(), constantOp.getValue(), constantOp.getType(), folder);
|
|
}
|
|
|
|
Value getOrCreateHostIndexConstant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
|
|
Builder builder(anchorOp->getContext());
|
|
return getOrCreateHostConstant(anchorOp, builder.getIndexAttr(value), builder.getIndexType(), folder);
|
|
}
|
|
|
|
Value getOrCreateHostI32Constant(Operation* anchorOp, int32_t value, OperationFolder& folder) {
|
|
Builder builder(anchorOp->getContext());
|
|
return getOrCreateHostConstant(anchorOp, builder.getI32IntegerAttr(value), builder.getI32Type(), folder);
|
|
}
|
|
|
|
Value getOrCreateHostI64Constant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
|
|
Builder builder(anchorOp->getContext());
|
|
return getOrCreateHostConstant(anchorOp, builder.getI64IntegerAttr(value), builder.getI64Type(), folder);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|