Tegra210B01: SE/SE2 and PKA1 context save (SW)

This change ports the software based SE context save routines.
The software implements the context save sequence for SE/SE2 and
PKA1. The context save routine is intended to be invoked from
the ATF SC7 entry.

Change-Id: I9aa156d6e7e22a394bb10cb0c3b05fc303f08807
Signed-off-by: Marvin Hsu <marvinh@nvidia.com>

2 files changed, 1002 insertions, 56 deletions

diff --git a/plat/nvidia/tegra/soc/t210/drivers/se/se_private.h b/plat/nvidia/tegra/soc/t210/drivers/se/se_private.h
index f76fa11d..be1f9cc3 100644
--- a/plat/nvidia/tegra/soc/t210/drivers/se/se_private.h
+++ b/plat/nvidia/tegra/soc/t210/drivers/se/se_private.h
@@ -16,14 +16,16 @@
  */
 
 /* Secure scratch registers */
-#define PMC_SECURE_SCRATCH4_OFFSET      	0xC0U
-#define PMC_SECURE_SCRATCH5_OFFSET      	0xC4U
-#define PMC_SECURE_SCRATCH6_OFFSET      	0x224U
-#define PMC_SECURE_SCRATCH7_OFFSET      	0x228U
-#define PMC_SECURE_SCRATCH120_OFFSET    	0xB38U
-#define PMC_SECURE_SCRATCH121_OFFSET    	0xB3CU
-#define PMC_SECURE_SCRATCH122_OFFSET    	0xB40U
-#define PMC_SECURE_SCRATCH123_OFFSET    	0xB44U
+#define PMC_SECURE_SCRATCH4_OFFSET		0xC0U
+#define PMC_SECURE_SCRATCH5_OFFSET		0xC4U
+#define PMC_SECURE_SCRATCH6_OFFSET		0x224U
+#define PMC_SECURE_SCRATCH7_OFFSET		0x228U
+#define PMC_SECURE_SCRATCH116_OFFSET		0xB28U
+#define PMC_SECURE_SCRATCH117_OFFSET		0xB2CU
+#define PMC_SECURE_SCRATCH120_OFFSET		0xB38U
+#define PMC_SECURE_SCRATCH121_OFFSET		0xB3CU
+#define PMC_SECURE_SCRATCH122_OFFSET		0xB40U
+#define PMC_SECURE_SCRATCH123_OFFSET		0xB44U
 
 /*
  * AHB arbitration memory write queue
@@ -32,6 +34,12 @@
 #define ARAHB_MST_ID_SE2_MASK			(0x1U << 13)
 #define ARAHB_MST_ID_SE_MASK			(0x1U << 14)
 
+/**
+ * SE registers
+ */
+#define TEGRA_SE_AES_KEYSLOT_COUNT		16
+#define SE_MAX_LAST_BLOCK_SIZE			0xFFFFF
+
 /* SE Status register */
 #define SE_STATUS_OFFSET			0x800U
 #define SE_STATUS_SHIFT				0
@@ -42,8 +50,24 @@
 #define SE_STATUS(x)	\
 		((x) & ((0x3U) << SE_STATUS_SHIFT))
 
+#define SE_MEM_INTERFACE_SHIFT			2
+#define SE_MEM_INTERFACE_IDLE			0
+#define SE_MEM_INTERFACE_BUSY			1
+#define SE_MEM_INTERFACE(x)	((x) << SE_STATUS_SHIFT)
+
+/* SE register definitions */
+#define SE_SECURITY_REG_OFFSET			0x0
+#define SE_SECURITY_TZ_LOCK_SOFT_SHIFT		5
+#define SE_SECURE				0x0
+#define SE_SECURITY_TZ_LOCK_SOFT(x)	((x) << SE_SECURITY_TZ_LOCK_SOFT_SHIFT)
+
+#define SE_SEC_ENG_DIS_SHIFT			1
+#define SE_DISABLE_FALSE			0
+#define SE_DISABLE_TRUE				1
+#define SE_SEC_ENG_DISABLE(x)((x) << SE_SEC_ENG_DIS_SHIFT)
+
 /* SE config register */
-#define SE_CONFIG_REG_OFFSET    		0x14U
+#define SE_CONFIG_REG_OFFSET			0x14U
 #define SE_CONFIG_ENC_ALG_SHIFT 		12
 #define SE_CONFIG_ENC_ALG_AES_ENC	\
 		((1U) << SE_CONFIG_ENC_ALG_SHIFT)
@@ -66,7 +90,7 @@
 #define SE_CONFIG_DEC_ALG(x)	\
 		((x) & ((0xFU) << SE_CONFIG_DEC_ALG_SHIFT))
 
-#define SE_CONFIG_DST_SHIFT     		2
+#define SE_CONFIG_DST_SHIFT	 		2
 #define SE_CONFIG_DST_MEMORY	\
 		((0U) << SE_CONFIG_DST_SHIFT)
 #define SE_CONFIG_DST_HASHREG	\
@@ -80,6 +104,47 @@
 #define SE_CONFIG_DST(x)	\
 		((x) & ((0x7U) << SE_CONFIG_DST_SHIFT))
 
