#include "rsmu.h"
 
 /*
- * 16-bit register address: the lower 8 bits of the register address come
- * from the offset addr byte and the upper 8 bits come from the page register.
+ * 32-bit register address: the lower 8 bits of the register address come
+ * from the offset addr byte and the upper 24 bits come from the page register.
  */
-#define        RSMU_CM_PAGE_ADDR               0xFD
-#define        RSMU_CM_PAGE_WINDOW             256
+#define        RSMU_CM_PAGE_ADDR               0xFC
+#define RSMU_CM_PAGE_MASK              0xFFFFFF00
+#define RSMU_CM_ADDRESS_MASK           0x000000FF
 
 /*
  * 15-bit register address: the lower 7 bits of the register address come
 #define        RSMU_SABRE_PAGE_ADDR            0x7F
 #define        RSMU_SABRE_PAGE_WINDOW          128
 
-static const struct regmap_range_cfg rsmu_cm_range_cfg[] = {
-       {
-               .range_min = 0,
-               .range_max = 0xD000,
-               .selector_reg = RSMU_CM_PAGE_ADDR,
-               .selector_mask = 0xFF,
-               .selector_shift = 0,
-               .window_start = 0,
-               .window_len = RSMU_CM_PAGE_WINDOW,
-       }
-};
-
 static const struct regmap_range_cfg rsmu_sabre_range_cfg[] = {
        {
                .range_min = 0,
        }
 };
 
-static bool rsmu_cm_volatile_reg(struct device *dev, unsigned int reg)
+static bool rsmu_sabre_volatile_reg(struct device *dev, unsigned int reg)
 {
        switch (reg) {
-       case RSMU_CM_PAGE_ADDR:
+       case RSMU_SABRE_PAGE_ADDR:
                return false;
        default:
                return true;
        }
 }
 
-static bool rsmu_sabre_volatile_reg(struct device *dev, unsigned int reg)
+static int rsmu_read_device(struct rsmu_ddata *rsmu, u8 reg, u8 *buf, u16 bytes)
 {
-       switch (reg) {
-       case RSMU_SABRE_PAGE_ADDR:
-               return false;
-       default:
-               return true;
+       struct i2c_client *client = to_i2c_client(rsmu->dev);
+       struct i2c_msg msg[2];
+       int cnt;
+
+       msg[0].addr = client->addr;
+       msg[0].flags = 0;
+       msg[0].len = 1;
+       msg[0].buf = ®
+
+       msg[1].addr = client->addr;
+       msg[1].flags = I2C_M_RD;
+       msg[1].len = bytes;
+       msg[1].buf = buf;
+
+       cnt = i2c_transfer(client->adapter, msg, 2);
+
+       if (cnt < 0) {
+               dev_err(rsmu->dev, "i2c_transfer failed at addr: %04x!", reg);
+               return cnt;
+       } else if (cnt != 2) {
+               dev_err(rsmu->dev,
+                       "i2c_transfer sent only %d of 2 messages", cnt);
+               return -EIO;
+       }
+
+       return 0;
+}
+
+static int rsmu_write_device(struct rsmu_ddata *rsmu, u8 reg, u8 *buf, u16 bytes)
+{
+       struct i2c_client *client = to_i2c_client(rsmu->dev);
+       u8 msg[RSMU_MAX_WRITE_COUNT + 1]; /* 1 Byte added for the device register */
+       int cnt;
+
+       if (bytes > RSMU_MAX_WRITE_COUNT)
+               return -EINVAL;
+
+       msg[0] = reg;
+       memcpy(&msg[1], buf, bytes);
+
+       cnt = i2c_master_send(client, msg, bytes + 1);
+
+       if (cnt < 0) {
+               dev_err(&client->dev,
+                       "i2c_master_send failed at addr: %04x!", reg);
+               return cnt;
        }
