diff --git a/drivers/gpu/drm/drm_atomic.c b/drivers/gpu/drm/drm_atomic.c
index 434915448ea0be99b7786c6daf2650cfb93b728a..f7d5166f89b24ef740e854175927ad934652fed4 100644
--- a/drivers/gpu/drm/drm_atomic.c
+++ b/drivers/gpu/drm/drm_atomic.c
@@ -1515,7 +1515,8 @@ retry:
 			copied_props++;
 		}
 
-		if (obj->type == DRM_MODE_OBJECT_PLANE && count_props) {
+		if (obj->type == DRM_MODE_OBJECT_PLANE && count_props &&
+		    !(arg->flags & DRM_MODE_ATOMIC_TEST_ONLY)) {
 			plane = obj_to_plane(obj);
 			plane_mask |= (1 << drm_plane_index(plane));
 			plane->old_fb = plane->fb;
@@ -1537,10 +1538,11 @@ retry:
 	}
 
 	if (arg->flags & DRM_MODE_ATOMIC_TEST_ONLY) {
+		/*
+		 * Unlike commit, check_only does not clean up state.
+		 * Below we call drm_atomic_state_free for it.
+		 */
 		ret = drm_atomic_check_only(state);
-		/* _check_only() does not free state, unlike _commit() */
-		if (!ret)
-			drm_atomic_state_free(state);
 	} else if (arg->flags & DRM_MODE_ATOMIC_NONBLOCK) {
 		ret = drm_atomic_async_commit(state);
 	} else {
@@ -1567,25 +1569,30 @@ out:
 		plane->old_fb = NULL;
 	}
 
+	if (ret && arg->flags & DRM_MODE_PAGE_FLIP_EVENT) {
+		/*
+		 * TEST_ONLY and PAGE_FLIP_EVENT are mutually exclusive,
+		 * if they weren't, this code should be called on success
+		 * for TEST_ONLY too.
+		 */
+
+		for_each_crtc_in_state(state, crtc, crtc_state, i) {
+			if (!crtc_state->event)
+				continue;
+
+			destroy_vblank_event(dev, file_priv,
+					     crtc_state->event);
+		}
+	}
+
 	if (ret == -EDEADLK) {
 		drm_atomic_state_clear(state);
 		drm_modeset_backoff(&ctx);
 		goto retry;
 	}
 
-	if (ret) {
-		if (arg->flags & DRM_MODE_PAGE_FLIP_EVENT) {
-			for_each_crtc_in_state(state, crtc, crtc_state, i) {
-				if (!crtc_state->event)
-					continue;
-
-				destroy_vblank_event(dev, file_priv,
-						     crtc_state->event);
-			}
-		}
-
+	if (ret || arg->flags & DRM_MODE_ATOMIC_TEST_ONLY)
 		drm_atomic_state_free(state);
-	}
 
 	drm_modeset_drop_locks(&ctx);
 	drm_modeset_acquire_fini(&ctx);
diff --git a/drivers/gpu/drm/drm_dp_helper.c b/drivers/gpu/drm/drm_dp_helper.c
index 80a02a412607f9c1b38e20b5494033b969c9e800..291734e87fca7457da9eb3ec8e0ba80f262bd232 100644
--- a/drivers/gpu/drm/drm_dp_helper.c
+++ b/drivers/gpu/drm/drm_dp_helper.c
@@ -159,6 +159,8 @@ int drm_dp_bw_code_to_link_rate(u8 link_bw)
 }
 EXPORT_SYMBOL(drm_dp_bw_code_to_link_rate);
 
