[Mlir-commits] [mlir] 095e6ac - [mlir] Fix memory explosion when converting global variable bodies in ModuleTranslation

[Kiran Chandramohan]

Author: Tung D. Le
Date: 2023-05-16T10:37:33Z
New Revision: 095e6ac9fd92d03dcb1e19b60cb06a8140aae69d

URL: https://github.com/llvm/llvm-project/commit/095e6ac9fd92d03dcb1e19b60cb06a8140aae69d
DIFF: https://github.com/llvm/llvm-project/commit/095e6ac9fd92d03dcb1e19b60cb06a8140aae69d.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    mlir/test/Target/LLVMIR/llvmir.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index e4cd94026a59e..70171e77eab90 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 6d340bc57fcd1..b950128b17c06 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>