Index: lib/Target/Mips/AsmParser/MipsAsmParser.cpp
===================================================================
--- lib/Target/Mips/AsmParser/MipsAsmParser.cpp
+++ lib/Target/Mips/AsmParser/MipsAsmParser.cpp
@@ -2633,16 +2633,17 @@
   unsigned SrcReg = Inst.getOperand(0).getReg();
   unsigned BaseReg = Inst.getOperand(1).getReg();
 
-  unsigned ATReg = getATReg(IDLoc);
-  if (!ATReg)
-    return;
-
   if (IsImmOpnd) {
     TOut.emitStoreWithImmOffset(Inst.getOpcode(), SrcReg, BaseReg,
-                                Inst.getOperand(2).getImm(), ATReg, IDLoc, STI);
+                                Inst.getOperand(2).getImm(),
+                                [&]() { return getATReg(IDLoc); }, IDLoc, STI);
     return;
   }
 
+  unsigned ATReg = getATReg(IDLoc);
+  if (!ATReg)
+    return;
+
   const MCExpr *ExprOffset = Inst.getOperand(2).getExpr();
   MCOperand LoOperand = MCOperand::createExpr(
       MipsMCExpr::create(MipsMCExpr::MEK_LO, ExprOffset, getContext()));
@@ -3626,12 +3627,8 @@
     return;
   }
 
-  unsigned ATReg = getATReg(IDLoc);
-  if (!ATReg)
-    return;
-
-  TOut.emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, StackOffset, ATReg,
-                              IDLoc, STI);
+  TOut.emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, StackOffset,
+                              [&]() { return getATReg(IDLoc); }, IDLoc, STI);
 }
 
 unsigned MipsAsmParser::checkTargetMatchPredicate(MCInst &Inst) {
Index: lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
===================================================================
--- lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
+++ lib/Target/Mips/MCTargetDesc/MipsTargetStreamer.cpp
@@ -228,7 +228,8 @@
 /// Emit a store instruction with an immediate offset.
 void MipsTargetStreamer::emitStoreWithImmOffset(
     unsigned Opcode, unsigned SrcReg, unsigned BaseReg, int64_t Offset,
-    unsigned ATReg, SMLoc IDLoc, const MCSubtargetInfo *STI) {
+    std::function<unsigned()> GetATReg, SMLoc IDLoc,
+    const MCSubtargetInfo *STI) {
   if (isInt<16>(Offset)) {
     emitRRI(Opcode, SrcReg, BaseReg, Offset, IDLoc, STI);
     return;
@@ -238,6 +239,10 @@
   //                      add $at, $at, $8
   //                      sw $8, %lo(offset)($at)
 
+  unsigned ATReg = GetATReg();
+  if (!ATReg)
+    return;
+
   unsigned LoOffset = Offset & 0x0000ffff;
   unsigned HiOffset = (Offset & 0xffff0000) >> 16;
 
@@ -1051,12 +1056,8 @@
   if (!Pic || (getABI().IsN32() || getABI().IsN64()))
     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,
+  emitStoreWithImmOffset(Mips::SW, Mips::GP, Mips::SP, Offset, GetATReg, IDLoc,
                          STI);
   return true;
 }
Index: lib/Target/Mips/MipsTargetStreamer.h
===================================================================
--- lib/Target/Mips/MipsTargetStreamer.h
+++ lib/Target/Mips/MipsTargetStreamer.h
@@ -122,9 +122,18 @@
   void emitEmptyDelaySlot(bool hasShortDelaySlot, SMLoc IDLoc,
                           const MCSubtargetInfo *STI);
   void emitNop(SMLoc IDLoc, const MCSubtargetInfo *STI);
+
+  /// Emit a store instruction with an offset. If the offset is out of range
+  /// then it will be synthesized using the assembler temporary.
+  ///
+  /// GetATReg() is a callback that can be used to obtain the current assembler
+  /// temporary and is only called when the assembler temporary is required. It
+  /// must handle the case where no assembler temporary is available (typically
+  /// by reporting an error).
   void emitStoreWithImmOffset(unsigned Opcode, unsigned SrcReg,
-                              unsigned BaseReg, int64_t Offset, unsigned ATReg,
-                              SMLoc IDLoc, const MCSubtargetInfo *STI);
+                              unsigned BaseReg, int64_t Offset,
+                              std::function<unsigned()> GetATReg, SMLoc IDLoc,
+                              const MCSubtargetInfo *STI);
   void emitStoreWithSymOffset(unsigned Opcode, unsigned SrcReg,
                               unsigned BaseReg, MCOperand &HiOperand,
                               MCOperand &LoOperand, unsigned ATReg, SMLoc IDLoc,
@@ -237,6 +246,14 @@
 
   // PIC support
   void emitDirectiveCpLoad(unsigned RegNo) override;
+
+  /// Emit a .cprestore directive.  If the offset is out of range then it will
+  /// be synthesized using the assembler temporary.
+  ///
+  /// GetATReg() is a callback that can be used to obtain the current assembler
+  /// temporary and is only called when the assembler temporary is required. It
+  /// must handle the case where no assembler temporary is available (typically
+  /// by reporting an error).
   bool emitDirectiveCpRestore(int Offset, std::function<unsigned()> GetATReg,
                               SMLoc IDLoc, const MCSubtargetInfo *STI) override;
   void emitDirectiveCpsetup(unsigned RegNo, int RegOrOffset,
Index: test/MC/Mips/cprestore-noreorder-noat.s
===================================================================
--- test/MC/Mips/cprestore-noreorder-noat.s
+++ test/MC/Mips/cprestore-noreorder-noat.s
@@ -24,7 +24,7 @@
 
   .cpload $25
   .cprestore 8
-# O32: :[[@LINE-1]]:3: error: pseudo-instruction requires $at, which is not available
+# O32-NOT: error: pseudo-instruction requires $at, which is not available
 # N32-NOT: error: pseudo-instruction requires $at, which is not available
 # N64-NOT: error: pseudo-instruction requires $at, which is not available
 # NO-STORE-NOT: sw  $gp, 8($sp)
@@ -34,3 +34,22 @@
   jal foo
 
   .end foo
+
+  .ent bar
+bar:
+  .frame  $sp, 0, $ra
+  .set noreorder
+  .set noat
+
+  .cpload $25
+  .cprestore 65536
+# O32: :[[@LINE-1]]:3: error: pseudo-instruction requires $at, which is not available
+# N32-NOT: error: pseudo-instruction requires $at, which is not available
+# N64-NOT: error: pseudo-instruction requires $at, which is not available
+# NO-STORE-NOT: sw $gp,
+
+  jal $25
+  jal $4, $25
+  jal bar
+
+  .end bar