diff --git a/drivers/fsi/fsi-sbefifo.c b/drivers/fsi/fsi-sbefifo.c
index fde1c34743a0..6ca5817910cd 100644
--- a/drivers/fsi/fsi-sbefifo.c
+++ b/drivers/fsi/fsi-sbefifo.c
@@ -1022,9 +1022,9 @@ static void sbefifo_free(struct device *dev)
  * Probe/remove
  */
 
-static int sbefifo_probe(struct device *dev)
+static int sbefifo_probe(struct fsi_device *fsi_dev)
 {
-	struct fsi_device *fsi_dev = to_fsi_dev(dev);
+	struct device *dev = &fsi_dev->dev;
 	struct sbefifo *sbefifo;
 	struct device_node *np;
 	struct platform_device *child;
@@ -1045,7 +1045,7 @@ static int sbefifo_probe(struct device *dev)
 
 	sbefifo->magic = SBEFIFO_MAGIC;
 	sbefifo->fsi_dev = fsi_dev;
-	dev_set_drvdata(dev, sbefifo);
+	fsi_set_drvdata(fsi_dev, sbefifo);
 	mutex_init(&sbefifo->lock);
 	sbefifo->timeout_in_cmd_ms = SBEFIFO_TIMEOUT_IN_CMD;
 	sbefifo->timeout_start_rsp_ms = SBEFIFO_TIMEOUT_START_RSP;
@@ -1101,9 +1101,10 @@ static int sbefifo_unregister_child(struct device *dev, void *data)
 	return 0;
 }
 
-static int sbefifo_remove(struct device *dev)
+static void sbefifo_remove(struct fsi_device *fsi_dev)
 {
-	struct sbefifo *sbefifo = dev_get_drvdata(dev);
+	struct device *dev = &fsi_dev->dev;
+	struct sbefifo *sbefifo = fsi_get_drvdata(fsi_dev);
 
 	dev_dbg(dev, "Removing sbefifo device...\n");
 
@@ -1117,8 +1118,6 @@ static int sbefifo_remove(struct device *dev)
 	fsi_free_minor(sbefifo->dev.devt);
 	device_for_each_child(dev, NULL, sbefifo_unregister_child);
 	put_device(&sbefifo->dev);
-
-	return 0;
 }
 
 static const struct fsi_device_id sbefifo_ids[] = {
@@ -1131,10 +1130,10 @@ static const struct fsi_device_id sbefifo_ids[] = {
 
 static struct fsi_driver sbefifo_drv = {
 	.id_table = sbefifo_ids,
+	.probe = sbefifo_probe,
+	.remove = sbefifo_remove,
 	.drv = {
 		.name = DEVICE_NAME,
-		.probe = sbefifo_probe,
-		.remove = sbefifo_remove,
 	}
 };