6 files changed, 34 insertions, 3 deletions
diff --git a/llvm/include/llvm-c/lto.h b/llvm/include/llvm-c/lto.h
index 1acd610f70a..090cd34af4e 100644
--- a/llvm/include/llvm-c/lto.h
+++ b/llvm/include/llvm-c/lto.h
@@ -44,7 +44,7 @@ typedef bool lto_bool_t;
  * @{
  */
 
-#define LTO_API_VERSION 22
+#define LTO_API_VERSION 23
 
 /**
  * \since prior to LTO_API_VERSION=3
@@ -828,6 +828,16 @@ extern void thinlto_codegen_set_cache_size_bytes(thinlto_code_gen_t cg,
                                                  unsigned max_size_bytes);
 
 /**
+ * Same as thinlto_codegen_set_cache_size_bytes, except the maximum size is in
+ * megabytes (2^20 bytes).
+ *
+ * \since LTO_API_VERSION=23
+ */
+extern void
+thinlto_codegen_set_cache_size_megabytes(thinlto_code_gen_t cg,
+                                         unsigned max_size_megabytes);
+
+/**
  * Sets the maximum number of files in the cache directory. An unspecified
  * default value will be applied. A value of 0 will be ignored.
  *
diff --git a/llvm/include/llvm/LTO/legacy/ThinLTOCodeGenerator.h b/llvm/include/llvm/LTO/legacy/ThinLTOCodeGenerator.h
index 8916ca6d3b3..2ff317a3227 100644
--- a/llvm/include/llvm/LTO/legacy/ThinLTOCodeGenerator.h
+++ b/llvm/include/llvm/LTO/legacy/ThinLTOCodeGenerator.h
@@ -187,7 +187,7 @@ public:
   /// Cache policy: the maximum size for the cache directory in bytes. A value
   /// over the amount of available space on the disk will be reduced to the
   /// amount of available space. A value of 0 will be ignored.
-  void setCacheMaxSizeBytes(unsigned MaxSizeBytes) {
+  void setCacheMaxSizeBytes(uint64_t MaxSizeBytes) {
     if (MaxSizeBytes)
       CacheOptions.Policy.MaxSizeBytes = MaxSizeBytes;
   }
diff --git a/llvm/test/ThinLTO/X86/cache.ll b/llvm/test/ThinLTO/X86/cache.ll
index 8cc015365f4..08eb5733dd6 100644
--- a/llvm/test/ThinLTO/X86/cache.ll
+++ b/llvm/test/ThinLTO/X86/cache.ll
@@ -121,6 +121,19 @@
 ; RUN: not ls %t.cache/llvmcache-foo-100k
 ; RUN: not ls %t.cache/llvmcache-foo-77k
 
+; Verify that specifying a max size > 4GB for the cache directory does not
+; prematurely prune, due to an integer overflow.
+; RUN: rm -Rf %t.cache && mkdir %t.cache
+; RUN: %python -c "with open(r'%t.cache/llvmcache-foo-10', 'w') as file: file.truncate(10)"
+; RUN: llvm-lto -thinlto-action=run -exported-symbol=globalfunc %t2.bc %t.bc -thinlto-cache-dir %t.cache --thinlto-cache-max-size-bytes 4294967297
+; RUN: ls %t.cache/llvmcache-foo-10
+
+; Verify that negative numbers aren't accepted for the
+; --thinlto-cache-max-size-bytes switch
+; RUN: rm -Rf %t.cache && mkdir %t.cache
+; RUN: not llvm-lto %t.bc --thinlto-cache-max-size-bytes -1 2>&1 | FileCheck %s
+; CHECK: -thinlto-cache-max-size-bytes option: '-1' value invalid
+
 ; Verify that specifying max number of files in the cache directory prunes
 ; it to this amount, removing the oldest files first.
 ; RUN: rm -Rf %t.cache && mkdir %t.cache
diff --git a/llvm/tools/llvm-lto/llvm-lto.cpp b/llvm/tools/llvm-lto/llvm-lto.cpp
index 75668a9dd8b..25f086fdac9 100644
--- a/llvm/tools/llvm-lto/llvm-lto.cpp
+++ b/llvm/tools/llvm-lto/llvm-lto.cpp
@@ -158,7 +158,7 @@ static cl::opt<int>
     ThinLTOCachePruningInterval("thinlto-cache-pruning-interval",
     cl::init(1200), cl::desc("Set ThinLTO cache pruning interval."));
 
-static cl::opt<int>
+static cl::opt<unsigned long long>
     ThinLTOCacheMaxSizeBytes("thinlto-cache-max-size-bytes",
     cl::desc("Set ThinLTO cache pruning directory maximum size in bytes."));
 
diff --git a/llvm/tools/lto/lto.cpp b/llvm/tools/lto/lto.cpp
index f9b518ca039..0a3ba98957e 100644
--- a/llvm/tools/lto/lto.cpp
+++ b/llvm/tools/lto/lto.cpp
@@ -591,6 +591,13 @@ void thinlto_codegen_set_cache_size_bytes(
   return unwrap(cg)->setCacheMaxSizeBytes(MaxSizeBytes);
 }
 
+void thinlto_codegen_set_cache_size_megabytes(
+    thinlto_code_gen_t cg, unsigned MaxSizeMegabytes) {
+  uint64_t MaxSizeBytes = MaxSizeMegabytes;
+  MaxSizeBytes *= 1024 * 1024;
+  return unwrap(cg)->setCacheMaxSizeBytes(MaxSizeBytes);
+}
+
 void thinlto_codegen_set_cache_size_files(
     thinlto_code_gen_t cg, unsigned MaxSizeFiles) {
   return unwrap(cg)->setCacheMaxSizeFiles(MaxSizeFiles);
diff --git a/llvm/tools/lto/lto.exports b/llvm/tools/lto/lto.exports
index abde3894cb2..c7a33868be3 100644
--- a/llvm/tools/lto/lto.exports
+++ b/llvm/tools/lto/lto.exports
@@ -58,6 +58,7 @@ thinlto_codegen_set_cache_pruning_interval
 thinlto_codegen_set_cache_entry_expiration
 thinlto_codegen_set_final_cache_size_relative_to_available_space
 thinlto_codegen_set_cache_size_bytes
+thinlto_codegen_set_cache_size_megabytes
 thinlto_codegen_set_cache_size_files
 thinlto_codegen_set_savetemps_dir
 thinlto_codegen_set_cpu