diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 1c811cb6e96cf..2ba1e016e03d2 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -172,6 +172,6 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     }
 
     private InputMethodManager getImm() {
-        return mController.getViewRoot().mContext.getSystemService(InputMethodManager.class);
+        return mController.getViewRoot().mDisplayContext.getSystemService(InputMethodManager.class);
     }
 }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 479d8068250c5..f3b7ad5e557c5 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -302,7 +302,13 @@ public final class ViewRootImpl implements ViewParent,
     @GuardedBy("mWindowCallbacks")
     final ArrayList<WindowCallbacks> mWindowCallbacks = new ArrayList<>();
     @UnsupportedAppUsage
-    public final Context mContext;
+    final Context mContext;
+    /**
+     * TODO(b/116349163): Check if we can merge this into {@link #mContext}.
+     * @hide
+     */
+    @NonNull
+    public Context mDisplayContext;
 
     @UnsupportedAppUsage
     final IWindowSession mWindowSession;
@@ -612,6 +618,7 @@ public final class ViewRootImpl implements ViewParent,
 
     public ViewRootImpl(Context context, Display display) {
         mContext = context;
+        mDisplayContext = context.createDisplayContext(display);
         mWindowSession = WindowManagerGlobal.getWindowSession();
         mDisplay = display;
         mBasePackageName = context.getBasePackageName();
@@ -1396,7 +1403,7 @@ public final class ViewRootImpl implements ViewParent,
         } else {
             mDisplay = preferredDisplay;
         }
-        mContext.updateDisplay(mDisplay.getDisplayId());
+        mDisplayContext = mContext.createDisplayContext(mDisplay);
     }
 
     void pokeDrawLockIfNeeded() {
@@ -2742,7 +2749,7 @@ public final class ViewRootImpl implements ViewParent,
                     .mayUseInputMethod(mWindowAttributes.flags);
             if (imTarget != mLastWasImTarget) {
                 mLastWasImTarget = imTarget;
-                InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
+                InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
                 if (imm != null && imTarget) {
                     imm.onPreWindowFocus(mView, hasWindowFocus);
                     imm.onPostWindowFocus(mView, mView.findFocus(),
@@ -2918,7 +2925,7 @@ public final class ViewRootImpl implements ViewParent,
             mLastWasImTarget = WindowManager.LayoutParams
                     .mayUseInputMethod(mWindowAttributes.flags);
 
-            InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
+            InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
             if (imm != null && mLastWasImTarget && !isInLocalFocusMode()) {
                 imm.onPreWindowFocus(mView, hasWindowFocus);
             }
@@ -4692,7 +4699,8 @@ public final class ViewRootImpl implements ViewParent,
                     enqueueInputEvent(event, null, 0, true);
                 } break;
                 case MSG_CHECK_FOCUS: {
-                    InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
+                    InputMethodManager imm =
+                            mDisplayContext.getSystemService(InputMethodManager.class);
                     if (imm != null) {
                         imm.checkFocus();
                     }
@@ -5237,7 +5245,7 @@ public final class ViewRootImpl implements ViewParent,
         @Override
         protected int onProcess(QueuedInputEvent q) {
             if (mLastWasImTarget && !isInLocalFocusMode()) {
-                InputMethodManager imm = mContext.getSystemService(InputMethodManager.class);
+                InputMethodManager imm = mDisplayContext.getSystemService(InputMethodManager.class);
                 if (imm != null) {
                     final InputEvent event = q.mEvent;
                     if (DEBUG_IMF) Log.v(mTag, "Sending input event to IME: " + event);
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 2cfd6226351e2..5e00425407ba8 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -510,7 +510,7 @@ public final class InputMethodManager {
             return null;
         }
         final InputMethodManager fallbackImm =
-                viewRootImpl.mContext.getSystemService(InputMethodManager.class);
+                viewRootImpl.mDisplayContext.getSystemService(InputMethodManager.class);
         if (fallbackImm == null) {
             Log.e(TAG, "b/117267690: Failed to get non-null fallback IMM. view=" + view);
             return null;