#include "jobdesc.h"
 #include "rsa_caam.h"
 
-#ifdef CONFIG_MX6
+#if defined(CONFIG_MX6) || defined(CONFIG_MX7)
 /*!
  * Secure memory run command
  *
 {
        uint32_t temp_reg;
 
-       sec_out32(CAAM_SMCJR0, sec_mem_cmd);
+       ccsr_sec_t *sec = (void *)CONFIG_SYS_FSL_SEC_ADDR;
+       uint32_t sm_vid = SM_VERSION(sec_in32(&sec->smvid));
+       uint32_t jr_id = 0;
+
+       sec_out32(CAAM_SMCJR(sm_vid, jr_id), sec_mem_cmd);
 
        do {
-               temp_reg = sec_in32(CAAM_SMCSJR0);
+               temp_reg = sec_in32(CAAM_SMCSJR(sm_vid, jr_id));
        } while (temp_reg & CMD_COMPLETE);
 
        return temp_reg;
 {
        uint32_t temp_reg;
 
+       ccsr_sec_t *sec = (void *)CONFIG_SYS_FSL_SEC_ADDR;
+       uint32_t sm_vid = SM_VERSION(sec_in32(&sec->smvid));
+       uint32_t jr_id = 0;
+
        /*
         * De-Allocate partition_num if already allocated to ARM core
         */
        }
 
        /* set the access rights to allow full access */
