diff --git a/drivers/aml_tee/optee/log.c b/drivers/aml_tee/optee/log.c
index e0315da8a..54adcc609 100644
--- a/drivers/aml_tee/optee/log.c
+++ b/drivers/aml_tee/optee/log.c
@@ -16,9 +16,11 @@
 #include <linux/string.h>
 #include <linux/types.h>
 #include <linux/workqueue.h>
+#include <linux/io.h>
 
 #include "optee_smc.h"
 #include "log.h"
+#include "optee_private.h"
 
 #define LOGGER_LOOPBUFFER_MAGIC       0xAA00AA00
 #define LOGGER_LOOPBUFFER_OFFSET      0x00000080
@@ -136,21 +138,49 @@ static void do_log_timer(struct work_struct *work)
 		pr_err("%s:%d Failed to join the workqueue\n", __func__, __LINE__);
 }
 
-int optee_log_init(void *va, phys_addr_t pa, size_t size)
+int optee_log_init(void *va)
 {
 	int rc = 0;
-	struct arm_smccc_res smccc;
+	size_t size = 0;
+	phys_addr_t begin = 0;
+	phys_addr_t end = 0;
+	struct arm_smccc_res smccc = { 0 };
 	struct loopbuffer_ctl_s *log_ctl = NULL;
 
 	arm_smccc_smc(OPTEE_SMC_ENABLE_LOGGER, 1, 0, 0, 0, 0, 0, 0,
 			&smccc);
 
-	log_buf_va = va;
-	log_ctl = (struct loopbuffer_ctl_s *)va;
+	if (smccc.a0 != TEEC_SUCCESS) {
+		pr_err("smc enable logger failed, res = 0x%lx\n", smccc.a0);
+		rc = -EACCES;
+		goto err;
+	}
+
+	if (smccc.a1 == LOGGER_SHM_ADDR_MAGIC) {
+		/*
+		 * for new logger solution
+		 * a1: logger get flags
+		 * a2: logger start address
+		 * a3: logger size
+		 */
+		begin = roundup(smccc.a2, PAGE_SIZE);
+		end = rounddown(smccc.a2 + smccc.a3, PAGE_SIZE);
+		size = end - begin;
+		log_buf_va = memremap(begin, size, MEMREMAP_WB);
+	} else {
+		/* for previous logger solution */
+		log_buf_va = va;
+	}
+
+	if (!log_buf_va) {
+		pr_err("shared memory logger ioremap failed\n");
+		return -ENOMEM;
+	}
+
+	log_ctl = (struct loopbuffer_ctl_s *)log_buf_va;
 	if (log_ctl->magic != LOGGER_LOOPBUFFER_MAGIC || log_ctl->inited != 1) {
 		pr_err("tee log buffer init failed\n");
-		rc = -1;
-
+		rc = -EINVAL;
 		goto err;
 	}
 
@@ -159,8 +189,7 @@ int optee_log_init(void *va, phys_addr_t pa, size_t size)
 	INIT_DELAYED_WORK(&log_work, do_log_timer);
 	if (queue_delayed_work(log_workqueue, &log_work, OPTEE_LOG_TIMER_INTERVAL * HZ) == 0) {
 		pr_err("%s:%d failed to join the workqueue.\n", __func__, __LINE__);
-		rc = -1;
-
+		rc = -EBUSY;
 		goto err;
 	}
 
diff --git a/drivers/aml_tee/optee/log.h b/drivers/aml_tee/optee/log.h
index 8679e82eb..b3e36da91 100644
--- a/drivers/aml_tee/optee/log.h
+++ b/drivers/aml_tee/optee/log.h
@@ -2,9 +2,11 @@
 #ifndef _OPTEE_LOG_H_
 #define _OPTEE_LOG_H_
 
+#define LOGGER_SHM_ADDR_MAGIC           0x1001
+
 #define DEF_LOGGER_SHM_SIZE             (256 * 1024)
 
-int optee_log_init(void *va, phys_addr_t pa, size_t size);
+int optee_log_init(void *va);
 
 void optee_log_uninit(void);
 
diff --git a/drivers/aml_tee/optee/optee_smc.h b/drivers/aml_tee/optee/optee_smc.h
index 890c66bb7..8caffe9be 100644
--- a/drivers/aml_tee/optee/optee_smc.h
+++ b/drivers/aml_tee/optee/optee_smc.h
@@ -556,18 +556,18 @@ struct optee_smc_disable_shm_cache_result {
 	OPTEE_SMC_RPC_VAL(OPTEE_SMC_RPC_FUNC_CMD)
 
 /*
- * enable logger
+ * Enable logger
  *
  * Call register usage:
  * a0      SMC Function ID, OPTEE_SMC_ENABLE_LOGGER
- * a1      enable logger: a1 > 0; disable logger: a1 = 0;
- * a2      logger share-mem phy addr
- * a3      logger share-mem size
- * a4-7    Not used
+ * a1      Enable logger: a1 > 0; disable logger: a1 = 0;
+ * a2-a7   Not used
  *
  * Normal return register usage:
- * a0      enable logger result
- * a1-7    Preserved
+ * a0      Enable logger result
+ * a1      Logger get flags
+ * a2      Logger share-mem phy addr
+ * a3      Logger share-mem size
  */
 #define OPTEE_SMC_FUNCID_ENABLE_LOGGER                     0xE001
 #define OPTEE_SMC_ENABLE_LOGGER  OPTEE_SMC_FAST_CALL_VAL(OPTEE_SMC_FUNCID_ENABLE_LOGGER)
diff --git a/drivers/aml_tee/optee/smc_abi.c b/drivers/aml_tee/optee/smc_abi.c
index b9d5b2ea5..aa09b5a14 100644
--- a/drivers/aml_tee/optee/smc_abi.c
+++ b/drivers/aml_tee/optee/smc_abi.c
@@ -1223,7 +1223,7 @@ static bool optee_msg_exchange_capabilities(optee_invoke_fn *invoke_fn,
 }
 
 static struct tee_shm_pool *
-optee_config_shm_memremap(optee_invoke_fn *invoke_fn, void **memremaped_shm)
+optee_config_shm_memremap(optee_invoke_fn *invoke_fn, void **memremaped_shm, void **log_va)
 {
 	union {
 		struct arm_smccc_res smccc;
@@ -1267,9 +1267,7 @@ optee_config_shm_memremap(optee_invoke_fn *invoke_fn, void **memremaped_shm)
 	else
 		*memremaped_shm = va;
 
-	optee_log_init(va + size - DEF_LOGGER_SHM_SIZE,
-			end - DEF_LOGGER_SHM_SIZE,
-			DEF_LOGGER_SHM_SIZE);
+	*log_va = va + size - DEF_LOGGER_SHM_SIZE;
 
 	return rc;
 }
@@ -1370,6 +1368,7 @@ static int optee_probe(struct platform_device *pdev)
 	struct tee_shm_pool *pool = ERR_PTR(-EINVAL);
 	struct optee *optee = NULL;
 	void *memremaped_shm = NULL;
+	void *shm_log_va = NULL;
 	unsigned int rpc_param_count;
 	struct tee_device *teedev;
 	struct tee_context *ctx;
@@ -1445,12 +1444,14 @@ static int optee_probe(struct platform_device *pdev)
 		 */
 		arg_cache_flags = OPTEE_SHM_ARG_SHARED |
 				  OPTEE_SHM_ARG_ALLOC_PRIV;
-		pool = optee_config_shm_memremap(invoke_fn, &memremaped_shm);
+		pool = optee_config_shm_memremap(invoke_fn, &memremaped_shm, &shm_log_va);
 	}
 
 	if (IS_ERR(pool))
 		return PTR_ERR(pool);
 
+	optee_log_init(shm_log_va);
+
 	optee = kzalloc(sizeof(*optee), GFP_KERNEL);
 	if (!optee) {
 		rc = -ENOMEM;