+#define SE_CONFIG_ENC_MODE_SHIFT		24
+#define SE_CONFIG_ENC_MODE_KEY128	\
+			((0UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_KEY192	\
+			((1UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_KEY256	\
+			((2UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_SHA1				\
+			((0UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_SHA224	\
+			((4UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_SHA256	\
+			((5UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_SHA384	\
+			((6UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE_SHA512	\
+			((7UL) << SE_CONFIG_ENC_MODE_SHIFT)
+#define SE_CONFIG_ENC_MODE(x)\
+			((x) & ((0xFFUL) << SE_CONFIG_ENC_MODE_SHIFT))
+
+#define SE_CONFIG_DEC_MODE_SHIFT		16
+#define SE_CONFIG_DEC_MODE_KEY128	\
+			((0UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_KEY192	\
+			((1UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_KEY256	\
+			((2UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_SHA1				\
+			((0UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_SHA224	\
+			((4UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_SHA256	\
+			((5UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_SHA384	\
+			((6UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE_SHA512	\
+			((7UL) << SE_CONFIG_DEC_MODE_SHIFT)
+#define SE_CONFIG_DEC_MODE(x)\
+			((x) & ((0xFFUL) << SE_CONFIG_DEC_MODE_SHIFT))
+
+
 /* DRBG random number generator config */
 #define SE_RNG_CONFIG_REG_OFFSET		0x340
 
@@ -104,9 +169,10 @@
 		((x) & ((0x3U) << DRBG_SRC_SHIFT))
 
 /* DRBG random number generator entropy config */
+
 #define SE_RNG_SRC_CONFIG_REG_OFFSET		0x344U
 
-#define DRBG_RO_ENT_SRC_SHIFT       		1
+#define DRBG_RO_ENT_SRC_SHIFT			1
 #define DRBG_RO_ENT_SRC_ENABLE	\
 		((1U) << DRBG_RO_ENT_SRC_SHIFT)
 #define DRBG_RO_ENT_SRC_DISABLE	\
@@ -114,7 +180,7 @@
 #define SE_RNG_SRC_CONFIG_RO_ENT_SRC(x)	\
 		((x) & ((0x1U) << DRBG_RO_ENT_SRC_SHIFT))
 
-#define DRBG_RO_ENT_SRC_LOCK_SHIFT  		0
+#define DRBG_RO_ENT_SRC_LOCK_SHIFT		0
 #define DRBG_RO_ENT_SRC_LOCK_ENABLE	\
 		((1U) << DRBG_RO_ENT_SRC_LOCK_SHIFT)
 #define DRBG_RO_ENT_SRC_LOCK_DISABLE	\
@@ -130,9 +196,97 @@
 #define SE_RNG_SRC_CONFIG_RO_ENT_IGNORE_MEM(x)	\
 		((x) & ((0x1U) << DRBG_RO_ENT_IGNORE_MEM_SHIFT))
 
+#define SE_RNG_RESEED_INTERVAL_REG_OFFSET	0x348
+
+/* SE CRYPTO */
+#define SE_CRYPTO_REG_OFFSET			0x304
+#define SE_CRYPTO_HASH_SHIFT			0
+#define SE_CRYPTO_HASH_DISABLE	\
+		((0U) << SE_CRYPTO_HASH_SHIFT)
+#define SE_CRYPTO_HASH_ENABLE	\
+		((1U) << SE_CRYPTO_HASH_SHIFT)
+
+#define SE_CRYPTO_XOR_POS_SHIFT			1
+#define SE_CRYPTO_XOR_BYPASS	\
+		((0U) << SE_CRYPTO_XOR_POS_SHIFT)
+#define SE_CRYPTO_XOR_TOP	\
+		((2U) << SE_CRYPTO_XOR_POS_SHIFT)
+#define SE_CRYPTO_XOR_BOTTOM	\
+		((3U) << SE_CRYPTO_XOR_POS_SHIFT)
+
+#define SE_CRYPTO_INPUT_SEL_SHIFT		3
+#define SE_CRYPTO_INPUT_AHB 	\
+		((0U) << SE_CRYPTO_INPUT_SEL_SHIFT)
+#define SE_CRYPTO_INPUT_RANDOM	\
+		((1U) << SE_CRYPTO_INPUT_SEL_SHIFT)
+#define SE_CRYPTO_INPUT_AESOUT	\
+		((2U) << SE_CRYPTO_INPUT_SEL_SHIFT)
+#define SE_CRYPTO_INPUT_LNR_CTR \
+		((3U) << SE_CRYPTO_INPUT_SEL_SHIFT)
+
+#define SE_CRYPTO_VCTRAM_SEL_SHIFT		 5
+#define SE_CRYPTO_VCTRAM_AHB	\
+		((0U) << SE_CRYPTO_VCTRAM_SEL_SHIFT)
+#define SE_CRYPTO_VCTRAM_AESOUT \
+		((2U) << SE_CRYPTO_VCTRAM_SEL_SHIFT)
+#define SE_CRYPTO_VCTRAM_PREVAHB	\
+		((3U) << SE_CRYPTO_VCTRAM_SEL_SHIFT)
+
+#define SE_CRYPTO_IV_SEL_SHIFT			 7
+#define SE_CRYPTO_IV_ORIGINAL	\
+		((0U) << SE_CRYPTO_IV_SEL_SHIFT)
+#define SE_CRYPTO_IV_UPDATED	\
+		((1U) << SE_CRYPTO_IV_SEL_SHIFT)
+
+#define SE_CRYPTO_CORE_SEL_SHIFT		8
+#define SE_CRYPTO_CORE_DECRYPT	\
+		((0U) << SE_CRYPTO_CORE_SEL_SHIFT)
+#define SE_CRYPTO_CORE_ENCRYPT	\
+		((1U) << SE_CRYPTO_CORE_SEL_SHIFT)
+
+#define SE_CRYPTO_KEY_INDEX_SHIFT		24
+#define SE_CRYPTO_KEY_INDEX(x) (x << SE_CRYPTO_KEY_INDEX_SHIFT)
+
+#define SE_CRYPTO_MEMIF_AHB	\
+		((0U) << SE_CRYPTO_MEMIF_SHIFT)
+#define SE_CRYPTO_MEMIF_MCCIF	\
+		((1U) << SE_CRYPTO_MEMIF_SHIFT)
+#define SE_CRYPTO_MEMIF_SHIFT			31
+
+/* KEY TABLE */
+#define SE_KEYTABLE_REG_OFFSET			0x31C
+
+/* KEYIV PKT - key slot */
+#define SE_KEYTABLE_SLOT_SHIFT			4
+#define SE_KEYTABLE_SLOT(x)	(x << SE_KEYTABLE_SLOT_SHIFT)
+
+/* KEYIV PKT - KEYIV select */
+#define SE_KEYIV_PKT_KEYIV_SEL_SHIFT		3
+#define SE_CRYPTO_KEYIV_KEY	\
+		((0U) << SE_KEYIV_PKT_KEYIV_SEL_SHIFT)
+#define SE_CRYPTO_KEYIV_IVS	\
+		((1U) << SE_KEYIV_PKT_KEYIV_SEL_SHIFT)
+
+/* KEYIV PKT - IV select */
+#define SE_KEYIV_PKT_IV_SEL_SHIFT		2
+#define SE_CRYPTO_KEYIV_IVS_OIV	\
+		((0U) << SE_KEYIV_PKT_IV_SEL_SHIFT)
+#define SE_CRYPTO_KEYIV_IVS_UIV \
+		((1U) << SE_KEYIV_PKT_IV_SEL_SHIFT)
+
+/* KEYIV PKT - key word */
+#define SE_KEYIV_PKT_KEY_WORD_SHIFT		0
+#define SE_KEYIV_PKT_KEY_WORD(x)	\
+		((x) << SE_KEYIV_PKT_KEY_WORD_SHIFT)
+
+/* KEYIV PKT - iv word */
+#define SE_KEYIV_PKT_IV_WORD_SHIFT		0
+#define SE_KEYIV_PKT_IV_WORD(x)		\
+		((x) << SE_KEYIV_PKT_IV_WORD_SHIFT)
+
 /* SE OPERATION */
 #define SE_OPERATION_REG_OFFSET 		0x8U
-#define SE_OPERATION_SHIFT      		0
+#define SE_OPERATION_SHIFT			0
 #define SE_OP_ABORT	\
 		((0x0U) << SE_OPERATION_SHIFT)
 #define SE_OP_START	\
@@ -146,11 +300,85 @@
 #define SE_OPERATION(x)	\
 		((x) & ((0x7U) << SE_OPERATION_SHIFT))
 
+/* SE CONTEXT */
+#define SE_CTX_SAVE_CONFIG_REG_OFFSET		0x70
+#define SE_CTX_SAVE_WORD_QUAD_SHIFT		0
+#define SE_CTX_SAVE_WORD_QUAD(x)	\
+		(x << SE_CTX_SAVE_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_WORD_QUAD_KEYS_0_3	\
+		((0U) << SE_CTX_SAVE_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_WORD_QUAD_KEYS_4_7	\
+		((1U) << SE_CTX_SAVE_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_WORD_QUAD_ORIG_IV	\
+		((2U) << SE_CTX_SAVE_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_WORD_QUAD_UPD_IV	\
+		((3U) << SE_CTX_SAVE_WORD_QUAD_SHIFT)
+
+#define SE_CTX_SAVE_KEY_INDEX_SHIFT		8
+#define SE_CTX_SAVE_KEY_INDEX(x)	(x << SE_CTX_SAVE_KEY_INDEX_SHIFT)
+
+#define SE_CTX_SAVE_STICKY_WORD_QUAD_SHIFT	24
+#define SE_CTX_SAVE_STICKY_WORD_QUAD_STICKY_0_3	\
+		((0U) << SE_CTX_SAVE_STICKY_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_STICKY_WORD_QUAD_STICKY_4_7	\
+		((1U) << SE_CTX_SAVE_STICKY_WORD_QUAD_SHIFT)
+#define SE_CTX_SAVE_STICKY_WORD_QUAD(x)	\
+		(x << SE_CTX_SAVE_STICKY_WORD_QUAD_SHIFT)
+
+#define SE_CTX_SAVE_SRC_SHIFT			29
+#define SE_CTX_SAVE_SRC_STICKY_BITS	\
+		((0U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_RSA_KEYTABLE	\
+		((1U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_AES_KEYTABLE	\
+		((2U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_PKA1_STICKY_BITS	\
+		((3U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_MEM	\
+		((4U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_SRK	\
+		((6U) << SE_CTX_SAVE_SRC_SHIFT)
+#define SE_CTX_SAVE_SRC_PKA1_KEYTABLE	\
+		((7U) << SE_CTX_SAVE_SRC_SHIFT)
+
+#define SE_CTX_STICKY_WORD_QUAD_SHIFT		24
+#define SE_CTX_STICKY_WORD_QUAD_WORDS_0_3 \
+		((0U) << SE_CTX_STICKY_WORD_QUAD_SHIFT)
+#define SE_CTX_STICKY_WORD_QUAD_WORDS_4_7 \
+		((1U) << SE_CTX_STICKY_WORD_QUAD_SHIFT)
+#define SE_CTX_STICKY_WORD_QUAD(x)	 (x << SE_CTX_STICKY_WORD_QUAD_SHIFT)
+
+#define SE_CTX_SAVE_RSA_KEY_INDEX_SHIFT		16
+#define SE_CTX_SAVE_RSA_KEY_INDEX(x)				\
+			(x << SE_CTX_SAVE_RSA_KEY_INDEX_SHIFT)
+
+#define SE_CTX_RSA_WORD_QUAD_SHIFT		12
+#define SE_CTX_RSA_WORD_QUAD(x)	\
+			(x << SE_CTX_RSA_WORD_QUAD_SHIFT)
+
+#define SE_CTX_PKA1_WORD_QUAD_L_SHIFT		0
+#define SE_CTX_PKA1_WORD_QUAD_L_SIZE		\
+			((true ? 4:0) - \
+			(false ? 4:0) + 1)
+#define SE_CTX_PKA1_WORD_QUAD_L(x)\
+			(((x) << SE_CTX_PKA1_WORD_QUAD_L_SHIFT) & 0x1f)
+
+#define SE_CTX_PKA1_WORD_QUAD_H_SHIFT		12
+#define SE_CTX_PKA1_WORD_QUAD_H(x)\
+			((((x) >> SE_CTX_PKA1_WORD_QUAD_L_SIZE) & 0xf) \
+			<< SE_CTX_PKA1_WORD_QUAD_H_SHIFT)
+
+#define SE_RSA_KEY_INDEX_SLOT0_EXP		0
+#define SE_RSA_KEY_INDEX_SLOT0_MOD		1
+#define SE_RSA_KEY_INDEX_SLOT1_EXP		2
+#define SE_RSA_KEY_INDEX_SLOT1_MOD		3
+
+
 /* SE_CTX_SAVE_AUTO */
 #define SE_CTX_SAVE_AUTO_REG_OFFSET 		0x74U
 
 /* Enable */
-#define SE_CTX_SAVE_AUTO_ENABLE_SHIFT  		0
+#define SE_CTX_SAVE_AUTO_ENABLE_SHIFT		0
 #define SE_CTX_SAVE_AUTO_DIS	\
 		((0U) << SE_CTX_SAVE_AUTO_ENABLE_SHIFT)
 #define SE_CTX_SAVE_AUTO_EN	\
@@ -167,20 +395,22 @@
 #define SE_CTX_SAVE_AUTO_LOCK(x)	\
 		((x) & ((0x1U) << SE_CTX_SAVE_AUTO_LOCK_SHIFT))
 
-/* Current context save number of blocks  */
+/* Current context save number of blocks*/
 #define SE_CTX_SAVE_AUTO_CURR_CNT_SHIFT		16
 #define SE_CTX_SAVE_AUTO_CURR_CNT_MASK 		0x3FFU
 #define SE_CTX_SAVE_GET_BLK_COUNT(x)	\
 		(((x) >> SE_CTX_SAVE_AUTO_CURR_CNT_SHIFT) & \
 		SE_CTX_SAVE_AUTO_CURR_CNT_MASK)
 
-#define SE_CTX_SAVE_SIZE_BLOCKS_SE1      	133
-#define SE_CTX_SAVE_SIZE_BLOCKS_SE2     	646
+#define SE_CTX_SAVE_SIZE_BLOCKS_SE1		133
+#define SE_CTX_SAVE_SIZE_BLOCKS_SE2	 	646
 
 /* SE TZRAM OPERATION - only for SE1 */
-#define SE_TZRAM_OPERATION      		0x540U
+#define SE_TZRAM_OPERATION			0x540U
 
-#define SE_TZRAM_OP_MODE_SHIFT  		1
+#define SE_TZRAM_OP_MODE_SHIFT			1
+#define SE_TZRAM_OP_COMMAND_INIT		1
+#define SE_TZRAM_OP_COMMAND_SHIFT		0
 #define SE_TZRAM_OP_MODE_SAVE		\
 		((0U) << SE_TZRAM_OP_MODE_SHIFT)
 #define SE_TZRAM_OP_MODE_RESTORE	\
@@ -188,7 +418,7 @@
 #define SE_TZRAM_OP_MODE(x)		\
 		((x) & ((0x1U) << SE_TZRAM_OP_MODE_SHIFT))
 
-#define SE_TZRAM_OP_BUSY_SHIFT  		2
+#define SE_TZRAM_OP_BUSY_SHIFT			2
 #define SE_TZRAM_OP_BUSY_OFF	\
 		((0U) << SE_TZRAM_OP_BUSY_SHIFT)
 #define SE_TZRAM_OP_BUSY_ON	\
@@ -196,7 +426,7 @@
 #define SE_TZRAM_OP_BUSY(x)	\
 		((x) & ((0x1U) << SE_TZRAM_OP_BUSY_SHIFT))
 
-#define SE_TZRAM_OP_REQ_SHIFT  			0
+#define SE_TZRAM_OP_REQ_SHIFT			0
 #define SE_TZRAM_OP_REQ_IDLE	\
 		((0U) << SE_TZRAM_OP_REQ_SHIFT)
 #define SE_TZRAM_OP_REQ_INIT	\
@@ -206,7 +436,7 @@
 
 /* SE Interrupt */
 #define SE_INT_STATUS_REG_OFFSET		0x10U
-#define SE_INT_OP_DONE_SHIFT    		4
+#define SE_INT_OP_DONE_SHIFT			4
 #define SE_INT_OP_DONE_CLEAR	\
 		((0U) << SE_INT_OP_DONE_SHIFT)
 #define SE_INT_OP_DONE_ACTIVE	\
@@ -214,19 +444,186 @@
 #define SE_INT_OP_DONE(x)	\
 		((x) & ((0x1U) << SE_INT_OP_DONE_SHIFT))
 
+/* SE TZRAM SECURITY */
+#define SE_TZRAM_SEC_REG_OFFSET			0x4
+
+#define SE_TZRAM_SEC_SETTING_SHIFT		 0
+#define SE_TZRAM_SECURE		\
+		((0UL) << SE_TZRAM_SEC_SETTING_SHIFT)
+#define SE_TZRAM_NONSECURE	 \
+		((1UL) << SE_TZRAM_SEC_SETTING_SHIFT)
+#define SE_TZRAM_SEC_SETTING(x)		\
+		((x) & ((0x1UL) << SE_TZRAM_SEC_SETTING_SHIFT))
+
+/* PKA1 KEY SLOTS */
+#define TEGRA_SE_PKA1_KEYSLOT_COUNT		4
+
+
 /* SE error status */
 #define SE_ERR_STATUS_REG_OFFSET		0x804U
+#define SE_CRYPTO_KEYTABLE_DST_REG_OFFSET	0x330
+#define SE_CRYPTO_KEYTABLE_DST_WORD_QUAD_SHIFT	0
+#define SE_CRYPTO_KEYTABLE_DST_WORD_QUAD(x)	\
+			(x << SE_CRYPTO_KEYTABLE_DST_WORD_QUAD_SHIFT)
+
+#define SE_KEY_INDEX_SHIFT			8
+#define SE_CRYPTO_KEYTABLE_DST_KEY_INDEX(x)	(x << SE_KEY_INDEX_SHIFT)
+
 
 /* SE linked list (LL) register */
 #define SE_IN_LL_ADDR_REG_OFFSET		0x18U
-#define SE_OUT_LL_ADDR_REG_OFFSET  		0x24U
-#define SE_BLOCK_COUNT_REG_OFFSET  		0x318U
+#define SE_OUT_LL_ADDR_REG_OFFSET		0x24U
+#define SE_BLOCK_COUNT_REG_OFFSET		0x318U
 
 /* AES data sizes */
+#define TEGRA_SE_KEY_256_SIZE			32
+#define TEGRA_SE_KEY_192_SIZE			24
+#define TEGRA_SE_KEY_128_SIZE			16
 #define TEGRA_SE_AES_BLOCK_SIZE 		16
-#define TEGRA_SE_AES_MIN_KEY_SIZE  		16
-#define TEGRA_SE_AES_MAX_KEY_SIZE  		32
-#define TEGRA_SE_AES_IV_SIZE    		16
+#define TEGRA_SE_AES_MIN_KEY_SIZE		16
+#define TEGRA_SE_AES_MAX_KEY_SIZE		32
+#define TEGRA_SE_AES_IV_SIZE			16
+
+#define TEGRA_SE_RNG_IV_SIZE			16
+#define TEGRA_SE_RNG_DT_SIZE			16
+#define TEGRA_SE_RNG_KEY_SIZE			16
+#define TEGRA_SE_RNG_SEED_SIZE	(TEGRA_SE_RNG_IV_SIZE + \
+									TEGRA_SE_RNG_KEY_SIZE + \
+									TEGRA_SE_RNG_DT_SIZE)
+#define TEGRA_SE_RSA512_DIGEST_SIZE		64
+#define TEGRA_SE_RSA1024_DIGEST_SIZE		128
+#define TEGRA_SE_RSA1536_DIGEST_SIZE		192
+#define TEGRA_SE_RSA2048_DIGEST_SIZE		256
+
+#define SE_KEY_TABLE_ACCESS_REG_OFFSET		0x284
+#define SE_KEY_READ_DISABLE_SHIFT		0
+
+#define SE_CTX_BUFER_SIZE			1072
+#define SE_CTX_DRBG_BUFER_SIZE			2112
+
+/* SE blobs size in bytes */
+#define SE_CTX_SAVE_RSA_KEY_LENGTH		1024
+#define SE_CTX_SAVE_RANDOM_DATA_SIZE		16
+#define SE_CTX_SAVE_STICKY_BITS_SIZE		16
+#define SE2_CONTEXT_SAVE_PKA1_STICKY_BITS_LENGTH 16
+#define SE2_CONTEXT_SAVE_PKA1_KEYS_LENGTH	8192
+#define SE_CTX_KNOWN_PATTERN_SIZE		16
+#define SE_CTX_KNOWN_PATTERN_SIZE_WORDS		(SE_CTX_KNOWN_PATTERN_SIZE/4)
+
+/* SE RSA */
+#define TEGRA_SE_RSA_KEYSLOT_COUNT		2
+#define SE_RSA_KEY_SIZE_REG_OFFSET		0x404
+#define SE_RSA_EXP_SIZE_REG_OFFSET		0x408
+#define SE_RSA_MAX_EXP_BIT_SIZE			2048
+#define SE_RSA_MAX_EXP_SIZE32	\
+		(SE_RSA_MAX_EXP_BIT_SIZE >> 5)
+#define SE_RSA_MAX_MOD_BIT_SIZE			2048
+#define SE_RSA_MAX_MOD_SIZE32	\
+		(SE_RSA_MAX_MOD_BIT_SIZE >> 5)
+
+/* SE_RSA_KEYTABLE_ADDR */
+#define SE_RSA_KEYTABLE_ADDR			0x420
+#define RSA_KEY_PKT_WORD_ADDR_SHIFT		0
+#define RSA_KEY_PKT_EXPMOD_SEL_SHIFT	\
+		((6U) << RSA_KEY_PKT_WORD_ADDR_SHIFT)
+#define RSA_KEY_MOD	\
+		((1U) << RSA_KEY_PKT_EXPMOD_SEL_SHIFT)
+#define RSA_KEY_EXP	\
+		((0U) << RSA_KEY_PKT_EXPMOD_SEL_SHIFT)
+#define RSA_KEY_PKT_SLOT_SHIFT			7
+#define RSA_KEY_SLOT_1	\
+		((0U) << RSA_KEY_PKT_SLOT_SHIFT)
+#define RSA_KEY_SLOT_2	\
+		((1U) << RSA_KEY_PKT_SLOT_SHIFT)
+#define RSA_KEY_PKT_INPUT_MODE_SHIFT		8
+#define RSA_KEY_REG_INPUT	\
+		((0U) << RSA_KEY_PKT_INPUT_MODE_SHIFT)
+#define RSA_KEY_DMA_INPUT	\
+		((1U) << RSA_KEY_PKT_INPUT_MODE_SHIFT)
+
+/* SE_RSA_KEYTABLE_DATA */
+#define SE_RSA_KEYTABLE_DATA			0x424
+
+/* SE_RSA_CONFIG register */
+#define SE_RSA_CONFIG				0x400
+#define RSA_KEY_SLOT_SHIFT			24
+#define RSA_KEY_SLOT(x) \
+		((x) << RSA_KEY_SLOT_SHIFT)
+
+/*******************************************************************************
+ * Structure definition
+ ******************************************************************************/
+
+/* SE context blob */
+#pragma pack(push, 1)
+typedef struct tegra_aes_key_slot {
+	/* 0 - 7 AES key */
+	uint32_t key[8];
+	/* 8 - 11 Original IV */
+	uint32_t oiv[4];
+	/* 12 - 15 Updated IV */
+	uint32_t uiv[4];
+} tegra_se_aes_key_slot_t;
+#pragma pack(pop)
+
+#pragma pack(push, 1)
+typedef struct tegra_se_context {
+	/* random number */
+	unsigned char rand_data[SE_CTX_SAVE_RANDOM_DATA_SIZE];
+	/* Sticky bits */
+	unsigned char sticky_bits[SE_CTX_SAVE_STICKY_BITS_SIZE * 2];
+	/* AES key slots */
+	tegra_se_aes_key_slot_t key_slots[TEGRA_SE_AES_KEYSLOT_COUNT];
+	/* RSA key slots */
+	unsigned char rsa_keys[SE_CTX_SAVE_RSA_KEY_LENGTH];
+} tegra_se_context_t;
+#pragma pack(pop)
+
+/* PKA context blob */
+#pragma pack(push, 1)
+typedef struct tegra_pka_context {
+	unsigned char sticky_bits[SE2_CONTEXT_SAVE_PKA1_STICKY_BITS_LENGTH];
+	unsigned char pka_keys[SE2_CONTEXT_SAVE_PKA1_KEYS_LENGTH];
+} tegra_pka_context_t;
+#pragma pack(pop)
+
+/* SE context blob */
+#pragma pack(push, 1)
+typedef struct tegra_se_context_blob {
+	/* SE context */
+	tegra_se_context_t se_ctx;
+	/* Known Pattern */
+	unsigned char known_pattern[SE_CTX_KNOWN_PATTERN_SIZE];
+} tegra_se_context_blob_t;
+#pragma pack(pop)
+
+/* SE2 and PKA1 context blob */
+#pragma pack(push, 1)
+typedef struct tegra_se2_context_blob {
+	/* SE2 context */
+	tegra_se_context_t se_ctx;
+	/* PKA1 context */
+	tegra_pka_context_t pka_ctx;
+	/* Known Pattern */
+	unsigned char known_pattern[SE_CTX_KNOWN_PATTERN_SIZE];
+} tegra_se2_context_blob_t;
+#pragma pack(pop)
+
+/* SE AES key type 128bit, 192bit, 256bit */
+typedef enum {
+	SE_AES_KEY128,
+	SE_AES_KEY192,
+	SE_AES_KEY256,
+} tegra_se_aes_key_type_t;
+
+/* SE RSA key slot */
+typedef struct tegra_se_rsa_key_slot {
+	/* 0 - 63 exponent key */
+	uint32_t exponent[SE_RSA_MAX_EXP_SIZE32];
+	/* 64 - 127 modulus key */
+	uint32_t modulus[SE_RSA_MAX_MOD_SIZE32];
+} tegra_se_rsa_key_slot_t;
+
 
 /*******************************************************************************
  * Inline functions definition
@@ -242,8 +639,21 @@ static inline void tegra_se_write_32(const tegra_se_dev_t *dev, uint32_t offset,
 	mmio_write_32(dev->se_base + offset, val);
 }
 
+static inline uint32_t tegra_pka_read_32(tegra_pka_dev_t *dev, uint32_t offset)
+{
+	return mmio_read_32(dev->pka_base + offset);
+}
+
+static inline void tegra_pka_write_32(tegra_pka_dev_t *dev, uint32_t offset,
+uint32_t val)
+{
+	mmio_write_32(dev->pka_base + offset, val);
+}
+
 /*******************************************************************************
  * Prototypes
  ******************************************************************************/
+int tegra_se_start_normal_operation(const tegra_se_dev_t *, uint32_t);
+int tegra_se_start_ctx_save_operation(const tegra_se_dev_t *, uint32_t);
 
 #endif /* SE_PRIVATE_H */
diff --git a/plat/nvidia/tegra/soc/t210/drivers/se/security_engine.c b/plat/nvidia/tegra/soc/t210/drivers/se/security_engine.c
index a995eac6..9650896f 100644
--- a/plat/nvidia/tegra/soc/t210/drivers/se/security_engine.c
+++ b/plat/nvidia/tegra/soc/t210/drivers/se/security_engine.c
@@ -21,6 +21,7 @@
  ******************************************************************************/
 
 #define TIMEOUT_100MS	100U	// Timeout in 100ms
+#define RNG_AES_KEY_INDEX   1
 
 /*******************************************************************************
  * Data structure and global variables
@@ -67,6 +68,15 @@
  * #--------------------------------#
  */
 
+/* Known pattern data */
+static const uint32_t se_ctx_known_pattern_data[SE_CTX_KNOWN_PATTERN_SIZE_WORDS] = {
+	/* 128 bit AES block */
+	0x0C0D0E0F,
+	0x08090A0B,
+	0x04050607,
+	0x00010203,
+};
+
 /* SE input and output linked list buffers */
 static tegra_se_io_lst_t se1_src_ll_buf;
 static tegra_se_io_lst_t se1_dst_ll_buf;
@@ -78,7 +88,7 @@ static tegra_se_io_lst_t se2_dst_ll_buf;
 /* SE1 security engine device handle */
 static tegra_se_dev_t se_dev_1 = {
 	.se_num = 1,
-	/* setup base address for se */
+	/* Setup base address for se */
 	.se_base = TEGRA_SE1_BASE,
 	/* Setup context size in AES blocks */
 	.ctx_size_blks = SE_CTX_SAVE_SIZE_BLOCKS_SE1,
@@ -86,12 +96,14 @@ static tegra_se_dev_t se_dev_1 = {
 	.src_ll_buf = &se1_src_ll_buf,
 	/* Setup DST buffers for SE operations */
 	.dst_ll_buf = &se1_dst_ll_buf,
+	/* Setup context save destination */
+	.ctx_save_buf = (uint32_t *)(TEGRA_TZRAM_CARVEOUT_BASE),
 };
 
 /* SE2 security engine device handle */
 static tegra_se_dev_t se_dev_2 = {
 	.se_num = 2,
-	/* setup base address for se */
+	/* Setup base address for se */
 	.se_base = TEGRA_SE2_BASE,
 	/* Setup context size in AES blocks */
 	.ctx_size_blks = SE_CTX_SAVE_SIZE_BLOCKS_SE2,
@@ -99,6 +111,8 @@ static tegra_se_dev_t se_dev_2 = {
 	.src_ll_buf = &se2_src_ll_buf,
 	/* Setup DST buffers for SE operations */
 	.dst_ll_buf = &se2_dst_ll_buf,
+	/* Setup context save destination */
+	.ctx_save_buf = (uint32_t *)(TEGRA_TZRAM_CARVEOUT_BASE + 0x1000),
 };
 
 /*******************************************************************************
@@ -189,16 +203,12 @@ static int32_t tegra_se_operation_complete(const tegra_se_dev_t *se_dev)
  * Returns true if the SE engine is configured to perform SE context save in
  * hardware.
  */
-static inline int32_t tegra_se_atomic_save_enabled(const tegra_se_dev_t *se_dev)
+static inline bool tegra_se_atomic_save_enabled(const tegra_se_dev_t *se_dev)
 {
 	uint32_t val;
-	int32_t ret = 0;
 
 	val = tegra_se_read_32(se_dev, SE_CTX_SAVE_AUTO_REG_OFFSET);
-	if (SE_CTX_SAVE_AUTO_ENABLE(val) == SE_CTX_SAVE_AUTO_EN)
-		ret = 1;
-
-	return ret;
+	return (SE_CTX_SAVE_AUTO_ENABLE(val) == SE_CTX_SAVE_AUTO_EN);
 }
 
 /*
@@ -243,14 +253,6 @@ static int32_t tegra_se_context_save_atomic(const tegra_se_dev_t *se_dev)
 	/* Check that previous operation is finalized */
 	ret = tegra_se_operation_prepare(se_dev);
 
-	/* Ensure HW atomic context save has been enabled
-	 * This should have been done at boot time.
-	 * SE_CTX_SAVE_AUTO.ENABLE == ENABLE
-	 */
-	if (ret == 0) {
-		ret = tegra_se_atomic_save_enabled(se_dev);
-	}
-
 	/* Read the context save progress counter: block_count
 	 * Ensure no previous context save has been triggered
 	 * SE_CTX_SAVE_AUTO.CURR_CNT == 0
@@ -309,7 +311,8 @@ static int32_t tegra_se_context_save_atomic(const tegra_se_dev_t *se_dev)
  * Security engine primitive operations, including normal operation
  * and the context save operation.
  */
-static int tegra_se_perform_operation(const tegra_se_dev_t *se_dev, uint32_t nbytes)
+static int tegra_se_perform_operation(const tegra_se_dev_t *se_dev, uint32_t nbytes,
+					bool context_save)
 {
 	uint32_t nblocks = nbytes / TEGRA_SE_AES_BLOCK_SIZE;
 	int ret = 0;
@@ -335,7 +338,10 @@ static int tegra_se_perform_operation(const tegra_se_dev_t *se_dev, uint32_t nby
 	tegra_se_make_data_coherent(se_dev);
 
 	/* Start hardware operation */
-	tegra_se_write_32(se_dev, SE_OPERATION_REG_OFFSET, SE_OP_START);
+	if (context_save)
+		tegra_se_write_32(se_dev, SE_OPERATION_REG_OFFSET, SE_OP_CTX_SAVE);
+	else
+		tegra_se_write_32(se_dev, SE_OPERATION_REG_OFFSET, SE_OP_START);
 
 	/* Wait for operation to finish */
 	ret = tegra_se_operation_complete(se_dev);
@@ -345,6 +351,22 @@ op_error:
 }
 
 /*
+ * Normal security engine operations other than the context save
+ */
+int tegra_se_start_normal_operation(const tegra_se_dev_t *se_dev, uint32_t nbytes)
+{
+	return tegra_se_perform_operation(se_dev, nbytes, false);
+}
+
+/*
+ * Security engine context save operation
+ */
+int tegra_se_start_ctx_save_operation(const tegra_se_dev_t *se_dev, uint32_t nbytes)
+{
+	return tegra_se_perform_operation(se_dev, nbytes, true);
+}
+
+/*
  * Security Engine sequence to generat SRK
  * SE and SE2 will generate different SRK by different
  * entropy seeds.
@@ -375,12 +397,501 @@ static int tegra_se_generate_srk(const tegra_se_dev_t *se_dev)
 	tegra_se_write_32(se_dev, SE_CONFIG_REG_OFFSET, val);
 
 	/* Perform hardware operation */
-	ret = tegra_se_perform_operation(se_dev, 0);
+	ret = tegra_se_start_normal_operation(se_dev, 0);
+
+	return ret;
+}
+
+/*
+ * Generate plain text random data to some memory location using
+ * SE/SE2's SP800-90 random number generator. The random data size
+ * must be some multiple of the AES block size (16 bytes).
+ */
+static int tegra_se_lp_generate_random_data(tegra_se_dev_t *se_dev)
+{
+	int ret = 0;
+	uint32_t val;
+
+	/* Set some arbitrary memory location to store the random data */
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	if (!se_dev->ctx_save_buf) {
+		ERROR("%s: ERR: context save buffer NULL pointer!\n", __func__);
+		return PSCI_E_NOT_PRESENT;
+	}
+	se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(((tegra_se_context_t *)
+					se_dev->ctx_save_buf)->rand_data)));
+	se_dev->dst_ll_buf->buffer[0].data_len = SE_CTX_SAVE_RANDOM_DATA_SIZE;
+
+
+	/* Confgure the following hardware register settings:
+	 * SE_CONFIG.DEC_ALG = NOP
+	 * SE_CONFIG.ENC_ALG = RNG
+	 * SE_CONFIG.ENC_MODE = KEY192
+	 * SE_CONFIG.DST = MEMORY
+	 */
+	val = (SE_CONFIG_ENC_ALG_RNG |
+		SE_CONFIG_DEC_ALG_NOP |
+		SE_CONFIG_ENC_MODE_KEY192 |
+		SE_CONFIG_DST_MEMORY);
+	tegra_se_write_32(se_dev, SE_CONFIG_REG_OFFSET, val);
+
+	/* Program the RNG options in SE_CRYPTO_CONFIG as follows:
+	 * XOR_POS = BYPASS
+	 * INPUT_SEL = RANDOM (Entropy or LFSR)
+	 * HASH_ENB = DISABLE
+	 */
+	val = (SE_CRYPTO_INPUT_RANDOM |
+		SE_CRYPTO_XOR_BYPASS |
+		SE_CRYPTO_CORE_ENCRYPT |
+		SE_CRYPTO_HASH_DISABLE |
+		SE_CRYPTO_KEY_INDEX(RNG_AES_KEY_INDEX) |
+		SE_CRYPTO_IV_ORIGINAL);
+	tegra_se_write_32(se_dev, SE_CRYPTO_REG_OFFSET, val);
+
+	/* Configure RNG */
+	val = (DRBG_MODE_FORCE_INSTANTION | DRBG_SRC_LFSR);
+	tegra_se_write_32(se_dev, SE_RNG_CONFIG_REG_OFFSET, val);
+
+	/* SE normal operation */
+	ret = tegra_se_start_normal_operation(se_dev, SE_CTX_SAVE_RANDOM_DATA_SIZE);
+
+	return ret;
+}
+
+/*
+ * Encrypt memory blocks with SRK as part of the security engine context.
+ * The data blocks include: random data and the known pattern data, where
+ * the random data is the first block and known pattern is the last block.
+ */
+static int tegra_se_lp_data_context_save(tegra_se_dev_t *se_dev,
+		uint64_t src_addr, uint64_t dst_addr, uint32_t data_size)
+{
+	int ret = 0;
+
+	se_dev->src_ll_buf->last_buff_num = 0;
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	se_dev->src_ll_buf->buffer[0].addr = src_addr;
+	se_dev->src_ll_buf->buffer[0].data_len = data_size;
+	se_dev->dst_ll_buf->buffer[0].addr = dst_addr;
+	se_dev->dst_ll_buf->buffer[0].data_len = data_size;
+
+	/* By setting the context source from memory and calling the context save
+	 * operation, the SE encrypts the memory data with SRK.
+	 */
+	tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET, SE_CTX_SAVE_SRC_MEM);
+
+	ret = tegra_se_start_ctx_save_operation(se_dev, data_size);
 
 	return ret;
 }
 
 /*
+ * Context save the key table access control sticky bits and
+ * security status of each key-slot. The encrypted sticky-bits are
+ * 32 bytes (2 AES blocks) and formatted as the following structure:
+ * {	bit in registers			bit in context save
+ *	SECURITY_0[4]				158
+ *	SE_RSA_KEYTABLE_ACCE4SS_1[2:0]		157:155
+ *	SE_RSA_KEYTABLE_ACCE4SS_0[2:0]		154:152
+ *	SE_RSA_SECURITY_PERKEY_0[1:0]		151:150
+ *	SE_CRYPTO_KEYTABLE_ACCESS_15[7:0]	149:142
+ *	...,
+ *	SE_CRYPTO_KEYTABLE_ACCESS_0[7:0]	29:22
+ *	SE_CRYPTO_SECURITY_PERKEY_0[15:0]	21:6
+ *	SE_TZRAM_SECURITY_0[1:0]		5:4
+ *	SE_SECURITY_0[16]			3:3
+ *	SE_SECURITY_0[2:0] }			2:0
+ */
+static int tegra_se_lp_sticky_bits_context_save(tegra_se_dev_t *se_dev)
+{
+	int ret = PSCI_E_INTERN_FAIL;
+	uint32_t val = 0;
+
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	if (!se_dev->ctx_save_buf) {
+		ERROR("%s: ERR: context save buffer NULL pointer!\n", __func__);
+		return PSCI_E_NOT_PRESENT;
+	}
+	se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(((tegra_se_context_t *)
+						se_dev->ctx_save_buf)->sticky_bits)));
+	se_dev->dst_ll_buf->buffer[0].data_len = SE_CTX_SAVE_STICKY_BITS_SIZE;
+
+	/*
+	 * The 1st AES block save the sticky-bits context 1 - 16 bytes (0 - 3 words).
+	 * The 2nd AES block save the sticky-bits context 17 - 32 bytes (4 - 7 words).
+	 */
+	for (int i = 0; i < 2; i++) {
+		val = SE_CTX_SAVE_SRC_STICKY_BITS |
+			SE_CTX_SAVE_STICKY_WORD_QUAD(i);
+		tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+		/* SE context save operation */
+		ret = tegra_se_start_ctx_save_operation(se_dev,
+				SE_CTX_SAVE_STICKY_BITS_SIZE);
+		if (ret)
+			break;
+		se_dev->dst_ll_buf->buffer[0].addr += SE_CTX_SAVE_STICKY_BITS_SIZE;
+	}
+
+	return ret;
+}
+
+static int tegra_se_aeskeytable_context_save(tegra_se_dev_t *se_dev)
+{
+	uint32_t val = 0;
+	int ret = 0;
+
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	if (!se_dev->ctx_save_buf) {
+		ERROR("%s: ERR: context save buffer NULL pointer!\n", __func__);
+		ret = -EINVAL;
+		goto aes_keytable_save_err;
+	}
+
+	/* AES key context save */
+	for (int slot = 0; slot < TEGRA_SE_AES_KEYSLOT_COUNT; slot++) {
+		se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+						((tegra_se_context_t *)se_dev->
+						 ctx_save_buf)->key_slots[slot].key)));
+		se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_KEY_128_SIZE;
+		for (int i = 0; i < 2; i++) {
+			val = SE_CTX_SAVE_SRC_AES_KEYTABLE |
+				SE_CTX_SAVE_KEY_INDEX(slot) |
+				SE_CTX_SAVE_WORD_QUAD(i);
+			tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+			/* SE context save operation */
+			ret = tegra_se_start_ctx_save_operation(se_dev,
+					TEGRA_SE_KEY_128_SIZE);
+			if (ret) {
+				ERROR("%s: ERR: AES key CTX_SAVE OP failed, "
+						"slot=%d, word_quad=%d.\n",
+						__func__, slot, i);
+				goto aes_keytable_save_err;
+			}
+			se_dev->dst_ll_buf->buffer[0].addr += TEGRA_SE_KEY_128_SIZE;
+		}
+
+		/* OIV context save */
+		se_dev->dst_ll_buf->last_buff_num = 0;
+		se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+						((tegra_se_context_t *)se_dev->
+						 ctx_save_buf)->key_slots[slot].oiv)));
+		se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_AES_IV_SIZE;
+
+		val = SE_CTX_SAVE_SRC_AES_KEYTABLE |
+			SE_CTX_SAVE_KEY_INDEX(slot) |
+			SE_CTX_SAVE_WORD_QUAD_ORIG_IV;
+		tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+		/* SE context save operation */
+		ret = tegra_se_start_ctx_save_operation(se_dev, TEGRA_SE_AES_IV_SIZE);
+		if (ret) {
+			ERROR("%s: ERR: OIV CTX_SAVE OP failed, slot=%d.\n",
+					__func__, slot);
+			goto aes_keytable_save_err;
+		}
+
+		/* UIV context save */
+		se_dev->dst_ll_buf->last_buff_num = 0;
+		se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+						((tegra_se_context_t *)se_dev->
+						 ctx_save_buf)->key_slots[slot].uiv)));
+		se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_AES_IV_SIZE;
+
+		val = SE_CTX_SAVE_SRC_AES_KEYTABLE |
+			SE_CTX_SAVE_KEY_INDEX(slot) |
+			SE_CTX_SAVE_WORD_QUAD_UPD_IV;
+		tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+		/* SE context save operation */
+		ret = tegra_se_start_ctx_save_operation(se_dev, TEGRA_SE_AES_IV_SIZE);
+		if (ret) {
+			ERROR("%s: ERR: UIV CTX_SAVE OP failed, slot=%d\n",
+					__func__, slot);
+			goto aes_keytable_save_err;
+		}
+	}
+
+aes_keytable_save_err:
+	return ret;
+}
+
+static int tegra_se_lp_rsakeytable_context_save(tegra_se_dev_t *se_dev)
+{
+	uint32_t val = 0;
+	int ret = 0;
+	/* First the modulus and then the exponent must be
+	 * encrypted and saved. This is repeated for SLOT 0
+	 * and SLOT 1. Hence the order:
+	 * SLOT 0 exponent : RSA_KEY_INDEX : 0
+	 * SLOT 0 modulus : RSA_KEY_INDEX : 1
+	 * SLOT 1 exponent : RSA_KEY_INDEX : 2
+	 * SLOT 1 modulus : RSA_KEY_INDEX : 3
+	 */
+	const unsigned int key_index_mod[TEGRA_SE_RSA_KEYSLOT_COUNT][2] = {
+		/* RSA key slot 0 */
+		{SE_RSA_KEY_INDEX_SLOT0_EXP, SE_RSA_KEY_INDEX_SLOT0_MOD},
+		/* RSA key slot 1 */
+		{SE_RSA_KEY_INDEX_SLOT1_EXP, SE_RSA_KEY_INDEX_SLOT1_MOD},
+	};
+
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+					((tegra_se_context_t *)se_dev->
+					 ctx_save_buf)->rsa_keys)));
+	se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_KEY_128_SIZE;
+
+	for (int slot = 0; slot < TEGRA_SE_RSA_KEYSLOT_COUNT; slot++) {
+		/* loop for modulus and exponent */
+		for (int index = 0; index < 2; index++) {
+			for (int word_quad = 0; word_quad < 16; word_quad++) {
+				val = SE_CTX_SAVE_SRC_RSA_KEYTABLE |
+					SE_CTX_SAVE_RSA_KEY_INDEX(
+						key_index_mod[slot][index]) |
+					SE_CTX_RSA_WORD_QUAD(word_quad);
+				tegra_se_write_32(se_dev,
+					SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+				/* SE context save operation */
+				ret = tegra_se_start_ctx_save_operation(se_dev,
+						TEGRA_SE_KEY_128_SIZE);
+				if (ret) {
+					ERROR("%s: ERR: slot=%d.\n",
+						__func__, slot);
+					goto rsa_keytable_save_err;
+				}
+
+				/* Update the pointer to the next word quad */
+				se_dev->dst_ll_buf->buffer[0].addr +=
+					TEGRA_SE_KEY_128_SIZE;
+			}
+		}
+	}
+
+rsa_keytable_save_err:
+	return ret;
+}
+
+static int tegra_se_pkakeytable_sticky_bits_save(tegra_se_dev_t *se_dev)
+{
+	int ret = 0;
+
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+					((tegra_se2_context_blob_t *)se_dev->
+					 ctx_save_buf)->pka_ctx.sticky_bits)));
+	se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_AES_BLOCK_SIZE;
+
+	/* PKA1 sticky bits are 1 AES block (16 bytes) */
+	tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET,
+			SE_CTX_SAVE_SRC_PKA1_STICKY_BITS |
+			SE_CTX_STICKY_WORD_QUAD_WORDS_0_3);
+
+	/* SE context save operation */
+	ret = tegra_se_start_ctx_save_operation(se_dev, 0);
+	if (ret) {
+		ERROR("%s: ERR: PKA1 sticky bits CTX_SAVE OP failed\n",
+				__func__);
+		goto pka_sticky_bits_save_err;
+	}
+
+pka_sticky_bits_save_err:
+	return ret;
+}
+
+static int tegra_se_pkakeytable_context_save(tegra_se_dev_t *se_dev)
+{
+	uint32_t val = 0;
+	int ret = 0;
+
+	se_dev->dst_ll_buf->last_buff_num = 0;
+	se_dev->dst_ll_buf->buffer[0].addr = ((uint64_t)(&(
+					((tegra_se2_context_blob_t *)se_dev->
+					 ctx_save_buf)->pka_ctx.pka_keys)));
+	se_dev->dst_ll_buf->buffer[0].data_len = TEGRA_SE_KEY_128_SIZE;
+
+	/* for each slot, save word quad 0-127 */
+	for (int slot = 0; slot < TEGRA_SE_PKA1_KEYSLOT_COUNT; slot++) {
+		for (int word_quad = 0; word_quad < 512/4; word_quad++) {
+			val = SE_CTX_SAVE_SRC_PKA1_KEYTABLE |
+				SE_CTX_PKA1_WORD_QUAD_L((slot * 128) +
+						word_quad) |
+				SE_CTX_PKA1_WORD_QUAD_H((slot * 128) +
+						word_quad);
+			tegra_se_write_32(se_dev,
+					SE_CTX_SAVE_CONFIG_REG_OFFSET, val);
+
+			/* SE context save operation */
+			ret = tegra_se_start_ctx_save_operation(se_dev,
+					TEGRA_SE_KEY_128_SIZE);
+			if (ret) {
+				ERROR("%s: ERR: pka1 keytable ctx save error\n",
+						__func__);
+				goto pka_keytable_save_err;
+			}
+
+			/* Update the pointer to the next word quad */
+			se_dev->dst_ll_buf->buffer[0].addr +=
+				TEGRA_SE_KEY_128_SIZE;
+		}
+	}
+
+pka_keytable_save_err:
+	return ret;
+}
+
+static int tegra_se_save_SRK(tegra_se_dev_t *se_dev)
+{
+	tegra_se_write_32(se_dev, SE_CTX_SAVE_CONFIG_REG_OFFSET,
+			SE_CTX_SAVE_SRC_SRK);
+
+	/* SE context save operation */
+	return tegra_se_start_ctx_save_operation(se_dev, 0);
+}
+
+/*
+ *  Lock both SE from non-TZ clients.
+ */
+static inline void tegra_se_lock(tegra_se_dev_t *se_dev)
+{
+	uint32_t val;
+
+	assert(se_dev);
+	val = tegra_se_read_32(se_dev, SE_SECURITY_REG_OFFSET);
+	val |= SE_SECURITY_TZ_LOCK_SOFT(SE_SECURE);
+	tegra_se_write_32(se_dev, SE_SECURITY_REG_OFFSET, val);
+}
+
+/*
+ * Use SRK to encrypt SE state and save to TZRAM carveout
+ */
+static int tegra_se_context_save_sw(tegra_se_dev_t *se_dev)
+{
+	int err = 0;
+
+	assert(se_dev);
+
+	/* Lock entire SE/SE2 as TZ protected */
+	tegra_se_lock(se_dev);
+
+	INFO("%s: generate SRK\n", __func__);
+	/* Generate SRK */
+	err = tegra_se_generate_srk(se_dev);
+	if (err) {
+		ERROR("%s: ERR: SRK generation failed\n", __func__);
+		return err;
+	}
+
+	INFO("%s: generate random data\n", __func__);
+	/* Generate random data */
+	err = tegra_se_lp_generate_random_data(se_dev);
+	if (err) {
+		ERROR("%s: ERR: LP random pattern generation failed\n", __func__);
+		return err;
+	}
+
+	INFO("%s: encrypt random data\n", __func__);
+	/* Encrypt the random data block */
+	err = tegra_se_lp_data_context_save(se_dev,
+		((uint64_t)(&(((tegra_se_context_t *)se_dev->
+					ctx_save_buf)->rand_data))),
+		((uint64_t)(&(((tegra_se_context_t *)se_dev->
+					ctx_save_buf)->rand_data))),
+		SE_CTX_SAVE_RANDOM_DATA_SIZE);
+	if (err) {
+		ERROR("%s: ERR: random pattern encryption failed\n", __func__);
+		return err;
+	}
+
+	INFO("%s: save SE sticky bits\n", __func__);
+	/* Save AES sticky bits context */
+	err = tegra_se_lp_sticky_bits_context_save(se_dev);
+	if (err) {
+		ERROR("%s: ERR: sticky bits context save failed\n", __func__);
+		return err;
+	}
+
+	INFO("%s: save AES keytables\n", __func__);
+	/* Save AES key table context */
+	err = tegra_se_aeskeytable_context_save(se_dev);
+	if (err) {
+		ERROR("%s: ERR: LP keytable save failed\n", __func__);
+		return err;
+	}
+
+	/* RSA key slot table context save */
+	INFO("%s: save RSA keytables\n", __func__);
+	err = tegra_se_lp_rsakeytable_context_save(se_dev);
+	if (err) {
+		ERROR("%s: ERR: rsa key table context save failed\n", __func__);
+		return err;
+	}
+
+	/* Only SE2 has an interface with PKA1; thus, PKA1's context is saved
+	 * via SE2.
+	 */
+	if (se_dev->se_num == 2) {
+		/* Encrypt PKA1 sticky bits on SE2 only */
+		INFO("%s: save PKA sticky bits\n", __func__);
+		err = tegra_se_pkakeytable_sticky_bits_save(se_dev);
+		if (err) {
+			ERROR("%s: ERR: PKA sticky bits context save failed\n", __func__);
+			return err;
+		}
+
+		/* Encrypt PKA1 keyslots on SE2 only */
+		INFO("%s: save PKA keytables\n", __func__);
+		err = tegra_se_pkakeytable_context_save(se_dev);
+		if (err) {
+			ERROR("%s: ERR: PKA key table context save failed\n", __func__);
+			return err;
+		}
+	}
+
+	/* Encrypt known pattern */
+	if (se_dev->se_num == 1) {
+		err = tegra_se_lp_data_context_save(se_dev,
+			((uint64_t)(&se_ctx_known_pattern_data)),
+			((uint64_t)(&(((tegra_se_context_blob_t *)se_dev->ctx_save_buf)->known_pattern))),
+			SE_CTX_KNOWN_PATTERN_SIZE);
+	} else if (se_dev->se_num == 2) {
+		err = tegra_se_lp_data_context_save(se_dev,
+			((uint64_t)(&se_ctx_known_pattern_data)),
+			((uint64_t)(&(((tegra_se2_context_blob_t *)se_dev->ctx_save_buf)->known_pattern))),
+			SE_CTX_KNOWN_PATTERN_SIZE);
+	}
+	if (err) {
+		ERROR("%s: ERR: save LP known pattern failure\n", __func__);
+		return err;
+	}
+
+	/* Write lp context buffer address into PMC scratch register */
+	if (se_dev->se_num == 1) {
+		/* SE context address */
+		mmio_write_32((uint64_t)TEGRA_PMC_BASE + PMC_SECURE_SCRATCH117_OFFSET,
+				((uint64_t)(se_dev->ctx_save_buf)));
+	} else if (se_dev->se_num == 2) {
+		/* SE2 & PKA1 context address */
+		mmio_write_32((uint64_t)TEGRA_PMC_BASE + PMC_SECURE_SCRATCH116_OFFSET,
+				((uint64_t)(se_dev->ctx_save_buf)));
+	}
+
+	/* Saves SRK to PMC secure scratch registers for BootROM, which
+	 * verifies and restores the security engine context on warm boot.
+	 */
+	err = tegra_se_save_SRK(se_dev);
+	if (err < 0) {
+		ERROR("%s: ERR: LP SRK save failure\n", __func__);
+		return err;
+	}
+
+	INFO("%s: SE context save done \n", __func__);
+
+	return err;
+}
+
+/*
  * Initialize the SE engine handle
  */
 void tegra_se_init(void)
@@ -451,18 +962,43 @@ int32_t tegra_se_suspend(void)
 
 	tegra_se_enable_clocks();
 
-	/* Atomic context save se2 and pka1 */
-	INFO("%s: SE2/PKA1 atomic context save\n", __func__);
-	ret = tegra_se_context_save_atomic(&se_dev_2);
+	if (tegra_se_atomic_save_enabled(&se_dev_2) &&
+			tegra_se_atomic_save_enabled(&se_dev_1)) {
+		/* Atomic context save se2 and pka1 */
+		INFO("%s: SE2/PKA1 atomic context save\n", __func__);
+		if (ret == 0) {
+			ret = tegra_se_context_save_atomic(&se_dev_2);
+		}
 
-	/* Atomic context save se */
-	if (ret == 0) {
-		INFO("%s: SE1 atomic context save\n", __func__);
-		ret = tegra_se_context_save_atomic(&se_dev_1);
-	}
+		/* Atomic context save se */
+		if (ret == 0) {
+			INFO("%s: SE1 atomic context save\n", __func__);
+			ret = tegra_se_context_save_atomic(&se_dev_1);
+		}
 
-	if (ret == 0) {
-		INFO("%s: SE atomic context save done\n", __func__);
+		if (ret == 0) {
+			INFO("%s: SE atomic context save done\n", __func__);
+		}
+	} else if (!tegra_se_atomic_save_enabled(&se_dev_2) &&
+			!tegra_se_atomic_save_enabled(&se_dev_1)) {
+		/* SW context save se2 and pka1 */
+		INFO("%s: SE2/PKA1 legacy(SW) context save\n", __func__);
+		if (ret == 0) {
+			ret = tegra_se_context_save_sw(&se_dev_2);
+		}
+
+		/* SW context save se */
+		if (ret == 0) {
+			INFO("%s: SE1 legacy(SW) context save\n", __func__);
+			ret = tegra_se_context_save_sw(&se_dev_1);
+		}
+
+		if (ret == 0) {
+			INFO("%s: SE SW context save done\n", __func__);
+		}
+	} else {
+		ERROR("%s: One SE set for atomic CTX save, the other is not\n",
+			 __func__);
 	}
 
 	tegra_se_disable_clocks();