diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
index 54bb777b2d0c..725e15178cdb 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
@@ -749,43 +749,6 @@ void dpu_crtc_commit_kickoff(struct drm_crtc *crtc)
 	DPU_ATRACE_END("crtc_commit");
 }
 
-/**
- * _dpu_crtc_vblank_enable_no_lock - update power resource and vblank request
- * @dpu_crtc: Pointer to dpu crtc structure
- * @enable: Whether to enable/disable vblanks
- */
-static void _dpu_crtc_vblank_enable_no_lock(
-		struct dpu_crtc *dpu_crtc, bool enable)
-{
-	struct drm_crtc *crtc = &dpu_crtc->base;
-	struct drm_device *dev = crtc->dev;
-	struct drm_encoder *enc;
-
-	if (enable) {
-		list_for_each_entry(enc, &dev->mode_config.encoder_list, head) {
-			if (enc->crtc != crtc)
-				continue;
-
-			trace_dpu_crtc_vblank_enable(DRMID(&dpu_crtc->base),
-						     DRMID(enc), enable,
-						     dpu_crtc);
-
-			dpu_encoder_assign_crtc(enc, crtc);
-		}
-	} else {
-		list_for_each_entry(enc, &dev->mode_config.encoder_list, head) {
-			if (enc->crtc != crtc)
-				continue;
-
-			trace_dpu_crtc_vblank_enable(DRMID(&dpu_crtc->base),
-						     DRMID(enc), enable,
-						     dpu_crtc);
-
-			dpu_encoder_assign_crtc(enc, NULL);
-		}
-	}
-}
-
 /**
  * dpu_crtc_duplicate_state - state duplicate hook
  * @crtc: Pointer to drm crtc structure
@@ -839,6 +802,10 @@ static void dpu_crtc_disable(struct drm_crtc *crtc,
 	/* Disable/save vblank irq handling */
 	drm_crtc_vblank_off(crtc);
 
+	drm_for_each_encoder_mask(encoder, crtc->dev,
+				  old_crtc_state->encoder_mask)
+		dpu_encoder_assign_crtc(encoder, NULL);
+
 	mutex_lock(&dpu_crtc->crtc_lock);
 
 	/* wait for frame_event_done completion */
@@ -848,9 +815,6 @@ static void dpu_crtc_disable(struct drm_crtc *crtc,
 				atomic_read(&dpu_crtc->frame_pending));
 
 	trace_dpu_crtc_disable(DRMID(crtc), false, dpu_crtc);
