refactorone

src/PIM/Common/IR/ConstantUtils.cpp

@@ -0,0 +1,62 @@
+#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