-       sec_out32(CAAM_SMAG1JR0(partition_num), 0xF);
-       sec_out32(CAAM_SMAG2JR0(partition_num), 0xF);
-       sec_out32(CAAM_SMAPJR0(partition_num), 0xFF);
+       sec_out32(CAAM_SMAG1JR(sm_vid, jr_id, partition_num), 0xF);
+       sec_out32(CAAM_SMAG2JR(sm_vid, jr_id, partition_num), 0xF);
+       sec_out32(CAAM_SMAPJR(sm_vid, jr_id, partition_num), 0xFF);
 
        /* Now need to allocate partition_num of secure RAM. */
        /* De-Allocate page_num by starting with a page inquiry command */
 int inline_cnstr_jobdesc_blob_dek(uint32_t *desc, const uint8_t *plain_txt,
                                       uint8_t *dek_blob, uint32_t in_sz)
 {
+       ccsr_sec_t *sec = (void *)CONFIG_SYS_FSL_SEC_ADDR;
+       uint32_t sm_vid = SM_VERSION(sec_in32(&sec->smvid));
+       uint32_t jr_id = 0;
+
        uint32_t ret = 0;
        u32 aad_w1, aad_w2;
        /* output blob will have 32 bytes key blob in beginning and
        flush_dcache_range(start, end);
 
        /* Now configure the access rights of the partition */
-       sec_out32(CAAM_SMAG1JR0(PARTITION_1), KS_G1); /* set group 1 */
-       sec_out32(CAAM_SMAG2JR0(PARTITION_1), 0);     /* clear group 2 */
-       sec_out32(CAAM_SMAPJR0(PARTITION_1), PERM);   /* set perm & locks */
+       sec_out32(CAAM_SMAG1JR(sm_vid, jr_id, PARTITION_1), KS_G1);
+       sec_out32(CAAM_SMAG2JR(sm_vid, jr_id, PARTITION_1), 0);
+       sec_out32(CAAM_SMAPJR(sm_vid, jr_id, PARTITION_1), PERM);
 
        /* construct aad for AES */
        aad_w1 = (in_sz << OP_ALG_ALGSEL_SHIFT) | KEY_AES_SRC | LD_CCM_MODE;
 
        u32     drr;            /* DECO Reset Register */
        u8      res5[0x4d8];
        struct rng4tst rng;     /* RNG Registers */
-       u8      res11[0x8a0];
+       u8      res6[0x8a0];
        u32     crnr_ms;        /* CHA Revision Number Register, MS */
        u32     crnr_ls;        /* CHA Revision Number Register, LS */
        u32     ctpr_ms;        /* Compile Time Parameters Register, MS */
        u32     ctpr_ls;        /* Compile Time Parameters Register, LS */
-       u8      res6[0x10];
+       u8      res7[0x10];
        u32     far_ms;         /* Fault Address Register, MS */
        u32     far_ls;         /* Fault Address Register, LS */
        u32     falr;           /* Fault Address LIODN Register */
        u32     fadr;           /* Fault Address Detail Register */
-       u8      res7[0x4];
+       u8      res8[0x4];
        u32     csta;           /* CAAM Status Register */
-       u8      res8[0x8];
+       u32     smpart;         /* Secure Memory Partition Parameters */
+       u32     smvid;          /* Secure Memory Version ID */
        u32     rvid;           /* Run Time Integrity Checking Version ID Reg.*/
        u32     ccbvid;         /* CHA Cluster Block Version ID Register */
        u32     chavid_ms;      /* CHA Version ID Register, MS */
 #define CONFIG_JRSTARTR_JR0            0x00000001
 
 struct jr_regs {
-#if defined(CONFIG_SYS_FSL_SEC_LE) && !defined(CONFIG_MX6)
+#if defined(CONFIG_SYS_FSL_SEC_LE) && \
+       !(defined(CONFIG_MX6) || defined(CONFIG_MX7))
        u32 irba_l;
        u32 irba_h;
 #else
        u32 irsa;
        u32 rsvd3;
        u32 irja;
-#if defined(CONFIG_SYS_FSL_SEC_LE) && !defined(CONFIG_MX6)
+#if defined(CONFIG_SYS_FSL_SEC_LE) && \
+       !(defined(CONFIG_MX6) || defined(CONFIG_MX7))
        u32 orba_l;
        u32 orba_h;
 #else
  * related information
  */
 struct sg_entry {
-#if defined(CONFIG_SYS_FSL_SEC_LE) && !defined(CONFIG_MX6)
+#if defined(CONFIG_SYS_FSL_SEC_LE) && \
+       !(defined(CONFIG_MX6) || defined(CONFIG_MX7))
        uint32_t addr_lo;       /* Memory Address - lo */
        uint32_t addr_hi;       /* Memory Address of start of buffer - hi */
 #else
 #define SG_ENTRY_OFFSET_SHIFT  0
 };
 
-#ifdef CONFIG_MX6
+#if defined(CONFIG_MX6) || defined(CONFIG_MX7)
+/* Job Ring Base Address */
+#define JR_BASE_ADDR(x) (CONFIG_SYS_FSL_SEC_ADDR + 0x1000 * (x + 1))
+/* Secure Memory Offset varies accross versions */
+#define SM_V1_OFFSET 0x0f4
+#define SM_V2_OFFSET 0xa00
+/*Secure Memory Versioning */
+#define SMVID_V2 0x20105
+#define SM_VERSION(x)  (x < SMVID_V2 ? 1 : 2)
+#define SM_OFFSET(x)  (x == 1 ? SM_V1_OFFSET : SM_V2_OFFSET)
 /* CAAM Job Ring 0 Registers */
 /* Secure Memory Partition Owner register */
 #define SMCSJR_PO              (3 << 6)
 /* JR Allocation Error */
 #define SMCSJR_AERR            (3 << 12)
 /* Secure memory partition 0 page 0 owner register */
-#define CAAM_SMPO_0            CONFIG_SYS_FSL_SEC_ADDR + 0x1FBC
+#define CAAM_SMPO_0        (CONFIG_SYS_FSL_SEC_ADDR + 0x1FBC)
 /* Secure memory command register */
-#define CAAM_SMCJR0            CONFIG_SYS_FSL_SEC_ADDR + 0x10f4
+#define CAAM_SMCJR(v, jr)   (JR_BASE_ADDR(jr) + SM_OFFSET(v) + SM_CMD(v))
 /* Secure memory command status register */
-#define CAAM_SMCSJR0           CONFIG_SYS_FSL_SEC_ADDR + 0x10fc
+#define CAAM_SMCSJR(v, jr)  (JR_BASE_ADDR(jr) + SM_OFFSET(v) + SM_STATUS(v))
 /* Secure memory access permissions register */
-#define CAAM_SMAPJR0(y)        (CONFIG_SYS_FSL_SEC_ADDR + 0x1104 + y*16)
+#define CAAM_SMAPJR(v, jr, y) \
+       (JR_BASE_ADDR(jr) + SM_OFFSET(v) + SM_PERM(v) + y * 16)
 /* Secure memory access group 2 register */
-#define CAAM_SMAG2JR0(y)       (CONFIG_SYS_FSL_SEC_ADDR + 0x1108 + y*16)
+#define CAAM_SMAG2JR(v, jr, y) \
+       (JR_BASE_ADDR(jr) + SM_OFFSET(v) + SM_GROUP2(v) + y * 16)
 /* Secure memory access group 1 register */
-#define CAAM_SMAG1JR0(y)       (CONFIG_SYS_FSL_SEC_ADDR + 0x110C + y*16)
+#define CAAM_SMAG1JR(v, jr, y)  \
+       (JR_BASE_ADDR(jr) + SM_OFFSET(v) + SM_GROUP1(v) + y * 16)
 
 /* Commands and macros for secure memory */
+#define SM_CMD(v)              (v == 1 ? 0x0 : 0x1E4)
+#define SM_STATUS(v)           (v == 1 ? 0x8 : 0x1EC)
+#define SM_PERM(v)             (v == 1 ?  0x10 : 0x4)
+#define SM_GROUP2(v)           (v == 1 ? 0x14 : 0x8)
+#define SM_GROUP1(v)           (v == 1 ? 0x18 : 0xC)
 #define CMD_PAGE_ALLOC         0x1
 #define CMD_PAGE_DEALLOC       0x2
 #define CMD_PART_DEALLOC       0x3