1 files changed, 9 insertions, 22 deletions

diff --git a/drivers/char/hw_random/n2-drv.c b/drivers/char/hw_random/n2-drv.c
index 20b962e1d83..292a5889f67 100644
--- a/drivers/char/hw_random/n2-drv.c
+++ b/drivers/char/hw_random/n2-drv.c
@@ -7,7 +7,6 @@
 #include <linux/module.h>
 #include <linux/types.h>
 #include <linux/delay.h>
-#include <linux/init.h>
 #include <linux/slab.h>
 #include <linux/workqueue.h>
 #include <linux/preempt.h>
@@ -633,7 +632,7 @@ static int n2rng_probe(struct platform_device *op)
 	multi_capable = (match->data != NULL);
 
 	n2rng_driver_version();
-	np = kzalloc(sizeof(*np), GFP_KERNEL);
+	np = devm_kzalloc(&op->dev, sizeof(*np), GFP_KERNEL);
 	if (!np)
 		goto out;
 	np->op = op;
@@ -654,7 +653,7 @@ static int n2rng_probe(struct platform_device *op)
 					 &np->hvapi_minor)) {
 			dev_err(&op->dev, "Cannot register suitable "
 				"HVAPI version.\n");
-			goto out_free;
+			goto out;
 		}
 	}
 
@@ -677,15 +676,16 @@ static int n2rng_probe(struct platform_device *op)
 	dev_info(&op->dev, "Registered RNG HVAPI major %lu minor %lu\n",
 		 np->hvapi_major, np->hvapi_minor);
 
-	np->units = kzalloc(sizeof(struct n2rng_unit) * np->num_units,
-			    GFP_KERNEL);
+	np->units = devm_kzalloc(&op->dev,
+				 sizeof(struct n2rng_unit) * np->num_units,
+				 GFP_KERNEL);
 	err = -ENOMEM;
 	if (!np->units)
 		goto out_hvapi_unregister;
 
 	err = n2rng_init_control(np);
 	if (err)
-		goto out_free_units;
+		goto out_hvapi_unregister;
 
 	dev_info(&op->dev, "Found %s RNG, units: %d\n",
 		 ((np->flags & N2RNG_FLAG_MULTI) ?
@@ -698,30 +698,24 @@ static int n2rng_probe(struct platform_device *op)
 
 	err = hwrng_register(&np->hwrng);
 	if (err)
-		goto out_free_units;
+		goto out_hvapi_unregister;
 
-	dev_set_drvdata(&op->dev, np);
+	platform_set_drvdata(op, np);
 
 	schedule_delayed_work(&np->work, 0);
 
 	return 0;
 
-out_free_units:
-	kfree(np->units);
-	np->units = NULL;
-
 out_hvapi_unregister:
 	sun4v_hvapi_unregister(HV_GRP_RNG);
 
-out_free:
-	kfree(np);
 out:
 	return err;
 }
 
 static int n2rng_remove(struct platform_device *op)
 {
-	struct n2rng *np = dev_get_drvdata(&op->dev);
+	struct n2rng *np = platform_get_drvdata(op);
 
 	np->flags |= N2RNG_FLAG_SHUTDOWN;
 
@@ -731,13 +725,6 @@ static int n2rng_remove(struct platform_device *op)
 
 	sun4v_hvapi_unregister(HV_GRP_RNG);
 
-	kfree(np->units);
-	np->units = NULL;
-
-	kfree(np);
-
-	dev_set_drvdata(&op->dev, NULL);
-
 	return 0;
 }