diff --git a/drivers/tee/optee/core.c b/drivers/tee/optee/core.c
index b898c32edc..828ab9b00a 100644
--- a/drivers/tee/optee/core.c
+++ b/drivers/tee/optee/core.c
@@ -12,6 +12,7 @@
 #include <linux/arm-smccc.h>
 #include <linux/err.h>
 #include <linux/io.h>
+#include <mmc.h>
 
 #include "optee_smc.h"
 #include "optee_msg.h"
@@ -607,14 +608,46 @@ static optee_invoke_fn *get_invoke_func(struct udevice *dev)
 	return ERR_PTR(-EINVAL);
 }
 
+static struct mmc *get_rpmb_dev(struct udevice *dev)
+{
+	struct udevice *mmc_dev;
+	const fdt32_t *phandle_p;
+	u32 phandle;
+	int ret = 0;
+
+	debug("optee: looking for rpmb device in DT.\n");
+
+	phandle_p  = ofnode_get_property(dev_ofnode(dev),
+					 "rpmb-dev", NULL);
+	if (!phandle_p) {
+		debug("optee: missing \"rpmb-dev\" property\n");
+		return NULL;
+	}
+
+	phandle = fdt32_to_cpu(*phandle_p);
+
+	ret = uclass_get_device_by_phandle_id(UCLASS_MMC, phandle, &mmc_dev);
+	if (ret) {
+		printf("optee: invalid phandle value in \"rpmb-dev\".\n");
+		return NULL;
+	}
+
+	debug("optee: using phandle %d from \"rpmd-dev\" property.\n",
+	      phandle);
+	return mmc_get_mmc_dev(mmc_dev);
+}
+
 static int optee_of_to_plat(struct udevice *dev)
 {
 	struct optee_pdata *pdata = dev_get_plat(dev);
+	struct optee_private *priv = dev_get_priv(dev);
 
 	pdata->invoke_fn = get_invoke_func(dev);
 	if (IS_ERR(pdata->invoke_fn))
 		return PTR_ERR(pdata->invoke_fn);
 
+	priv->rpmb_mmc = get_rpmb_dev(dev);
+
 	return 0;
 }
 
diff --git a/drivers/tee/optee/optee_private.h b/drivers/tee/optee/optee_private.h
index 1f07a27ee4..8e5a280622 100644
--- a/drivers/tee/optee/optee_private.h
+++ b/drivers/tee/optee/optee_private.h
@@ -19,8 +19,8 @@
  */
 struct optee_private {
 	struct mmc *rpmb_mmc;
-	int rpmb_dev_id;
 	int rpmb_original_part;
+	bool rpmb_inited;
 };
 
 struct optee_msg_arg;
diff --git a/drivers/tee/optee/rpmb.c b/drivers/tee/optee/rpmb.c
index 0804fc963c..0137c44be1 100644
--- a/drivers/tee/optee/rpmb.c
+++ b/drivers/tee/optee/rpmb.c
@@ -45,55 +45,67 @@ static void release_mmc(struct optee_private *priv)
 {
 	int rc;
 
-	if (!priv->rpmb_mmc)
+	if (!priv->rpmb_mmc || !priv->rpmb_inited)
 		return;
 
-	rc = blk_select_hwpart_devnum(IF_TYPE_MMC, priv->rpmb_dev_id,
-				      priv->rpmb_original_part);
+	rc = mmc_switch_part(priv->rpmb_mmc, priv->rpmb_original_part);
 	if (rc)
 		debug("%s: blk_select_hwpart_devnum() failed: %d\n",
 		      __func__, rc);
 
-	priv->rpmb_mmc = NULL;
+	priv->rpmb_inited = false;
+}
+
+static int check_mmc(struct mmc *mmc)
+{
+	if (!mmc) {
+		debug("Cannot find RPMB device\n");
+		return -ENODEV;
+	}
+	if (!(mmc->version & MMC_VERSION_MMC)) {
+		debug("Device id is not an eMMC device\n");
+		return -ENODEV;
+	}
+	if (mmc->version < MMC_VERSION_4_41) {
+		debug("RPMB is not supported before version 4.41\n");
+		return -ENODEV;
+	}
+
+	return 0;
 }
 
 static struct mmc *get_mmc(struct optee_private *priv, int dev_id)
 {
-	struct mmc *mmc;
 	int rc;
 
-	if (priv->rpmb_mmc && priv->rpmb_dev_id == dev_id)
+	if (priv->rpmb_mmc && priv->rpmb_inited)
 		return priv->rpmb_mmc;
 
 	release_mmc(priv);
 
-	mmc = find_mmc_device(dev_id);
-	if (!mmc) {
-		debug("Cannot find RPMB device\n");
-		return NULL;
-	}
-	if (!(mmc->version & MMC_VERSION_MMC)) {
-		debug("Device id %d is not an eMMC device\n", dev_id);
-		return NULL;
-	}
-	if (mmc->version < MMC_VERSION_4_41) {
-		debug("Device id %d: RPMB not supported before version 4.41\n",
-		      dev_id);
+	/*
+	 * Check if priv->rpmb_mmc was already set from DT node,
+	 * otherwise use dev_id provided by OP-TEE OS
+	 * and find mmc device by its dev_id
+	 */
+	if (!priv->rpmb_mmc)
+		priv->rpmb_mmc = find_mmc_device(dev_id);
+
+	rc = check_mmc(priv->rpmb_mmc);
+	if (rc)
 		return NULL;
-	}
 
-	priv->rpmb_original_part = mmc_get_blk_desc(mmc)->hwpart;
+	priv->rpmb_original_part = mmc_get_blk_desc(priv->rpmb_mmc)->hwpart;
 
-	rc = blk_select_hwpart_devnum(IF_TYPE_MMC, dev_id, MMC_PART_RPMB);
+	rc = mmc_switch_part(priv->rpmb_mmc, MMC_PART_RPMB);
 	if (rc) {
 		debug("Device id %d: cannot select RPMB partition: %d\n",
 		      dev_id, rc);
 		return NULL;
 	}
 
-	priv->rpmb_mmc = mmc;
-	priv->rpmb_dev_id = dev_id;
-	return mmc;
+	priv->rpmb_inited = true;
+	return priv->rpmb_mmc;
 }
 
 static u32 rpmb_get_dev_info(u16 dev_id, struct rpmb_dev_info *info)