-	if (dpu_crtc->enabled && dpu_crtc->vblank_requested) {
-		_dpu_crtc_vblank_enable_no_lock(dpu_crtc, false);
-	}
 	dpu_crtc->enabled = false;
 
 	if (atomic_read(&dpu_crtc->frame_pending)) {
@@ -908,13 +872,13 @@ static void dpu_crtc_enable(struct drm_crtc *crtc,
 
 	mutex_lock(&dpu_crtc->crtc_lock);
 	trace_dpu_crtc_enable(DRMID(crtc), true, dpu_crtc);
-	if (!dpu_crtc->enabled && dpu_crtc->vblank_requested) {
-		_dpu_crtc_vblank_enable_no_lock(dpu_crtc, true);
-	}
 	dpu_crtc->enabled = true;
 
 	mutex_unlock(&dpu_crtc->crtc_lock);
 
+	drm_for_each_encoder_mask(encoder, crtc->dev, crtc->state->encoder_mask)
+		dpu_encoder_assign_crtc(encoder, crtc);
+
 	/* Enable/restore vblank irq handling */
 	drm_crtc_vblank_on(crtc);
 }
@@ -1159,10 +1123,33 @@ static int dpu_crtc_atomic_check(struct drm_crtc *crtc,
 int dpu_crtc_vblank(struct drm_crtc *crtc, bool en)
 {
 	struct dpu_crtc *dpu_crtc = to_dpu_crtc(crtc);
+	struct drm_encoder *enc;
 
-	mutex_lock(&dpu_crtc->crtc_lock);
 	trace_dpu_crtc_vblank(DRMID(&dpu_crtc->base), en, dpu_crtc);
-	_dpu_crtc_vblank_enable_no_lock(dpu_crtc, en);
+
+	/*
+	 * Normally we would iterate through encoder_mask in crtc state to find
+	 * attached encoders. In this case, we might be disabling vblank _after_
+	 * encoder_mask has been cleared.
+	 *
+	 * Instead, we "assign" a crtc to the encoder in enable and clear it in
+	 * disable (which is also after encoder_mask is cleared). So instead of
+	 * using encoder mask, we'll ask the encoder to toggle itself iff it's
+	 * currently assigned to our crtc.
+	 *
+	 * Note also that this function cannot be called while crtc is disabled
+	 * since we use drm_crtc_vblank_on/off. So we don't need to worry
+	 * about the assigned crtcs being inconsistent with the current state
+	 * (which means no need to worry about modeset locks).
+	 */
+	list_for_each_entry(enc, &crtc->dev->mode_config.encoder_list, head) {
+		trace_dpu_crtc_vblank_enable(DRMID(crtc), DRMID(enc), en,
+					     dpu_crtc);
+
+		dpu_encoder_toggle_vblank_for_crtc(enc, crtc, en);
+	}
+
+	mutex_lock(&dpu_crtc->crtc_lock);
 	dpu_crtc->vblank_requested = en;
 	mutex_unlock(&dpu_crtc->crtc_lock);
 
diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
index fd6514f681ae..90b77d98dde8 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
@@ -1267,22 +1267,29 @@ void dpu_encoder_assign_crtc(struct drm_encoder *drm_enc, struct drm_crtc *crtc)
 {
 	struct dpu_encoder_virt *dpu_enc = to_dpu_encoder_virt(drm_enc);
 	unsigned long lock_flags;
-	bool enable;
-	int i;
-
-	enable = crtc ? true : false;
-
-	if (!drm_enc) {
-		DPU_ERROR("invalid encoder\n");
-		return;
-	}
-	trace_dpu_enc_vblank_cb(DRMID(drm_enc), enable);
 
 	spin_lock_irqsave(&dpu_enc->enc_spinlock, lock_flags);
 	/* crtc should always be cleared before re-assigning */
 	WARN_ON(crtc && dpu_enc->crtc);
 	dpu_enc->crtc = crtc;
 	spin_unlock_irqrestore(&dpu_enc->enc_spinlock, lock_flags);
+}
+
+void dpu_encoder_toggle_vblank_for_crtc(struct drm_encoder *drm_enc,
+					struct drm_crtc *crtc, bool enable)
+{
+	struct dpu_encoder_virt *dpu_enc = to_dpu_encoder_virt(drm_enc);
+	unsigned long lock_flags;
+	int i;
+
+	trace_dpu_enc_vblank_cb(DRMID(drm_enc), enable);
+
+	spin_lock_irqsave(&dpu_enc->enc_spinlock, lock_flags);
+	if (dpu_enc->crtc != crtc) {
+		spin_unlock_irqrestore(&dpu_enc->enc_spinlock, lock_flags);
+		return;
+	}
+	spin_unlock_irqrestore(&dpu_enc->enc_spinlock, lock_flags);
 
 	for (i = 0; i < dpu_enc->num_phys_encs; i++) {
 		struct dpu_encoder_phys *phys = dpu_enc->phys_encs[i];
diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h
index be1d80867834..6896ea2ab854 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h
@@ -62,6 +62,16 @@ void dpu_encoder_get_hw_resources(struct drm_encoder *encoder,
 void dpu_encoder_assign_crtc(struct drm_encoder *encoder,
 			     struct drm_crtc *crtc);
 
+/**
+ * dpu_encoder_toggle_vblank_for_crtc - Toggles vblank interrupts on or off if
+ *	the encoder is assigned to the given crtc
+ * @encoder:	encoder pointer
+ * @crtc:	crtc pointer
+ * @enable:	true if vblank should be enabled
+ */
+void dpu_encoder_toggle_vblank_for_crtc(struct drm_encoder *encoder,
+					struct drm_crtc *crtc, bool enable);
+
 /**
  * dpu_encoder_register_frame_event_callback - provide callback to encoder that
  *	will be called after the request is complete, or other events.
