[PATCH] mtd: spi-nor: Store stacked memories sizes in parameter array

[Prasad Kummari]

Each flash connected in stacked mode should have a array of integers
to the nor structure to store each flash size. This array stores the
sizes parameters of each flash connected in a stacked or parallel mode.

The current stacked implementation assumes a maximum of two flashes
connected in stacked mode, with both flashes being of the same make
but potentially differing in size. Thus, except for the sizes, all
other flash parameters of both flashes are identical.

In the current implementation of stacked mode operation, the full
chip erase in stack mode does not work correctly for the upper flash.
This problem occurs because spi-nor cannot set the SPI_XFER_U_PAGE
flag. It only sets this flag for the lower flash and
calls spi_nor_erase_chip(), which erases only the lower flash.
Also, spi-nor does not know which chip select values to use.
It will select the upper flash if SPI_XFER_U_PAGE is set, otherwise,
it will select the lower flash.

To fix this issue, spi-nor needs to send the full chip erase command
to both the upper and lower flash. Since spi-nor does not know the
chip select values, it will determine which flash to erase based on
the individual flash sizes stored in array of params. The code changes
will loop through the number of flashes and set the SPI_XFER_U_PAGE
flag to indicate the upper flash. Using the individual index,
spi-nor will call spi_nor_erase_chip() for each flash and wait until
the chip is ready.

Need to rebase top of this commit:
https://lore.kernel.org/u-boot/20241018082644.22495-1-venkatesh.abbarapu@amd.com/T/#u

Signed-off-by: Prasad Kummari <prasad.kummari at amd.com>
---
 drivers/mtd/spi/spi-nor-core.c | 47 ++++++++++++++++++++++++++++++----
 include/linux/mtd/spi-nor.h    |  3 +++
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/drivers/mtd/spi/spi-nor-core.c b/drivers/mtd/spi/spi-nor-core.c
index 2034d84251..0d96373a1e 100644
--- a/drivers/mtd/spi/spi-nor-core.c
+++ b/drivers/mtd/spi/spi-nor-core.c
@@ -1063,7 +1063,7 @@ static int spi_nor_erase_chip(struct spi_nor *nor)
 	if (ret)
 		return ret;
 
-	return nor->mtd.size;
+	return ret;
 }
 
 /*
@@ -1104,7 +1104,7 @@ static int spi_nor_erase(struct mtd_info *mtd, struct erase_info *instr)
 	struct spi_nor *nor = mtd_to_spi_nor(mtd);
 	u32 addr, len, rem, offset, max_size;
 	bool addr_known = false;
-	int ret, err;
+	int ret, err, idx = 0, rc = 0;
 
 	dev_dbg(nor->dev, "at 0x%llx, len %lld\n", (long long)instr->addr,
 		(long long)instr->len);
@@ -1155,7 +1155,35 @@ static int spi_nor_erase(struct mtd_info *mtd, struct erase_info *instr)
 
 		if (len == mtd->size &&
 		    !(nor->flags & SNOR_F_NO_OP_CHIP_ERASE)) {
-			ret = spi_nor_erase_chip(nor);
+			if (nor->flags & SNOR_F_HAS_STACKED) {
+
+				ret = 0; idx = 0;
+				while (idx < nor->num_flash) {
+					if (idx == 1) {
+						nor->spi->flags |= SPI_XFER_U_PAGE;
+						rc = write_enable(nor);
+						if (rc < 0)
+							goto erase_err;
+					}
+					rc = spi_nor_erase_chip(nor);
+					if (rc < 0)
+						goto erase_err;
+
+					rc = spi_nor_erase_chip_wait_till_ready(nor,
+										nor->cs_params[idx]);
+					if (rc)
+						goto erase_err;
+
+					ret += nor->cs_params[idx];
+					idx++;
+				}
+
+			} else {
+				ret = spi_nor_erase_chip(nor);
+				if (ret < 0)
+					goto erase_err;
+				ret = nor->mtd.size;
+			}
 		} else {
 			ret = spi_nor_erase_sector(nor, offset);
 		}
@@ -1166,7 +1194,7 @@ static int spi_nor_erase(struct mtd_info *mtd, struct erase_info *instr)
 		len -= ret;
 
 		if (max_size == mtd->size &&
-		    !(nor->flags & SNOR_F_NO_OP_CHIP_ERASE)) {
+		    !(nor->flags & SNOR_F_NO_OP_CHIP_ERASE) && !idx) {
 			ret = spi_nor_erase_chip_wait_till_ready(nor, mtd->size);
 		} else {
 			ret = spi_nor_wait_till_ready(nor);
@@ -3204,6 +3232,9 @@ static int spi_nor_init_params(struct spi_nor *nor,
 		spi_nor_post_sfdp_fixups(nor, params);
 	}
 #if CONFIG_IS_ENABLED(DM_SPI) && CONFIG_IS_ENABLED(SPI_ADVANCE)
+
+	nor->num_flash = 0;
+
 	/*
 	 * The flashes that are connected in stacked mode should be of same make.
 	 * Except the flash size all other properties are identical for all the
@@ -3225,6 +3256,7 @@ static int spi_nor_init_params(struct spi_nor *nor,
 			if (!(nor->spi->flags & SPI_XFER_STACKED))
 				nor->spi->flags |= SPI_XFER_STACKED;
 		}
+		nor->num_flash++;
 	}
 
 	i = 0;
@@ -3242,13 +3274,18 @@ static int spi_nor_init_params(struct spi_nor *nor,
 			if (!(nor->flags & SNOR_F_HAS_PARALLEL))
 				nor->flags |= SNOR_F_HAS_PARALLEL;
 		}
+		nor->num_flash++;
 	}
 
 	if (nor->flags & (SNOR_F_HAS_STACKED | SNOR_F_HAS_PARALLEL)) {
 		params->size = 0;
-		for (idx = 0; idx < SNOR_FLASH_CNT_MAX; idx++)
+
+		for (idx = 0; idx < nor->num_flash; idx++) {
+			nor->cs_params[idx] = flash_size[idx];
 			params->size += flash_size[idx];
+		}
 	}
+
 	/*
 	 * In parallel-memories the erase operation is
 	 * performed on both the flashes simultaneously
diff --git a/include/linux/mtd/spi-nor.h b/include/linux/mtd/spi-nor.h
index 1ae586b2e5..f0719206c5 100644
--- a/include/linux/mtd/spi-nor.h
+++ b/include/linux/mtd/spi-nor.h
@@ -616,6 +616,9 @@ struct spi_nor {
 	u32 size;
 	u32 sector_size;
 	u32 erase_size;
+/* Number of flashes connected in parallel or stacked mode parameters */
+	u32 num_flash;
+	u64 cs_params[SNOR_FLASH_CNT_MAX];
 };
 
 #ifndef __UBOOT__
-- 
2.25.1