diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index 79589bf13853c..91c5fbec888a3 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -331,6 +331,12 @@
       "group": "WM_DEBUG_RECENTS_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RecentsAnimation.java"
     },
+    "-1410260105": {
+      "message": "Schedule IME show for %s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "-1391944764": {
       "message": "SURFACE DESTROY: %s. %s",
       "level": "INFO",
@@ -733,12 +739,6 @@
       "group": "WM_DEBUG_ORIENTATION",
       "at": "com\/android\/server\/wm\/DisplayRotation.java"
     },
-    "-554834595": {
-      "message": "Display id=%d is frozen, return %d",
-      "level": "VERBOSE",
-      "group": "WM_DEBUG_ORIENTATION",
-      "at": "com\/android\/server\/wm\/DisplayContent.java"
-    },
     "-549028919": {
       "message": "enableScreenIfNeededLocked: mDisplayEnabled=%b mForceDisplayEnabled=%b mShowingBootMessages=%b mSystemBooted=%b. %s",
       "level": "INFO",
@@ -1063,6 +1063,12 @@
       "group": "WM_SHOW_TRANSACTIONS",
       "at": "com\/android\/server\/wm\/WindowSurfaceController.java"
     },
+    "140319294": {
+      "message": "IME target changed within ActivityRecord",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "146871307": {
       "message": "Tried to remove starting window but startingWindow was null: %s",
       "level": "VERBOSE",
@@ -1267,6 +1273,12 @@
       "group": "WM_ERROR",
       "at": "com\/android\/server\/wm\/WindowManagerService.java"
     },
+    "438102669": {
+      "message": "call showInsets(ime) on %s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "457951957": {
       "message": "\tNot visible=%s",
       "level": "DEBUG",
@@ -1387,6 +1399,12 @@
       "group": "WM_ERROR",
       "at": "com\/android\/server\/wm\/WindowManagerService.java"
     },
+    "632168013": {
+      "message": "dcTarget: %s mImeTargetFromIme: %s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "633654009": {
       "message": "SURFACE POS (setPositionInTransaction) @ (%f,%f): %s",
       "level": "INFO",
@@ -1693,6 +1711,12 @@
       "group": "WM_DEBUG_ADD_REMOVE",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
+    "1373000889": {
+      "message": "abortShowImePostLayout",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "1401700824": {
       "message": "Window drawn win=%s",
       "level": "DEBUG",
@@ -1987,6 +2011,12 @@
       "group": "WM_DEBUG_ADD_REMOVE",
       "at": "com\/android\/server\/wm\/WindowState.java"
     },
+    "1928325128": {
+      "message": "Run showImeRunner",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_IME",
+      "at": "com\/android\/server\/wm\/ImeInsetsSourceProvider.java"
+    },
     "1947239194": {
       "message": "Deferring rotation, still finishing previous rotation",
       "level": "VERBOSE",
@@ -2139,6 +2169,9 @@
     "WM_DEBUG_FOCUS_LIGHT": {
       "tag": "WindowManager"
     },
+    "WM_DEBUG_IME": {
+      "tag": "WindowManager"
+    },
     "WM_DEBUG_KEEP_SCREEN_ON": {
       "tag": "WindowManager"
     },
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 04fc1d2f202bd..e777c84380f98 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -122,7 +122,6 @@ import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_FOCUS;
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_FOCUS_LIGHT;
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_ORIENTATION;
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_SCREEN_ON;
-import static com.android.server.wm.ProtoLogGroup.WM_ERROR;
 import static com.android.server.wm.ProtoLogGroup.WM_SHOW_TRANSACTIONS;
 import static com.android.server.wm.RootWindowContainer.TAG_STATES;
 import static com.android.server.wm.WindowContainer.AnimationFlags.PARENTS;
@@ -3468,10 +3467,22 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         mInputMethodTargetWaitingAnim = targetWaitingAnim;
         assignWindowLayers(false /* setLayoutNeeded */);
         mInputMethodControlTarget = computeImeControlTarget();
-        mInsetsStateController.onImeTargetChanged(mInputMethodControlTarget);
+        mInsetsStateController.onImeControlTargetChanged(mInputMethodControlTarget);
         updateImeParent();
     }
 
+    /**
+     * IME control target is the window that controls the IME visibility and animation.
+     * This window is same as the window on which startInput is called.
+     * @param target the window that receives IME control.
+     *
+     * @see #getImeControlTarget()
+     */
+    void updateImeControlTarget(WindowState target) {
+        mInputMethodControlTarget = target;
+        mInsetsStateController.onImeControlTargetChanged(mInputMethodControlTarget);
+    }
+
     private void updateImeParent() {
         // Force attaching IME to the display when magnifying, or it would be magnified with
         // target app together.
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index fb97ecf8bbc4f..3ff1b95dc7166 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -16,9 +16,13 @@
 
 package com.android.server.wm;
 
+import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_IME;
+
 import android.view.InsetsSource;
 import android.view.WindowInsets;
 
+import com.android.server.protolog.common.ProtoLog;
+
 /**
  * Controller for IME inset source on the server. It's called provider as it provides the
  * {@link InsetsSource} to the client that uses it in {@link InsetsSourceConsumer}.
@@ -35,40 +39,40 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
     }
 
     /**
-     * Called when a layout pass has occurred.
-     */
-    void onPostLayout() {
-        super.onPostLayout();
-
-        if (mImeTargetFromIme != null
-                && isImeTargetFromDisplayContentAndImeSame()
-                && mWin != null
-                && mWin.isDrawnLw()
-                && !mWin.mGivenInsetsPending) {
-            mIsImeLayoutDrawn = true;
-        }
-    }
-
-    /**
-     * Called when Insets have been dispatched to client.
+     * Called when Insets have been dispatched to client. This gets called just after onPostLayout.
      */
     void onPostInsetsDispatched() {
-        if (mIsImeLayoutDrawn && mShowImeRunner != null) {
-            // Show IME if InputMethodService requested to be shown and it's layout has finished.
-            mShowImeRunner.run();
-        }
+        checkShowImePostLayout();
     }
 
     /**
      * Called from {@link WindowManagerInternal#showImePostLayout} when {@link InputMethodService}
      * requests to show IME on {@param imeTarget}.
+     *
      * @param imeTarget imeTarget on which IME request is coming from.
      */
     void scheduleShowImePostLayout(WindowState imeTarget) {
+        boolean targetChanged = mImeTargetFromIme != imeTarget
+                && mImeTargetFromIme != null && imeTarget != null && mShowImeRunner != null
+                && mImeTargetFromIme.mActivityRecord == imeTarget.mActivityRecord;
         mImeTargetFromIme = imeTarget;
+        if (targetChanged) {
+            // target changed, check if new target can show IME.
+            ProtoLog.d(WM_DEBUG_IME, "IME target changed within ActivityRecord");
+            checkShowImePostLayout();
+            // if IME cannot be shown at this time, it is scheduled to be shown.
+            // once window that called IMM.showSoftInput() and DisplayContent's ImeTarget match,
+            // it will be shown.
+            return;
+        }
+
+        ProtoLog.d(WM_DEBUG_IME, "Schedule IME show for %s", mImeTargetFromIme.getName());
         mShowImeRunner = () -> {
+            ProtoLog.d(WM_DEBUG_IME, "Run showImeRunner");
             // Target should still be the same.
             if (isImeTargetFromDisplayContentAndImeSame()) {
+                ProtoLog.d(WM_DEBUG_IME, "call showInsets(ime) on %s",
+                        mDisplayContent.mInputMethodControlTarget.getWindow().getName());
                 mDisplayContent.mInputMethodControlTarget.showInsets(
                         WindowInsets.Type.ime(), true /* fromIme */);
             }
@@ -76,10 +80,27 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
         };
     }
 
+    private void checkShowImePostLayout() {
+        // check if IME is drawn
+        if (mIsImeLayoutDrawn
+                || (mImeTargetFromIme != null
+                && isImeTargetFromDisplayContentAndImeSame()
+                && mWin != null
+                && mWin.isDrawnLw()
+                && !mWin.mGivenInsetsPending)) {
+            mIsImeLayoutDrawn = true;
+            // show IME if InputMethodService requested it to be shown.
+            if (mShowImeRunner != null) {
+                mShowImeRunner.run();
+            }
+        }
+    }
+
     /**
      * Abort any pending request to show IME post layout.
      */
     void abortShowImePostLayout() {
+        ProtoLog.d(WM_DEBUG_IME, "abortShowImePostLayout");
         mImeTargetFromIme = null;
         mIsImeLayoutDrawn = false;
         mShowImeRunner = null;
@@ -96,12 +117,14 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
         // TODO(b/139861270): Remove the child & sublayer check once IMMS is aware of
         //  actual IME target.
         final WindowState dcTarget = mDisplayContent.mInputMethodTarget;
-        if (dcTarget == null) {
+        if (dcTarget == null || mImeTargetFromIme == null) {
             return false;
         }
+        ProtoLog.d(WM_DEBUG_IME, "dcTarget: %s mImeTargetFromIme: %s",
+                dcTarget.getName(), mImeTargetFromIme.getName());
+
         return (!dcTarget.isClosing() && mImeTargetFromIme == dcTarget)
                 || (mImeTargetFromIme != null && dcTarget.getParentWindow() == mImeTargetFromIme
-                        && dcTarget.mSubLayer > mImeTargetFromIme.mSubLayer);
+                && dcTarget.mSubLayer > mImeTargetFromIme.mSubLayer);
     }
-
 }
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index b2234d17984e4..aa30b09f179e7 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -167,7 +167,7 @@ class InsetsStateController {
         return mTypeFakeControlTargetMap.get(type) == target;
     }
 
-    void onImeTargetChanged(@Nullable InsetsControlTarget imeTarget) {
+    void onImeControlTargetChanged(@Nullable InsetsControlTarget imeTarget) {
         onControlChanged(ITYPE_IME, imeTarget);
         notifyPendingInsetsControlChanged();
     }
diff --git a/services/core/java/com/android/server/wm/ProtoLogGroup.java b/services/core/java/com/android/server/wm/ProtoLogGroup.java
index 5c2830e13b751..51725cecbc744 100644
--- a/services/core/java/com/android/server/wm/ProtoLogGroup.java
+++ b/services/core/java/com/android/server/wm/ProtoLogGroup.java
@@ -59,6 +59,8 @@ public enum ProtoLogGroup implements IProtoLogGroup {
             Consts.TAG_WM),
     WM_DEBUG_WINDOW_MOVEMENT(Consts.ENABLE_DEBUG, Consts.ENABLE_LOG_TO_PROTO_DEBUG, false,
             Consts.TAG_WM),
+    WM_DEBUG_IME(Consts.ENABLE_DEBUG, Consts.ENABLE_LOG_TO_PROTO_DEBUG, false,
+            Consts.TAG_WM),
     TEST_GROUP(true, true, false, "WindowManagetProtoLogTest");
 
     private final boolean mEnabled;
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index f15e9eb80403f..229bd2f4afe1a 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7298,6 +7298,10 @@ public class WindowManagerService extends IWindowManager.Stub
                 Slog.w(TAG_WM, "updateInputMethodTargetWindow: imeToken=" + imeToken
                         + " imeTargetWindowToken=" + imeTargetWindowToken);
             }
+            final WindowState imeTarget = mWindowMap.get(imeTargetWindowToken);
+            if (imeTarget != null) {
+                imeTarget.getDisplayContent().updateImeControlTarget(imeTarget);
+            }
         }
 
         @Override