+#define AUX_RETRY_INTERVAL 500 /* us */
+
 /**
  * DOC: dp helpers
  *
@@ -213,7 +215,7 @@ static int drm_dp_dpcd_access(struct drm_dp_aux *aux, u8 request,
 			return -EIO;
 
 		case DP_AUX_NATIVE_REPLY_DEFER:
-			usleep_range(400, 500);
+			usleep_range(AUX_RETRY_INTERVAL, AUX_RETRY_INTERVAL + 100);
 			break;
 		}
 	}
@@ -422,6 +424,90 @@ static u32 drm_dp_i2c_functionality(struct i2c_adapter *adapter)
 	       I2C_FUNC_10BIT_ADDR;
 }
 
+#define AUX_PRECHARGE_LEN 10 /* 10 to 16 */
+#define AUX_SYNC_LEN (16 + 4) /* preamble + AUX_SYNC_END */
+#define AUX_STOP_LEN 4
+#define AUX_CMD_LEN 4
+#define AUX_ADDRESS_LEN 20
+#define AUX_REPLY_PAD_LEN 4
+#define AUX_LENGTH_LEN 8
+
+/*
+ * Calculate the duration of the AUX request/reply in usec. Gives the
+ * "best" case estimate, ie. successful while as short as possible.
+ */
+static int drm_dp_aux_req_duration(const struct drm_dp_aux_msg *msg)
+{
+	int len = AUX_PRECHARGE_LEN + AUX_SYNC_LEN + AUX_STOP_LEN +
+		AUX_CMD_LEN + AUX_ADDRESS_LEN + AUX_LENGTH_LEN;
+
+	if ((msg->request & DP_AUX_I2C_READ) == 0)
+		len += msg->size * 8;
+
+	return len;
+}
+
+static int drm_dp_aux_reply_duration(const struct drm_dp_aux_msg *msg)
+{
+	int len = AUX_PRECHARGE_LEN + AUX_SYNC_LEN + AUX_STOP_LEN +
+		AUX_CMD_LEN + AUX_REPLY_PAD_LEN;
+
+	/*
+	 * For read we expect what was asked. For writes there will
+	 * be 0 or 1 data bytes. Assume 0 for the "best" case.
+	 */
+	if (msg->request & DP_AUX_I2C_READ)
+		len += msg->size * 8;
+
+	return len;
+}
+
+#define I2C_START_LEN 1
+#define I2C_STOP_LEN 1
+#define I2C_ADDR_LEN 9 /* ADDRESS + R/W + ACK/NACK */
+#define I2C_DATA_LEN 9 /* DATA + ACK/NACK */
+
+/*
+ * Calculate the length of the i2c transfer in usec, assuming
+ * the i2c bus speed is as specified. Gives the the "worst"
+ * case estimate, ie. successful while as long as possible.
+ * Doesn't account the the "MOT" bit, and instead assumes each
+ * message includes a START, ADDRESS and STOP. Neither does it
+ * account for additional random variables such as clock stretching.
+ */
+static int drm_dp_i2c_msg_duration(const struct drm_dp_aux_msg *msg,
+				   int i2c_speed_khz)
+{
+	/* AUX bitrate is 1MHz, i2c bitrate as specified */
+	return DIV_ROUND_UP((I2C_START_LEN + I2C_ADDR_LEN +
+			     msg->size * I2C_DATA_LEN +
+			     I2C_STOP_LEN) * 1000, i2c_speed_khz);
+}
+
+/*
+ * Deterine how many retries should be attempted to successfully transfer
+ * the specified message, based on the estimated durations of the
+ * i2c and AUX transfers.
+ */
+static int drm_dp_i2c_retry_count(const struct drm_dp_aux_msg *msg,
+			      int i2c_speed_khz)
+{
+	int aux_time_us = drm_dp_aux_req_duration(msg) +
+		drm_dp_aux_reply_duration(msg);
+	int i2c_time_us = drm_dp_i2c_msg_duration(msg, i2c_speed_khz);
+
+	return DIV_ROUND_UP(i2c_time_us, aux_time_us + AUX_RETRY_INTERVAL);
+}
+
+/*
+ * FIXME currently assumes 10 kHz as some real world devices seem
+ * to require it. We should query/set the speed via DPCD if supported.
+ */
+static int dp_aux_i2c_speed_khz __read_mostly = 10;
+module_param_unsafe(dp_aux_i2c_speed_khz, int, 0644);
+MODULE_PARM_DESC(dp_aux_i2c_speed_khz,
+		 "Assumed speed of the i2c bus in kHz, (1-400, default 10)");
+
 /*
  * Transfer a single I2C-over-AUX message and handle various error conditions,
  * retrying the transaction as appropriate.  It is assumed that the
@@ -434,13 +520,16 @@ static int drm_dp_i2c_do_msg(struct drm_dp_aux *aux, struct drm_dp_aux_msg *msg)
 {
 	unsigned int retry, defer_i2c;
 	int ret;
-
 	/*
 	 * DP1.2 sections 2.7.7.1.5.6.1 and 2.7.7.1.6.6.1: A DP Source device
 	 * is required to retry at least seven times upon receiving AUX_DEFER
 	 * before giving up the AUX transaction.
+	 *
+	 * We also try to account for the i2c bus speed.
 	 */
-	for (retry = 0, defer_i2c = 0; retry < (7 + defer_i2c); retry++) {
+	int max_retries = max(7, drm_dp_i2c_retry_count(msg, dp_aux_i2c_speed_khz));
+
+	for (retry = 0, defer_i2c = 0; retry < (max_retries + defer_i2c); retry++) {
 		mutex_lock(&aux->hw_mutex);
 		ret = aux->transfer(aux, msg);
 		mutex_unlock(&aux->hw_mutex);
@@ -476,7 +565,7 @@ static int drm_dp_i2c_do_msg(struct drm_dp_aux *aux, struct drm_dp_aux_msg *msg)
 			 * For now just defer for long enough to hopefully be
 			 * safe for all use-cases.
 			 */
-			usleep_range(500, 600);
+			usleep_range(AUX_RETRY_INTERVAL, AUX_RETRY_INTERVAL + 100);
 			continue;
 
 		default:
@@ -506,7 +595,7 @@ static int drm_dp_i2c_do_msg(struct drm_dp_aux *aux, struct drm_dp_aux_msg *msg)
 			aux->i2c_defer_count++;
 			if (defer_i2c < 7)
 				defer_i2c++;
-			usleep_range(400, 500);
+			usleep_range(AUX_RETRY_INTERVAL, AUX_RETRY_INTERVAL + 100);
 			continue;
 
 		default: