struct xgene_reboot_context {
        struct device *dev;
-       void *csr;
+       void __iomem *csr;
        u32 mask;
        struct notifier_block restart_handler;
 };
        if (!ctx)
                return -ENOMEM;
 
-       ctx->csr = of_iomap(dev->of_node, 0);
+       ctx->csr = devm_platform_ioremap_resource(pdev, 0);
        if (!ctx->csr) {
                dev_err(dev, "can not map resource\n");
                return -ENODEV;
        ctx->restart_handler.notifier_call = xgene_restart_handler;
        ctx->restart_handler.priority = 128;
        err = register_restart_handler(&ctx->restart_handler);
-       if (err) {
-               iounmap(ctx->csr);
+       if (err)
                dev_err(dev, "cannot register restart handler (err=%d)\n", err);
-       }
 
        return err;
 }