+
+       return 0;
+}
+
+static int rsmu_write_page_register(struct rsmu_ddata *rsmu, u32 reg)
+{
+       u32 page = reg & RSMU_CM_PAGE_MASK;
+       u8 buf[4];
+       int err;
+
+       /* Do not modify offset register for none-scsr registers */
+       if (reg < RSMU_CM_SCSR_BASE)
+               return 0;
+
+       /* Simply return if we are on the same page */
+       if (rsmu->page == page)
+               return 0;
+
+       buf[0] = 0x0;
+       buf[1] = (u8)((page >> 8) & 0xFF);
+       buf[2] = (u8)((page >> 16) & 0xFF);
+       buf[3] = (u8)((page >> 24) & 0xFF);
+
+       err = rsmu_write_device(rsmu, RSMU_CM_PAGE_ADDR, buf, sizeof(buf));
+       if (err)
+               dev_err(rsmu->dev, "Failed to set page offset 0x%x\n", page);
+       else
+               /* Remember the last page */
+               rsmu->page = page;
+
+       return err;
+}
+
+static int rsmu_reg_read(void *context, unsigned int reg, unsigned int *val)
+{
+       struct rsmu_ddata *rsmu = i2c_get_clientdata((struct i2c_client *)context);
+       u8 addr = (u8)(reg & RSMU_CM_ADDRESS_MASK);
+       int err;
+
+       err = rsmu_write_page_register(rsmu, reg);
+       if (err)
+               return err;
+
+       err = rsmu_read_device(rsmu, addr, (u8 *)val, 1);
+       if (err)
+               dev_err(rsmu->dev, "Failed to read offset address 0x%x\n", addr);
+
+       return err;
+}
+
+static int rsmu_reg_write(void *context, unsigned int reg, unsigned int val)
+{
+       struct rsmu_ddata *rsmu = i2c_get_clientdata((struct i2c_client *)context);
+       u8 addr = (u8)(reg & RSMU_CM_ADDRESS_MASK);
+       u8 data = (u8)val;
+       int err;
+
+       err = rsmu_write_page_register(rsmu, reg);
+       if (err)
+               return err;
+
+       err = rsmu_write_device(rsmu, addr, &data, 1);
+       if (err)
+               dev_err(rsmu->dev,
+                       "Failed to write offset address 0x%x\n", addr);
+
+       return err;
 }
 
 static const struct regmap_config rsmu_cm_regmap_config = {
-       .reg_bits = 8,
+       .reg_bits = 32,
        .val_bits = 8,
-       .max_register = 0xD000,
-       .ranges = rsmu_cm_range_cfg,
-       .num_ranges = ARRAY_SIZE(rsmu_cm_range_cfg),
-       .volatile_reg = rsmu_cm_volatile_reg,
-       .cache_type = REGCACHE_RBTREE,
-       .can_multi_write = true,
+       .max_register = 0x20120000,
+       .reg_read = rsmu_reg_read,
+       .reg_write = rsmu_reg_write,
+       .cache_type = REGCACHE_NONE,
 };
 
 static const struct regmap_config rsmu_sabre_regmap_config = {
        .reg_bits = 16,
        .val_bits = 8,
        .reg_format_endian = REGMAP_ENDIAN_BIG,
-       .max_register = 0x339,
+       .max_register = 0x340,
        .cache_type = REGCACHE_NONE,
        .can_multi_write = true,
 };
 
-static int rsmu_i2c_probe(struct i2c_client *client)
+static int rsmu_i2c_probe(struct i2c_client *client,
+                         const struct i2c_device_id *id)
 {
-       const struct i2c_device_id *id = i2c_client_get_device_id(client);
        const struct regmap_config *cfg;
        struct rsmu_ddata *rsmu;
        int ret;
                dev_err(rsmu->dev, "Unsupported RSMU device type: %d\n", rsmu->type);
                return -ENODEV;
        }
-       rsmu->regmap = devm_regmap_init_i2c(client, cfg);
+
+       if (rsmu->type == RSMU_CM)
+               rsmu->regmap = devm_regmap_init(&client->dev, NULL, client, cfg);
+       else
+               rsmu->regmap = devm_regmap_init_i2c(client, cfg);
        if (IS_ERR(rsmu->regmap)) {
                ret = PTR_ERR(rsmu->regmap);
                dev_err(rsmu->dev, "Failed to allocate register map: %d\n", ret);
                .name = "rsmu-i2c",
                .of_match_table = of_match_ptr(rsmu_i2c_of_match),
        },
-       .probe_new = rsmu_i2c_probe,
+       .probe = rsmu_i2c_probe,
        .remove = rsmu_i2c_remove,
        .id_table = rsmu_i2c_id,
 };
 
 
 #define        RSMU_CM_PAGE_ADDR               0x7C
 #define        RSMU_SABRE_PAGE_ADDR            0x7F
