From c8a7f14b276fcea68c50ee6b9007680867d2393e Mon Sep 17 00:00:00 2001
From: weiwei chen <weiwei.chen@modular.com>
Date: Thu, 7 Nov 2024 10:05:19 -0500
Subject: [PATCH] [Backend] Add clearSubtargetMap API for TargetMachine.
 (#112383)

- [x] Add `clearSubtargetInfo` API to TargetMachine and each backend to
make it possible to release memory used in each backend's
`SubtargetInfo` map if needed. Keep this API as `protected` so that it
will be used with precautions.
---
 llvm/include/llvm/Target/TargetMachine.h         | 3 +++
 llvm/lib/Target/AArch64/AArch64TargetMachine.cpp | 2 ++
 llvm/lib/Target/AArch64/AArch64TargetMachine.h   | 3 +++
 llvm/lib/Target/ARM/ARMTargetMachine.cpp         | 2 ++
 llvm/lib/Target/ARM/ARMTargetMachine.h           | 3 +++
 llvm/lib/Target/X86/X86TargetMachine.cpp         | 2 ++
 llvm/lib/Target/X86/X86TargetMachine.h           | 3 +++
 7 files changed, 18 insertions(+)

diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index c3e9d41315f61..fa3ab58a21ddc 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -448,6 +448,9 @@ class LLVMTargetMachine : public TargetMachine {
 
   void initAsmInfo();
 
+  /// Reset internal state.
+  virtual void reset() {};
+
 public:
   /// Get a TargetTransformInfo implementation for the target.
   ///
diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
index c7bd0390b6562..e24a874e74970 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
@@ -272,6 +272,8 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAArch64Target() {
   initializeAArch64GlobalsTaggingPass(*PR);
 }
 
+void AArch64TargetMachine::reset() { SubtargetMap.clear(); }
+
 //===----------------------------------------------------------------------===//
 // AArch64 Lowering public interface.
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.h b/llvm/lib/Target/AArch64/AArch64TargetMachine.h
index 1a470ca87127c..e4d0aff50d8f6 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.h
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.h
@@ -26,6 +26,9 @@ class AArch64TargetMachine : public LLVMTargetMachine {
   std::unique_ptr<TargetLoweringObjectFile> TLOF;
   mutable StringMap<std::unique_ptr<AArch64Subtarget>> SubtargetMap;
 
+  /// Reset internal state.
+  void reset() override;
+
 public:
   AArch64TargetMachine(const Target &T, const Triple &TT, StringRef CPU,
                        StringRef FS, const TargetOptions &Options,
diff --git a/llvm/lib/Target/ARM/ARMTargetMachine.cpp b/llvm/lib/Target/ARM/ARMTargetMachine.cpp
index a58c63dcf762d..74a63361d341b 100644
--- a/llvm/lib/Target/ARM/ARMTargetMachine.cpp
+++ b/llvm/lib/Target/ARM/ARMTargetMachine.cpp
@@ -642,3 +642,5 @@ bool ARMBaseTargetMachine::parseMachineFunctionInfo(
   MF.getInfo<ARMFunctionInfo>()->initializeBaseYamlFields(YamlMFI);
   return false;
 }
+
+void ARMBaseTargetMachine::reset() { SubtargetMap.clear(); }
diff --git a/llvm/lib/Target/ARM/ARMTargetMachine.h b/llvm/lib/Target/ARM/ARMTargetMachine.h
index 75ee50f0e93c8..5b3594a4dcca8 100644
--- a/llvm/lib/Target/ARM/ARMTargetMachine.h
+++ b/llvm/lib/Target/ARM/ARMTargetMachine.h
@@ -38,6 +38,9 @@ class ARMBaseTargetMachine : public LLVMTargetMachine {
   bool isLittle;
   mutable StringMap<std::unique_ptr<ARMSubtarget>> SubtargetMap;
 
+  /// Reset internal state.
+  void reset() override;
+
 public:
   ARMBaseTargetMachine(const Target &T, const Triple &TT, StringRef CPU,
                        StringRef FS, const TargetOptions &Options,
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 4ba0ac11d2095..dcdd059b4c0cf 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -376,6 +376,8 @@ bool X86TargetMachine::isNoopAddrSpaceCast(unsigned SrcAS,
   return SrcAS < 256 && DestAS < 256;
 }
 
+void X86TargetMachine::reset() { SubtargetMap.clear(); }
+
 //===----------------------------------------------------------------------===//
 // X86 TTI query.
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Target/X86/X86TargetMachine.h b/llvm/lib/Target/X86/X86TargetMachine.h
index ec4a93e9c9d4b..d27c8a4d10116 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.h
+++ b/llvm/lib/Target/X86/X86TargetMachine.h
@@ -31,6 +31,9 @@ class X86TargetMachine final : public LLVMTargetMachine {
   // True if this is used in JIT.
   bool IsJIT;
 
+  /// Reset internal state.
+  void reset() override;
+
 public:
   X86TargetMachine(const Target &T, const Triple &TT, StringRef CPU,
                    StringRef FS, const TargetOptions &Options,