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

[Xiang Li]

================
@@ -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);
+          else
+            iter->second += numUsers;
+        }
+        for (Value v : op.getOperands()) {
+          auto cst = dyn_cast<llvm::ConstantAggregate>(lookupValue(v));
+          if (!cst)
+            continue;
+          auto iter = constantAggregateUseMap.find(cst);
+          assert(iter != constantAggregateUseMap.end() && "constant not found");
+          iter->second--;
+          if (iter->second == 0) {
+            cst->removeDeadConstantUsers();
+            if (cst->user_empty()) {
+              cst->destroyConstant();
+              numConstantsErased++;
+            }
+            constantAggregateUseMap.erase(iter);
+          }
+        }
       }
----------------
python3kgae wrote:


The constant will fold, but the input constant is still in memory.
And too many temp constants are created which eats all the memory.

I tried to avoid creating these insertvalue which will be converted to temp constant.

Here's what I got with experiment

```

  llvm.mlir.global private @krnl_global_28(dense<["", "ab", "xyz", ...]> : tensor<75525x!krnl.string>) {addr_space = 0 : i32, alignment = 8 : i64} : !llvm.array<75525 x i64>

```
It crashes because !llvm.array<75525 x i64> is not correct type.

Any suggestion for how to express array of string for llvm dialect?




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