-#define        RSMU_HIGHER_ADDR_MASK           0xFF80
-#define        RSMU_HIGHER_ADDR_SHIFT          7
-#define        RSMU_LOWER_ADDR_MASK            0x7F
+#define        RSMU_PAGE_MASK                  0xFFFFFF80
+#define        RSMU_ADDR_MASK                  0x7F
 
 static int rsmu_read_device(struct rsmu_ddata *rsmu, u8 reg, u8 *buf, u16 bytes)
 {
        struct spi_device *client = to_spi_device(rsmu->dev);
        struct spi_transfer xfer = {0};
        struct spi_message msg;
-       u8 cmd[256] = {0};
-       u8 rsp[256] = {0};
+       u8 cmd[RSMU_MAX_READ_COUNT + 1] = {0};
+       u8 rsp[RSMU_MAX_READ_COUNT + 1] = {0};
        int ret;
 
+       if (bytes > RSMU_MAX_READ_COUNT)
+               return -EINVAL;
+
        cmd[0] = reg | 0x80;
        xfer.rx_buf = rsp;
        xfer.len = bytes + 1;
        struct spi_device *client = to_spi_device(rsmu->dev);
        struct spi_transfer xfer = {0};
        struct spi_message msg;
-       u8 cmd[256] = {0};
+       u8 cmd[RSMU_MAX_WRITE_COUNT + 1] = {0};
+
+       if (bytes > RSMU_MAX_WRITE_COUNT)
+               return -EINVAL;
 
        cmd[0] = reg;
        memcpy(&cmd[1], buf, bytes);
  * 16-bit register address: the lower 7 bits of the register address come
  * from the offset addr byte and the upper 9 bits come from the page register.
  */
-static int rsmu_write_page_register(struct rsmu_ddata *rsmu, u16 reg)
+static int rsmu_write_page_register(struct rsmu_ddata *rsmu, u32 reg)
 {
        u8 page_reg;
-       u8 buf[2];
+       u8 buf[4];
        u16 bytes;
-       u16 page;
+       u32 page;
        int err;
 
        switch (rsmu->type) {
        case RSMU_CM:
+               /* Do not modify page register for none-scsr registers */
+               if (reg < RSMU_CM_SCSR_BASE)
+                       return 0;
                page_reg = RSMU_CM_PAGE_ADDR;
-               page = reg & RSMU_HIGHER_ADDR_MASK;
+               page = reg & RSMU_PAGE_MASK;
                buf[0] = (u8)(page & 0xff);
                buf[1] = (u8)((page >> 8) & 0xff);
-               bytes = 2;
+               buf[2] = (u8)((page >> 16) & 0xff);
+               buf[3] = (u8)((page >> 24) & 0xff);
+               bytes = 4;
                break;
        case RSMU_SABRE:
+               /* Do not modify page register if reg is page register itself */
+               if ((reg & RSMU_ADDR_MASK) == RSMU_ADDR_MASK)
+                       return 0;
                page_reg = RSMU_SABRE_PAGE_ADDR;
-               page = reg >> RSMU_HIGHER_ADDR_SHIFT;
-               buf[0] = (u8)(page & 0xff);
+               page = reg & RSMU_PAGE_MASK;
+               /* The three page bits are located in the single Page Register */
+               buf[0] = (u8)((page >> 7) & 0x7);
                bytes = 1;
                break;
        default:
 
 static int rsmu_reg_read(void *context, unsigned int reg, unsigned int *val)
 {
-       struct rsmu_ddata *rsmu = spi_get_drvdata(context);
-       u8 addr = (u8)(reg & RSMU_LOWER_ADDR_MASK);
+       struct rsmu_ddata *rsmu = spi_get_drvdata((struct spi_device *)context);
+       u8 addr = (u8)(reg & RSMU_ADDR_MASK);
        int err;
 
        err = rsmu_write_page_register(rsmu, reg);
 
 static int rsmu_reg_write(void *context, unsigned int reg, unsigned int val)
 {
-       struct rsmu_ddata *rsmu = spi_get_drvdata(context);
-       u8 addr = (u8)(reg & RSMU_LOWER_ADDR_MASK);
+       struct rsmu_ddata *rsmu = spi_get_drvdata((struct spi_device *)context);
+       u8 addr = (u8)(reg & RSMU_ADDR_MASK);
        u8 data = (u8)val;
        int err;
 
 }
 
 static const struct regmap_config rsmu_cm_regmap_config = {
-       .reg_bits = 16,
+       .reg_bits = 32,
        .val_bits = 8,
-       .max_register = 0xD000,
+       .max_register = 0x20120000,
        .reg_read = rsmu_reg_read,
        .reg_write = rsmu_reg_write,
        .cache_type = REGCACHE_NONE,