diff --git a/drivers/thermal/qcom-spmi-temp-alarm.c b/drivers/thermal/qcom-spmi-temp-alarm.c
index ad4f3a8d6560..936e4dde4298 100644
--- a/drivers/thermal/qcom-spmi-temp-alarm.c
+++ b/drivers/thermal/qcom-spmi-temp-alarm.c
@@ -23,6 +23,8 @@
 #include <linux/regmap.h>
 #include <linux/thermal.h>
 
+#include "thermal_core.h"
+
 #define QPNP_TM_REG_TYPE		0x04
 #define QPNP_TM_REG_SUBTYPE		0x05
 #define QPNP_TM_REG_STATUS		0x08
@@ -37,9 +39,11 @@
 #define STATUS_GEN2_STATE_MASK		GENMASK(6, 4)
 #define STATUS_GEN2_STATE_SHIFT		4
 
-#define SHUTDOWN_CTRL1_OVERRIDE_MASK	GENMASK(7, 6)
+#define SHUTDOWN_CTRL1_OVERRIDE_S2	BIT(6)
 #define SHUTDOWN_CTRL1_THRESHOLD_MASK	GENMASK(1, 0)
 
+#define SHUTDOWN_CTRL1_RATE_25HZ	BIT(3)
+
 #define ALARM_CTRL_FORCE_ENABLE		BIT(7)
 
 /*
@@ -56,12 +60,17 @@
 #define TEMP_THRESH_STEP		5000	/* Threshold step: 5 C */
 
 #define THRESH_MIN			0
+#define THRESH_MAX			3
+
+#define STAGE2_THRESHOLD_MIN		125000	/* Stage 2 Threshold Min: 125 C */
+#define STAGE2_THRESHOLD_MAX		140000	/* Stage 2 Threshold Max: 140 C */
 
 /* Temperature in Milli Celsius reported during stage 0 if no ADC is present */
 #define DEFAULT_TEMP			37000
 
 struct qpnp_tm_chip {
 	struct regmap			*map;
+	struct device			*dev;
 	struct thermal_zone_device	*tz_dev;
 	unsigned int			subtype;
 	long				temp;
@@ -69,6 +78,10 @@ struct qpnp_tm_chip {
 	unsigned int			stage;
 	unsigned int			prev_stage;
 	unsigned int			base;
+	/* protects .thresh, .stage and chip registers */
+	struct mutex			lock;
+	bool				initialized;
+
 	struct iio_channel		*adc;
 };
 
@@ -125,6 +138,8 @@ static int qpnp_tm_update_temp_no_adc(struct qpnp_tm_chip *chip)
 	unsigned int stage, stage_new, stage_old;
 	int ret;
 
+	WARN_ON(!mutex_is_locked(&chip->lock));
+
 	ret = qpnp_tm_get_temp_stage(chip);
 	if (ret < 0)
 		return ret;
@@ -163,8 +178,15 @@ static int qpnp_tm_get_temp(void *data, int *temp)
 	if (!temp)
 		return -EINVAL;
 
-	if (!chip->adc) {
+	if (!chip->initialized) {
+		*temp = DEFAULT_TEMP;
+		return 0;
+	}
+
+	if (!chip->adc)) {
+		mutex_lock(&chip->lock);
 		ret = qpnp_tm_update_temp_no_adc(chip);
+		mutex_unlock(&chip->lock);
 		if (ret < 0)
 			return ret;
 	} else {
@@ -180,8 +202,77 @@ static int qpnp_tm_get_temp(void *data, int *temp)
 	return 0;
 }
 
+static int qpnp_tm_update_critical_trip_temp(struct qpnp_tm_chip *chip,
+					     int temp)
+{
+	u8 reg;
+	bool disable_s2_shutdown = false;
+	int ret;
+
+	WARN_ON(!mutex_is_locked(&chip->lock));
+
+	/*
+	 * Default: S2 and S3 shutdown enabled, thresholds at
+	 * 105C/125C/145C, monitoring at 25Hz
+	 */
+	reg = SHUTDOWN_CTRL1_RATE_25HZ;
+
+	if ((temp == THERMAL_TEMP_INVALID) ||
+	    (temp < STAGE2_THRESHOLD_MIN)) {
+		chip->thresh = THRESH_MIN;
+		goto skip;
+	}
+
+	if (temp <= STAGE2_THRESHOLD_MAX) {
+		chip->thresh = THRESH_MAX -
+			((STAGE2_THRESHOLD_MAX - temp) /
+			 TEMP_THRESH_STEP);
+		disable_s2_shutdown = true;
+	} else {
+		chip->thresh = THRESH_MAX;
+
+		if (!IS_ERR(chip->adc))
+			disable_s2_shutdown = true;
+		else
+			dev_warn(chip->dev,
+				 "No ADC is configured and critical temperature is above the maximum stage 2 threshold of 140°C! Configuring stage 2 shutdown at 140°C.\n");
+	}
+
+skip:
+	reg |= chip->thresh;
+	if (disable_s2_shutdown)
+		reg |= SHUTDOWN_CTRL1_OVERRIDE_S2;
+
+	ret = qpnp_tm_write(chip, QPNP_TM_REG_SHUTDOWN_CTRL1, reg);
+	if (ret < 0)
+		return ret;
+
+	return ret;
+}
+
+static int qpnp_tm_set_trip_temp(void *data, int trip, int temp)
+{
+	struct qpnp_tm_chip *chip = data;
+	const struct thermal_trip *trip_points;
+	int ret;
+
+	trip_points = of_thermal_get_trip_points(chip->tz_dev);
+	if (!trip_points)
+		return -EINVAL;
+
+	if (trip_points[trip].type != THERMAL_TRIP_CRITICAL)
+		return 0;
+
+	mutex_lock(&chip->lock);
+	ret = qpnp_tm_update_critical_trip_temp(chip, temp);
+	mutex_unlock(&chip->lock);
+
+	return ret;
+}
+
 static const struct thermal_zone_of_device_ops qpnp_tm_sensor_ops = {
 	.get_temp = qpnp_tm_get_temp,
+	.set_trip_temp = qpnp_tm_set_trip_temp,
 };
 
 static irqreturn_t qpnp_tm_isr(int irq, void *data)
@@ -193,6 +284,29 @@ static irqreturn_t qpnp_tm_isr(int irq, void *data)
 	return IRQ_HANDLED;
 }
 
+static int qpnp_tm_get_critical_trip_temp(struct qpnp_tm_chip *chip)
+{
+	int ntrips;
+	const struct thermal_trip *trips;
+	int i;
+
+	ntrips = of_thermal_get_ntrips(chip->tz_dev);
+	if (ntrips <= 0)
+		return THERMAL_TEMP_INVALID;
+
+	trips = of_thermal_get_trip_points(chip->tz_dev);
+	if (!trips)
+		return THERMAL_TEMP_INVALID;
+
+	for (i = 0; i < ntrips; i++) {
+		if (of_thermal_is_trip_valid(chip->tz_dev, i) &&
+		    (trips[i].type == THERMAL_TRIP_CRITICAL))
+			return trips[i].temperature;
+	}
+
+	return THERMAL_TEMP_INVALID;
+}
+
 /*
  * This function initializes the internal temp value based on only the
  * current thermal stage and threshold. Setup threshold control and
@@ -203,17 +317,20 @@ static int qpnp_tm_init(struct qpnp_tm_chip *chip)
 	unsigned int stage;
 	int ret;
 	u8 reg = 0;
+	int crit_temp;
+
+	mutex_lock(&chip->lock);
 
 	ret = qpnp_tm_read(chip, QPNP_TM_REG_SHUTDOWN_CTRL1, &reg);
 	if (ret < 0)
-		return ret;
+		goto out;
 
 	chip->thresh = reg & SHUTDOWN_CTRL1_THRESHOLD_MASK;
 	chip->temp = DEFAULT_TEMP;
 
 	ret = qpnp_tm_get_temp_stage(chip);
 	if (ret < 0)
-		return ret;
+		goto out;
 	chip->stage = ret;
 
 	stage = chip->subtype == QPNP_TM_SUBTYPE_GEN1
@@ -224,21 +341,17 @@ static int qpnp_tm_init(struct qpnp_tm_chip *chip)
 			     (stage - 1) * TEMP_STAGE_STEP +
 			     TEMP_THRESH_MIN;
 
-	/*
-	 * Set threshold and disable software override of stage 2 and 3
-	 * shutdowns.
-	 */
-	chip->thresh = THRESH_MIN;
-	reg &= ~(SHUTDOWN_CTRL1_OVERRIDE_MASK | SHUTDOWN_CTRL1_THRESHOLD_MASK);
-	reg |= chip->thresh & SHUTDOWN_CTRL1_THRESHOLD_MASK;
-	ret = qpnp_tm_write(chip, QPNP_TM_REG_SHUTDOWN_CTRL1, reg);
+	crit_temp = qpnp_tm_get_critical_trip_temp(chip);
+	ret = qpnp_tm_update_critical_trip_temp(chip, crit_temp);
 	if (ret < 0)
-		return ret;
+		goto out;
 
 	/* Enable the thermal alarm PMIC module in always-on mode. */
 	reg = ALARM_CTRL_FORCE_ENABLE;
 	ret = qpnp_tm_write(chip, QPNP_TM_REG_ALARM_CTRL, reg);
 
+out:
+	mutex_unlock(&chip->lock);
 	return ret;
 }
 
@@ -257,6 +370,9 @@ static int qpnp_tm_probe(struct platform_device *pdev)
 		return -ENOMEM;
 
 	dev_set_drvdata(&pdev->dev, chip);
+	chip->dev = &pdev->dev;
+
+	mutex_init(&chip->lock);
 
 	chip->map = dev_get_regmap(pdev->dev.parent, NULL);
 	if (!chip->map)
@@ -302,6 +418,18 @@ static int qpnp_tm_probe(struct platform_device *pdev)
 
 	chip->subtype = subtype;
 
+	/*
+	 * Register the sensor before initializing the hardware to be able to
+	 * read the trip points. get_temp() returns the default temperature
+	 * before the hardware initialization is completed.
+	 */
+	chip->tz_dev = devm_thermal_zone_of_sensor_register(&pdev->dev, 0, chip,
+							&qpnp_tm_sensor_ops);
+	if (IS_ERR(chip->tz_dev)) {
+		dev_err(&pdev->dev, "failed to register sensor\n");
+		return PTR_ERR(chip->tz_dev);
+	}
+
 	ret = qpnp_tm_init(chip);
 	if (ret < 0) {
 		dev_err(&pdev->dev, "init failed\n");
@@ -313,12 +441,7 @@ static int qpnp_tm_probe(struct platform_device *pdev)
 	if (ret < 0)
 		return ret;
 
-	chip->tz_dev = devm_thermal_zone_of_sensor_register(&pdev->dev, 0, chip,
-							&qpnp_tm_sensor_ops);
-	if (IS_ERR(chip->tz_dev)) {
-		dev_err(&pdev->dev, "failed to register sensor\n");
-		return PTR_ERR(chip->tz_dev);
-	}
+	chip->initialized = true;
 
 	return 0;
 }
