"ad0", "ad1", "ad2",
 };
 
+static void ad7266_reg_disable(void *reg)
+{
+       regulator_disable(reg);
+}
+
 static int ad7266_probe(struct spi_device *spi)
 {
        struct ad7266_platform_data *pdata = spi->dev.platform_data;
                if (ret)
                        return ret;
 
+               ret = devm_add_action_or_reset(&spi->dev, ad7266_reg_disable, st->reg);
+               if (ret)
+                       return ret;
+
                ret = regulator_get_voltage(st->reg);
                if (ret < 0)
-                       goto error_disable_reg;
+                       return ret;
 
                st->vref_mv = ret / 1000;
        } else {
                                                      GPIOD_OUT_LOW);
                                if (IS_ERR(st->gpios[i])) {
                                        ret = PTR_ERR(st->gpios[i]);
-                                       goto error_disable_reg;
+                                       return ret;
                                }
                        }
                }
                st->mode = AD7266_MODE_DIFF;
        }
 
-       spi_set_drvdata(spi, indio_dev);
        st->spi = spi;
 
        indio_dev->name = spi_get_device_id(spi)->name;
        spi_message_add_tail(&st->single_xfer[1], &st->single_msg);
        spi_message_add_tail(&st->single_xfer[2], &st->single_msg);
 
-       ret = iio_triggered_buffer_setup(indio_dev, &iio_pollfunc_store_time,
+       ret = devm_iio_triggered_buffer_setup(&spi->dev, indio_dev, &iio_pollfunc_store_time,
                &ad7266_trigger_handler, &iio_triggered_buffer_setup_ops);
        if (ret)
-               goto error_disable_reg;
-
-       ret = iio_device_register(indio_dev);
-       if (ret)
-               goto error_buffer_cleanup;
-
-       return 0;
-
-error_buffer_cleanup:
-       iio_triggered_buffer_cleanup(indio_dev);
-error_disable_reg:
-       if (!IS_ERR(st->reg))
-               regulator_disable(st->reg);
-
-       return ret;
-}
-
-static void ad7266_remove(struct spi_device *spi)
-{
-       struct iio_dev *indio_dev = spi_get_drvdata(spi);
-       struct ad7266_state *st = iio_priv(indio_dev);
+               return ret;
 
-       iio_device_unregister(indio_dev);
-       iio_triggered_buffer_cleanup(indio_dev);
-       if (!IS_ERR(st->reg))
-               regulator_disable(st->reg);
+       return devm_iio_device_register(&spi->dev, indio_dev);
 }
 
 static const struct spi_device_id ad7266_id[] = {
                .name   = "ad7266",
        },
        .probe          = ad7266_probe,
-       .remove         = ad7266_remove,
        .id_table       = ad7266_id,
 };
 module_spi_driver(ad7266_driver);