.read_raw = adc128_read_raw,
 };
 
+static void adc128_disable_regulator(void *reg)
+{
+       regulator_disable(reg);
+}
+
 static int adc128_probe(struct spi_device *spi)
 {
        struct iio_dev *indio_dev;
        adc = iio_priv(indio_dev);
        adc->spi = spi;
 
-       spi_set_drvdata(spi, indio_dev);
-
        indio_dev->name = spi_get_device_id(spi)->name;
        indio_dev->modes = INDIO_DIRECT_MODE;
        indio_dev->info = &adc128_info;
        ret = regulator_enable(adc->reg);
        if (ret < 0)
                return ret;
-
-       mutex_init(&adc->lock);
-
-       ret = iio_device_register(indio_dev);
+       ret = devm_add_action_or_reset(&spi->dev, adc128_disable_regulator,
+                                      adc->reg);
        if (ret)
-               goto err_disable_regulator;
-
-       return 0;
-
-err_disable_regulator:
-       regulator_disable(adc->reg);
-       return ret;
-}
-
-static int adc128_remove(struct spi_device *spi)
-{
-       struct iio_dev *indio_dev = spi_get_drvdata(spi);
-       struct adc128 *adc = iio_priv(indio_dev);
+               return ret;
 
-       iio_device_unregister(indio_dev);
-       regulator_disable(adc->reg);
+       mutex_init(&adc->lock);
 
-       return 0;
+       return devm_iio_device_register(&spi->dev, indio_dev);
 }
 
 static const struct of_device_id adc128_of_match[] = {
                .acpi_match_table = ACPI_PTR(adc128_acpi_match),
        },
        .probe = adc128_probe,
-       .remove = adc128_remove,
        .id_table = adc128_id,
 };
 module_spi_driver(adc128_driver);