diff --git a/drivers/thermal/hisi_thermal.c b/drivers/thermal/hisi_thermal.c
index c2cb280..b862506 100644
--- a/drivers/thermal/hisi_thermal.c
+++ b/drivers/thermal/hisi_thermal.c
@@ -23,6 +23,7 @@
 #include <linux/module.h>
 #include <linux/platform_device.h>
 #include <linux/io.h>
+#include <linux/of_device.h>
 
 #include "thermal_core.h"
 
@@ -30,7 +31,7 @@
 #define HI6220_TEMP0_TH				(0x4)
 #define HI6220_TEMP0_RST_TH			(0x8)
 #define HI6220_TEMP0_CFG			(0xC)
-#define HI6220_TEMP0_CFG_SS_MSK		(0xF000)
+#define HI6220_TEMP0_CFG_SS_MSK			(0xF000)
 #define HI6220_TEMP0_CFG_HDAK_MSK		(0x30)
 #define HI6220_TEMP0_EN				(0x10)
 #define HI6220_TEMP0_INT_EN			(0x14)
@@ -41,7 +42,7 @@
 #define HI6220_TEMP_BASE			(-60000)
 #define HI6220_TEMP_RESET			(100000)
 #define HI6220_TEMP_STEP			(785)
-#define HI6220_TEMP_LAG			(3500)
+#define HI6220_TEMP_LAG				(3500)
 
 #define HI6220_DEFAULT_SENSOR		2
 
@@ -52,6 +53,10 @@ struct hisi_thermal_sensor {
 };
 
 struct hisi_thermal_data {
+	int (*get_temp)(struct hisi_thermal_data *data);
+	int (*enable_sensor)(struct hisi_thermal_data *data);
+	int (*disable_sensor)(struct hisi_thermal_data *data);
+	int (*irq_handler)(struct hisi_thermal_data *data);
 	struct platform_device *pdev;
 	struct clk *clk;
 	struct hisi_thermal_sensor sensor;
@@ -192,7 +197,18 @@ static inline void hi6220_thermal_hdak_set(void __iomem *addr, int value)
 	       (value << 4), addr + HI6220_TEMP0_CFG);
 }
 
-static void hi6220_thermal_disable_sensor(struct hisi_thermal_data *data)
+static int hi6220_thermal_irq_handler(struct hisi_thermal_data *data)
+{
+	hi6220_thermal_alarm_clear(data->regs, 1);
+	return 0;
+}
+
+static int hi6220_thermal_get_temp(struct hisi_thermal_data *data)
+{
+	return hi6220_thermal_get_temperature(data->regs);
+}
+
+static int hi6220_thermal_disable_sensor(struct hisi_thermal_data *data)
 {
 	/* disable sensor module */
 	hi6220_thermal_enable(data->regs, 0);
@@ -200,6 +216,8 @@ static void hi6220_thermal_disable_sensor(struct hisi_thermal_data *data)
 	hi6220_thermal_reset_enable(data->regs, 0);
 
 	clk_disable_unprepare(data->clk);
+
+	return 0;
 }
 
 static int hi6220_thermal_enable_sensor(struct hisi_thermal_data *data)
@@ -240,12 +258,48 @@ static int hi6220_thermal_enable_sensor(struct hisi_thermal_data *data)
 	return 0;
 }
 
