[Mlir-commits] [mlir] [MLIR][LLVM] Fix memory explosion when converting global variable bodies in ModuleTranslation (PR #82708)

[Mehdi Amini]

================
@@ -1042,17 +1046,77 @@ LogicalResult ModuleTranslation::convertGlobals() {
   for (auto op : getModuleBody(mlirModule).getOps<LLVM::GlobalOp>()) {
     if (Block *initializer = op.getInitializerBlock()) {
       llvm::IRBuilder<> builder(llvmModule->getContext());
+
+      int numConstantsHit = 0;
+      int numConstantsErased = 0;
+      DenseMap<llvm::ConstantAggregate *, int> constantAggregateUseMap;
+
       for (auto &op : initializer->without_terminator()) {
         if (failed(convertOperation(op, builder)) ||
             !isa<llvm::Constant>(lookupValue(op.getResult(0))))
           return emitError(op.getLoc(), "unemittable constant value");
+
+        // When emitting an LLVM constant, a new constant is created and the old
+        // constant may become dangling and take space. We should remove the
+        // dangling constants to avoid memory explosion especially for constant
+        // arrays whose number of elements is large.
+        // Because multiple operations may refer to the same constant, we need
+        // to count the number of uses of each constant array and remove it only
+        // when the count becomes zero.
+        if (op.getNumResults() == 1) {
+          Value result = op.getResult(0);
+          auto cst = dyn_cast<llvm::ConstantAggregate>(lookupValue(result));
+          if (!cst)
+            continue;
+          numConstantsHit++;
+          auto iter = constantAggregateUseMap.find(cst);
+          int numUsers = std::distance(result.use_begin(), result.use_end());
+          if (iter == constantAggregateUseMap.end())
+            constantAggregateUseMap.try_emplace(cst, numUsers);
----------------
joker-eph wrote:

Instead of `find` followed by an insertion, you can just try to insert and update the existing entry if the insertion "failed": that means a single lookup.

https://github.com/llvm/llvm-project/pull/82708