Index: lib/Target/Mips/AsmParser/MipsAsmParser.cpp
===================================================================
--- lib/Target/Mips/AsmParser/MipsAsmParser.cpp
+++ lib/Target/Mips/AsmParser/MipsAsmParser.cpp
@@ -5448,11 +5448,9 @@
     return false;
   }
 
-  unsigned ATReg = getATReg(Loc);
-  if (!ATReg)
+  if (!getTargetStreamer().emitDirectiveCpRestore(
+          CpRestoreOffset, [&]() { return getATReg(Loc); }, Loc, STI))
     return true;
-
-  getTargetStreamer().emitDirectiveCpRestore(CpRestoreOffset, ATReg, Loc, STI);
   Parser.Lex(); // Consume the EndOfStatement.
   return false;
 }
Index: lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
===================================================================
--- lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
+++ lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
@@ -96,10 +96,11 @@
 void MipsTargetStreamer::emitDirectiveSetDsp() { forbidModuleDirective(); }
 void MipsTargetStreamer::emitDirectiveSetNoDsp() { forbidModuleDirective(); }
 void MipsTargetStreamer::emitDirectiveCpLoad(unsigned RegNo) {}
-void MipsTargetStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                SMLoc IDLoc,
-                                                const MCSubtargetInfo *STI) {
+bool MipsTargetStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
   forbidModuleDirective();
+  return true;
 }
 void MipsTargetStreamer::emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                                               const MCSymbol &Sym, bool IsReg) {
@@ -578,11 +579,12 @@
   forbidModuleDirective();
 }
 
-void MipsTargetAsmStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                   SMLoc IDLoc,
-                                                   const MCSubtargetInfo *STI) {
-  MipsTargetStreamer::emitDirectiveCpRestore(Offset, ATReg, IDLoc, STI);
+bool MipsTargetAsmStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
+  MipsTargetStreamer::emitDirectiveCpRestore(Offset, GetATReg, IDLoc, STI);
   OS << "\t.cprestore\t" << Offset << "\n";
+  return true;
 }
 
 void MipsTargetAsmStreamer::emitDirectiveCpsetup(unsigned RegNo,
@@ -1034,10 +1036,10 @@
   forbidModuleDirective();
 }
 
-void MipsTargetELFStreamer::emitDirectiveCpRestore(int Offset, unsigned ATReg,
-                                                   SMLoc IDLoc,
-                                                   const MCSubtargetInfo *STI) {
-  MipsTargetStreamer::emitDirectiveCpRestore(Offset, ATReg, IDLoc, STI);
+bool MipsTargetELFStreamer::emitDirectiveCpRestore(
+    int Offset, std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
+  MipsTargetStreamer::emitDirectiveCpRestore(Offset, GetATReg, IDLoc, STI);
   // .cprestore offset
   // When PIC mode is enabled and the O32 ABI is used, this directive expands
   // to:
@@ -1047,11 +1049,16 @@
   // Note that .cprestore is ignored if used with the N32 and N64 ABIs or if it
   // is used in non-PIC mode.
   if (!Pic || (getABI().IsN32() || getABI().IsN64()))
-    return;
+    return true;
+
+  unsigned ATReg = GetATReg();
+  if (!ATReg)
+    return false;
 
   // Store the $gp on the stack.
   emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, Offset, ATReg, IDLoc,
                          STI);
+  return true;
 }
 
 void MipsTargetELFStreamer::emitDirectiveCpsetup(unsigned RegNo,
Index: lib/Target/Mips/MipsTargetStreamer.h
===================================================================
--- lib/Target/Mips/MipsTargetStreamer.h
+++ lib/Target/Mips/MipsTargetStreamer.h
@@ -81,8 +81,9 @@
 
   // PIC support
   virtual void emitDirectiveCpLoad(unsigned RegNo);
-  virtual void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                                      const MCSubtargetInfo *STI);
+  virtual bool emitDirectiveCpRestore(int Offset,
+                                      std::function<unsigned()> GetATReg,
+                                      SMLoc IDLoc, const MCSubtargetInfo *STI);
   virtual void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                                     const MCSymbol &Sym, bool IsReg);
   virtual void emitDirectiveCpreturn(unsigned SaveLocation,
@@ -235,8 +236,8 @@
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
-  void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                              const MCSubtargetInfo *STI) override;
+  bool emitDirectiveCpRestore(int Offset, std::function<unsigned()> GetATReg,
+                              SMLoc IDLoc, const MCSubtargetInfo *STI) override;
   void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                             const MCSymbol &Sym, bool IsReg) override;
   void emitDirectiveCpreturn(unsigned SaveLocation,
@@ -290,8 +291,8 @@
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
-  void emitDirectiveCpRestore(int Offset, unsigned ATReg, SMLoc IDLoc,
-                              const MCSubtargetInfo *STI) override;
+  bool emitDirectiveCpRestore(int Offset, std::function<unsigned()> GetATReg,
+                              SMLoc IDLoc, const MCSubtargetInfo *STI) override;
   void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
                             const MCSymbol &Sym, bool IsReg) override;
   void emitDirectiveCpreturn(unsigned SaveLocation,