Better implementation of MergeComputeNodesPass

src/PIM/Common/IR/ConstantUtils.cpp

@@ -10,6 +10,32 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+ConstantPool::ConstantPool(Operation *constantAnchor, OpBuilder &builder)
+  : anchor(constantAnchor), block(getConstantInsertionBlock(constantAnchor)),
+    builder(builder) {
+  for (Operation &op : *block)
+    if (auto constant = dyn_cast<arith::ConstantOp>(&op))
+      cache.try_emplace(
+        std::make_pair(constant.getType(), constant.getValue()),
+        constant.getResult());
+}
+
+Value ConstantPool::getIndex(int64_t value) {
+  return get(builder.getIndexType(), builder.getIndexAttr(value));
+}
+
+Value ConstantPool::get(Type type, Attribute value) {
+  auto key = std::make_pair(type, value);
+  if (Value existing = cache.lookup(key))
+    return existing;
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(block);
+  Value constant = arith::ConstantOp::create(
+    builder, anchor->getLoc(), type, cast<TypedAttr>(value)).getResult();
+  cache.try_emplace(key, constant);
+  return constant;
+}
+
 static std::optional<int64_t> getIndexConstantValue(arith::ConstantOp constantOp) {
   if (!constantOp.getType().isIndex())
     return std::nullopt;