[mlir] Fix memory explosion when converting global variable bodies in ModuleTranslation

There is memory explosion when converting the body or initializer region of a large global variable, e.g. a constant array.

For example, when translating a constant array of 100000 strings:
```
llvm.mlir.global internal constant @cats_strings() {addr_space = 0 : i32, alignment = 16 : i64} : !llvm.array<100000 x ptr<i8>> {
    %0 = llvm.mlir.undef : !llvm.array<100000 x ptr<i8>>
    %1 = llvm.mlir.addressof @om_1 : !llvm.ptr<array<1 x i8>>
    %2 = llvm.getelementptr %1[0, 0] : (!llvm.ptr<array<1 x i8>>) -> !llvm.ptr<i8>
    %3 = llvm.insertvalue %2, %0[0] : !llvm.array<100000 x ptr<i8>>
    %4 = llvm.mlir.addressof @om_2 : !llvm.ptr<array<1 x i8>>
    %5 = llvm.getelementptr %4[0, 0] : (!llvm.ptr<array<1 x i8>>) -> !llvm.ptr<i8>
    %6 = llvm.insertvalue %5, %3[1] : !llvm.array<100000 x ptr<i8>>
    %7 = llvm.mlir.addressof @om_3 : !llvm.ptr<array<1 x i8>>
    %8 = llvm.getelementptr %7[0, 0] : (!llvm.ptr<array<1 x i8>>) -> !llvm.ptr<i8>
    %9 = llvm.insertvalue %8, %6[2] : !llvm.array<100000 x ptr<i8>>
    %10 = llvm.mlir.addressof @om_4 : !llvm.ptr<array<1 x i8>>
    %11 = llvm.getelementptr %10[0, 0] : (!llvm.ptr<array<1 x i8>>) -> !llvm.ptr<i8>
    %12 = llvm.insertvalue %11, %9[3] : !llvm.array<100000 x ptr<i8>>

    ... (ignore the remaining part)
}
```

where `@om_1`, `@om_2`, ... are string global constants.

Each time an operation is converted to LLVM, a new constant is created.
When it comes to `llvm.insertvalue`, a new constant array of 100000 elements is created and the old constant array (input) is not destroyed.
This causes memory explosion. We observed that, on a system with 128 GB memory, the translation of 100000 elements got killed due to using up all the memory.
On a system with 64 GB, 65536 elements was enough to cause the translation killed.

This patch fixes the issue by checking generated constants and destroyed them if there is no use.
By the fix, the translation of 100000 elements only takes about 1.6 GB memory, and finishes without any error.

Reviewed By: ftynse, kiranchandramohan

Differential Revision: https://reviews.llvm.org/D148487

2 files changed, 25 insertions, 0 deletions

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index e4cd94026a59..70171e77eab9 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -741,6 +741,22 @@ LogicalResult ModuleTranslation::convertGlobals() {
         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.
+        // TODO: handle ops other than InsertValueOp with ConstantArray.
+        if (auto ivOp = dyn_cast<LLVM::InsertValueOp>(op)) {
+          Value container = ivOp.getContainer();
+          if (auto cst =
+                  dyn_cast<llvm::ConstantArray>(lookupValue(container))) {
+            // GlobalValue shouldn't be treated like other constants.
+            if (isa<llvm::GlobalValue>(cst))
+              continue;
+            if (cst->hasZeroLiveUses())
+              cst->destroyConstant();
+          }
+        }
       }
       ReturnOp ret = cast<ReturnOp>(initializer->getTerminator());
       llvm::Constant *cst =
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 6d340bc57fcd..b950128b17c0 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -53,6 +53,15 @@ llvm.mlir.global internal constant @int_gep() : !llvm.ptr<i32> {
   llvm.return %gepinit : !llvm.ptr<i32>
 }
 
+// CHECK: @string_array = internal constant [1 x ptr] [ptr @string_const]
+llvm.mlir.global internal constant @string_array() : !llvm.array<1 x ptr<i8>> {
+    %0 = llvm.mlir.undef : !llvm.array<1 x ptr<i8>>
+    %1 = llvm.mlir.addressof @string_const: !llvm.ptr<array<6 x i8>>
+    %2 = llvm.getelementptr %1[0, 0] : (!llvm.ptr<array<6 x i8>>) -> !llvm.ptr<i8>
+    %3 = llvm.insertvalue %2, %0[0] : !llvm.array<1 x ptr<i8>>
+    llvm.return %3 : !llvm.array<1 x ptr<i8>>
+}
+
 // CHECK{LITERAL}: @dense_float_vector = internal global <3 x float> <float 1.000000e+00, float 2.000000e+00, float 3.000000e+00>
 llvm.mlir.global internal @dense_float_vector(dense<[1.0, 2.0, 3.0]> : vector<3xf32>) : vector<3xf32>