diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
index 141ed1b0e90a..4d7f9ff1e9f4 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.c
@@ -293,9 +293,8 @@ enum dpu_intf_mode dpu_crtc_get_intf_mode(struct drm_crtc *crtc)
 	return INTF_MODE_NONE;
 }
 
-static void dpu_crtc_vblank_cb(void *data)
+void dpu_crtc_vblank_callback(struct drm_crtc *crtc)
 {
-	struct drm_crtc *crtc = (struct drm_crtc *)data;
 	struct dpu_crtc *dpu_crtc = to_dpu_crtc(crtc);
 
 	/* keep statistics on vblank callback - with auto reset via debugfs */
@@ -771,8 +770,7 @@ static void _dpu_crtc_vblank_enable_no_lock(
 						     DRMID(enc), enable,
 						     dpu_crtc);
 
-			dpu_encoder_register_vblank_callback(enc,
-					dpu_crtc_vblank_cb, (void *)crtc);
+			dpu_encoder_assign_crtc(enc, crtc);
 		}
 	} else {
 		list_for_each_entry(enc, &dev->mode_config.encoder_list, head) {
@@ -783,7 +781,7 @@ static void _dpu_crtc_vblank_enable_no_lock(
 						     DRMID(enc), enable,
 						     dpu_crtc);
 
-			dpu_encoder_register_vblank_callback(enc, NULL, NULL);
+			dpu_encoder_assign_crtc(enc, NULL);
 		}
 	}
 }
diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.h b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.h
index 93d21a61a040..54595cc29be5 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.h
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_crtc.h
@@ -270,6 +270,12 @@ static inline int dpu_crtc_frame_pending(struct drm_crtc *crtc)
  */
 int dpu_crtc_vblank(struct drm_crtc *crtc, bool en);
 
+/**
+ * dpu_crtc_vblank_callback - called on vblank irq, issues completion events
+ * @crtc: Pointer to drm crtc object
+ */
+void dpu_crtc_vblank_callback(struct drm_crtc *crtc);
+
 /**
  * dpu_crtc_commit_kickoff - trigger kickoff of the commit for this crtc
  * @crtc: Pointer to drm crtc object
diff --git a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
index d89ac520f7e6..fd6514f681ae 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.c
@@ -142,9 +142,11 @@ enum dpu_enc_rc_states {
  * @intfs_swapped	Whether or not the phys_enc interfaces have been swapped
  *			for partial update right-only cases, such as pingpong
  *			split where virtual pingpong does not generate IRQs
- * @crtc_vblank_cb:	Callback into the upper layer / CRTC for
- *			notification of the VBLANK
- * @crtc_vblank_cb_data:	Data from upper layer for VBLANK notification
+ * @crtc:		Pointer to the currently assigned crtc. Normally you
+ *			would use crtc->state->encoder_mask to determine the
+ *			link between encoder/crtc. However in this case we need
+ *			to track crtc in the disable() hook which is called
+ *			_after_ encoder_mask is cleared.
  * @crtc_kickoff_cb:		Callback into CRTC that will flush & start
  *				all CTL paths
  * @crtc_kickoff_cb_data:	Opaque user data given to crtc_kickoff_cb
@@ -186,8 +188,7 @@ struct dpu_encoder_virt {
 
 	bool intfs_swapped;
 
-	void (*crtc_vblank_cb)(void *);
-	void *crtc_vblank_cb_data;
+	struct drm_crtc *crtc;
 
 	struct dentry *debugfs_root;
 	struct mutex enc_lock;
@@ -1241,8 +1242,8 @@ static void dpu_encoder_vblank_callback(struct drm_encoder *drm_enc,
 	dpu_enc = to_dpu_encoder_virt(drm_enc);
 
 	spin_lock_irqsave(&dpu_enc->enc_spinlock, lock_flags);
-	if (dpu_enc->crtc_vblank_cb)
-		dpu_enc->crtc_vblank_cb(dpu_enc->crtc_vblank_cb_data);
+	if (dpu_enc->crtc)
+		dpu_crtc_vblank_callback(dpu_enc->crtc);
 	spin_unlock_irqrestore(&dpu_enc->enc_spinlock, lock_flags);
 
 	atomic_inc(&phy_enc->vsync_cnt);
@@ -1262,15 +1263,14 @@ static void dpu_encoder_underrun_callback(struct drm_encoder *drm_enc,
 	DPU_ATRACE_END("encoder_underrun_callback");
 }
 
-void dpu_encoder_register_vblank_callback(struct drm_encoder *drm_enc,
-		void (*vbl_cb)(void *), void *vbl_data)
+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 = vbl_cb ? true : false;
+	enable = crtc ? true : false;
 
 	if (!drm_enc) {
 		DPU_ERROR("invalid encoder\n");
@@ -1279,8 +1279,9 @@ void dpu_encoder_register_vblank_callback(struct drm_encoder *drm_enc,
 	trace_dpu_enc_vblank_cb(DRMID(drm_enc), enable);
 
 	spin_lock_irqsave(&dpu_enc->enc_spinlock, lock_flags);
-	dpu_enc->crtc_vblank_cb = vbl_cb;
-	dpu_enc->crtc_vblank_cb_data = vbl_data;
+	/* 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);
 
 	for (i = 0; i < dpu_enc->num_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 aa4f135218fa..be1d80867834 100644
--- a/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h
+++ b/drivers/gpu/drm/msm/disp/dpu1/dpu_encoder.h
@@ -55,14 +55,12 @@ void dpu_encoder_get_hw_resources(struct drm_encoder *encoder,
 				  struct dpu_encoder_hw_resources *hw_res);
 
 /**
- * dpu_encoder_register_vblank_callback - provide callback to encoder that
- *	will be called on the next vblank.
+ * dpu_encoder_assign_crtc - Link the encoder to the crtc it's assigned to
  * @encoder:	encoder pointer
- * @cb:		callback pointer, provide NULL to deregister and disable IRQs
- * @data:	user data provided to callback
+ * @crtc:	crtc pointer
  */
-void dpu_encoder_register_vblank_callback(struct drm_encoder *encoder,
-		void (*cb)(void *), void *data);
+void dpu_encoder_assign_crtc(struct drm_encoder *encoder,
+			     struct drm_crtc *crtc);
 
 /**
  * dpu_encoder_register_frame_event_callback - provide callback to encoder that