+static int hi6220_thermal_probe(struct hisi_thermal_data *data)
+{
+	struct platform_device *pdev = data->pdev;
+	struct device *dev = &pdev->dev;
+	struct resource *res;
+	int ret;
+
+	data->get_temp = hi6220_thermal_get_temp;
+	data->enable_sensor = hi6220_thermal_enable_sensor;
+	data->disable_sensor = hi6220_thermal_disable_sensor;
+	data->irq_handler = hi6220_thermal_irq_handler;
+
+	res = platform_get_resource(pdev, IORESOURCE_MEM, 0);
+	data->regs = devm_ioremap_resource(dev, res);
+	if (IS_ERR(data->regs)) {
+		dev_err(dev, "failed to get io address\n");
+		return PTR_ERR(data->regs);
+	}
+
+	data->clk = devm_clk_get(dev, "thermal_clk");
+	if (IS_ERR(data->clk)) {
+		ret = PTR_ERR(data->clk);
+		if (ret != -EPROBE_DEFER)
+			dev_err(dev, "failed to get thermal clk: %d\n", ret);
+		return ret;
+	}
+
+	data->irq = platform_get_irq(pdev, 0);
+	if (data->irq < 0)
+		return data->irq;
+
+	data->sensor.id = HI6220_DEFAULT_SENSOR;
+
+	return 0;
+}
+
 static int hisi_thermal_get_temp(void *__data, int *temp)
 {
 	struct hisi_thermal_data *data = __data;
 	struct hisi_thermal_sensor *sensor = &data->sensor;
 
-	*temp = hi6220_thermal_get_temperature(data->regs);
+	*temp = data->get_temp(data);
 
 	dev_dbg(&data->pdev->dev, "id=%d, temp=%d, thres=%d\n",
 		sensor->id, *temp, sensor->thres_temp);
@@ -263,7 +317,7 @@ static irqreturn_t hisi_thermal_alarm_irq_thread(int irq, void *dev)
 	struct hisi_thermal_sensor *sensor = &data->sensor;
 	int temp = 0;
 
-	hi6220_thermal_alarm_clear(data->regs, 1);
+	data->irq_handler(data);
 
 	hisi_thermal_get_temp(data, &temp);
 
@@ -284,14 +338,11 @@ static irqreturn_t hisi_thermal_alarm_irq_thread(int irq, void *dev)
 
 static int hisi_thermal_register_sensor(struct platform_device *pdev,
 					struct hisi_thermal_data *data,
-					struct hisi_thermal_sensor *sensor,
-					int index)
+					struct hisi_thermal_sensor *sensor)
 {
 	int ret, i;
 	const struct thermal_trip *trip;
 
-	sensor->id = index;
-
 	sensor->tzd = devm_thermal_zone_of_sensor_register(&pdev->dev,
 							   sensor->id, data,
 							   &hisi_of_thermal_ops);
@@ -316,7 +367,7 @@ static int hisi_thermal_register_sensor(struct platform_device *pdev,
 }
 
 static const struct of_device_id of_hisi_thermal_match[] = {
-	{ .compatible = "hisilicon,tsensor" },
+	{ .compatible = "hisilicon,tsensor", .data = hi6220_thermal_probe },
 	{ /* end */ }
 };
 MODULE_DEVICE_TABLE(of, of_hisi_thermal_match);
@@ -333,58 +384,48 @@ static void hisi_thermal_toggle_sensor(struct hisi_thermal_sensor *sensor,
 static int hisi_thermal_probe(struct platform_device *pdev)
 {
 	struct hisi_thermal_data *data;
-	struct resource *res;
+	int const (*platform_probe)(struct hisi_thermal_data *);
+	struct device *dev = &pdev->dev;
 	int ret;
 
-	data = devm_kzalloc(&pdev->dev, sizeof(*data), GFP_KERNEL);
+	data = devm_kzalloc(dev, sizeof(*data), GFP_KERNEL);
 	if (!data)
 		return -ENOMEM;
 
 	data->pdev = pdev;
-
-	res = platform_get_resource(pdev, IORESOURCE_MEM, 0);
-	data->regs = devm_ioremap_resource(&pdev->dev, res);
-	if (IS_ERR(data->regs)) {
-		dev_err(&pdev->dev, "failed to get io address\n");
-		return PTR_ERR(data->regs);
-	}
-
-	data->irq = platform_get_irq(pdev, 0);
-	if (data->irq < 0)
-		return data->irq;
-
 	platform_set_drvdata(pdev, data);
 
-	data->clk = devm_clk_get(&pdev->dev, "thermal_clk");
-	if (IS_ERR(data->clk)) {
-		ret = PTR_ERR(data->clk);
-		if (ret != -EPROBE_DEFER)
-			dev_err(&pdev->dev,
-				"failed to get thermal clk: %d\n", ret);
-		return ret;
+	platform_probe = of_device_get_match_data(dev);
+	if (!platform_probe) {
+		dev_err(dev, "failed to get probe func\n");
+		return -EINVAL;
 	}
 
+	ret = platform_probe(data);
+	if (ret)
+		return ret;
+
 	ret = hisi_thermal_register_sensor(pdev, data,
-					   &data->sensor,
-					   HI6220_DEFAULT_SENSOR);
+					   &data->sensor);
 	if (ret) {
-		dev_err(&pdev->dev, "failed to register thermal sensor: %d\n",
-			ret);
+		dev_err(dev, "failed to register thermal sensor: %d\n", ret);
 		return ret;
 	}
 
-	ret = hi6220_thermal_enable_sensor(data);
+	ret = data->enable_sensor(data);
 	if (ret) {
-		dev_err(&pdev->dev, "Failed to setup the sensor: %d\n", ret);
+		dev_err(dev, "Failed to setup the sensor: %d\n", ret);
 		return ret;
 	}
 
-	ret = devm_request_threaded_irq(&pdev->dev, data->irq, NULL,
-					hisi_thermal_alarm_irq_thread,
-					IRQF_ONESHOT, "hisi_thermal", data);
-	if (ret < 0) {
-		dev_err(&pdev->dev, "failed to request alarm irq: %d\n", ret);
-		return ret;
+	if (data->irq) {
+		ret = devm_request_threaded_irq(dev, data->irq, NULL,
+				hisi_thermal_alarm_irq_thread,
+				IRQF_ONESHOT, "hisi_thermal", data);
+		if (ret < 0) {
+			dev_err(dev, "failed to request alarm irq: %d\n", ret);
+			return ret;
+		}
 	}
 
 	hisi_thermal_toggle_sensor(&data->sensor, true);
@@ -398,7 +439,8 @@ static int hisi_thermal_remove(struct platform_device *pdev)
 	struct hisi_thermal_sensor *sensor = &data->sensor;
 
 	hisi_thermal_toggle_sensor(sensor, false);
-	hi6220_thermal_disable_sensor(data);
+
+	data->disable_sensor(data);
 
 	return 0;
 }
@@ -408,7 +450,7 @@ static int hisi_thermal_suspend(struct device *dev)
 {
 	struct hisi_thermal_data *data = dev_get_drvdata(dev);
 
-	hi6220_thermal_disable_sensor(data);
+	data->disable_sensor(data);
 
 	return 0;
 }
@@ -417,7 +459,7 @@ static int hisi_thermal_resume(struct device *dev)
 {
 	struct hisi_thermal_data *data = dev_get_drvdata(dev);
 
-	return hi6220_thermal_enable_sensor(data);
+	return data->enable_sensor(data);
 }
 #endif