1 files changed, 508 insertions, 105 deletions

diff --git a/drivers/soc/tegra/pmc.c b/drivers/soc/tegra/pmc.c
index bc34cf7482fb..bb173456bbff 100644
--- a/drivers/soc/tegra/pmc.c
+++ b/drivers/soc/tegra/pmc.c
@@ -28,12 +28,16 @@
 #include <linux/export.h>
 #include <linux/init.h>
 #include <linux/io.h>
+#include <linux/iopoll.h>
 #include <linux/of.h>
 #include <linux/of_address.h>
+#include <linux/of_platform.h>
 #include <linux/platform_device.h>
+#include <linux/pm_domain.h>
 #include <linux/reboot.h>
 #include <linux/reset.h>
 #include <linux/seq_file.h>
+#include <linux/slab.h>
 #include <linux/spinlock.h>
 
 #include <soc/tegra/common.h>
@@ -101,6 +105,16 @@
 
 #define GPU_RG_CNTRL			0x2d4
 
+struct tegra_powergate {
+	struct generic_pm_domain genpd;
+	struct tegra_pmc *pmc;
+	unsigned int id;
+	struct clk **clks;
+	unsigned int num_clks;
+	struct reset_control **resets;
+	unsigned int num_resets;
+};
+
 struct tegra_pmc_soc {
 	unsigned int num_powergates;
 	const char *const *powergates;
@@ -113,8 +127,11 @@ struct tegra_pmc_soc {
 
 /**
  * struct tegra_pmc - NVIDIA Tegra PMC
+ * @dev: pointer to PMC device structure
  * @base: pointer to I/O remapped register region
  * @clk: pointer to pclk clock
+ * @soc: pointer to SoC data structure
+ * @debugfs: pointer to debugfs entry
  * @rate: currently configured rate of pclk
  * @suspend_mode: lowest suspend mode available
  * @cpu_good_time: CPU power good time (in microseconds)
@@ -128,12 +145,14 @@ struct tegra_pmc_soc {
  * @cpu_pwr_good_en: CPU power good signal is enabled
  * @lp0_vec_phys: physical base address of the LP0 warm boot code
  * @lp0_vec_size: size of the LP0 warm boot code
+ * @powergates_available: Bitmap of available power gates
  * @powergates_lock: mutex for power gate register access
  */
 struct tegra_pmc {
 	struct device *dev;
 	void __iomem *base;
 	struct clk *clk;
+	struct dentry *debugfs;
 
 	const struct tegra_pmc_soc *soc;
 
@@ -151,6 +170,7 @@ struct tegra_pmc {
 	bool cpu_pwr_good_en;
 	u32 lp0_vec_phys;
 	u32 lp0_vec_size;
+	DECLARE_BITMAP(powergates_available, TEGRA_POWERGATE_MAX);
 
 	struct mutex powergates_lock;
 };
@@ -160,6 +180,12 @@ static struct tegra_pmc *pmc = &(struct tegra_pmc) {
 	.suspend_mode = TEGRA_SUSPEND_NONE,
 };
 
+static inline struct tegra_powergate *
+to_powergate(struct generic_pm_domain *domain)
+{
+	return container_of(domain, struct tegra_powergate, genpd);
+}
+
 static u32 tegra_pmc_readl(unsigned long offset)
 {
 	return readl(pmc->base + offset);
@@ -170,38 +196,287 @@ static void tegra_pmc_writel(u32 value, unsigned long offset)
 	writel(value, pmc->base + offset);
 }
 
+static inline bool tegra_powergate_state(int id)
+{
+	if (id == TEGRA_POWERGATE_3D && pmc->soc->has_gpu_clamps)
+		return (tegra_pmc_readl(GPU_RG_CNTRL) & 0x1) == 0;
+	else
+		return (tegra_pmc_readl(PWRGATE_STATUS) & BIT(id)) != 0;
+}
+
+static inline bool tegra_powergate_is_valid(int id)
+{
+	return (pmc->soc && pmc->soc->powergates[id]);
+}
+
+static inline bool tegra_powergate_is_available(int id)
+{
+	return test_bit(id, pmc->powergates_available);
+}
+
+static int tegra_powergate_lookup(struct tegra_pmc *pmc, const char *name)
+{
+	unsigned int i;
+
+	if (!pmc || !pmc->soc || !name)
+		return -EINVAL;
+
+	for (i = 0; i < pmc->soc->num_powergates; i++) {
+		if (!tegra_powergate_is_valid(i))
+			continue;
+
+		if (!strcmp(name, pmc->soc->powergates[i]))
+			return i;
+	}
+
+	dev_err(pmc->dev, "powergate %s not found\n", name);
+
+	return -ENODEV;
+}
+
 /**
  * tegra_powergate_set() - set the state of a partition
  * @id: partition ID
  * @new_state: new state of the partition
  */
-static int tegra_powergate_set(int id, bool new_state)
+static int tegra_powergate_set(unsigned int id, bool new_state)
 {
 	bool status;
+	int err;
 
-	mutex_lock(&pmc->powergates_lock);
+	if (id == TEGRA_POWERGATE_3D && pmc->soc->has_gpu_clamps)
+		return -EINVAL;
 
-	status = tegra_pmc_readl(PWRGATE_STATUS) & (1 << id);
+	mutex_lock(&pmc->powergates_lock);
 
-	if (status == new_state) {
+	if (tegra_powergate_state(id) == new_state) {
 		mutex_unlock(&pmc->powergates_lock);
 		return 0;
 	}
 
 	tegra_pmc_writel(PWRGATE_TOGGLE_START | id, PWRGATE_TOGGLE);
 
+	err = readx_poll_timeout(tegra_powergate_state, id, status,
+				 status == new_state, 10, 100000);
+
+	mutex_unlock(&pmc->powergates_lock);
+
+	return err;
+}
+
+static int __tegra_powergate_remove_clamping(unsigned int id)
+{
+	u32 mask;
+
+	mutex_lock(&pmc->powergates_lock);
+
+	/*
+	 * On Tegra124 and later, the clamps for the GPU are controlled by a
+	 * separate register (with different semantics).
+	 */
+	if (id == TEGRA_POWERGATE_3D) {
+		if (pmc->soc->has_gpu_clamps) {
+			tegra_pmc_writel(0, GPU_RG_CNTRL);
+			goto out;
+		}
+	}
+
+	/*
+	 * Tegra 2 has a bug where PCIE and VDE clamping masks are
+	 * swapped relatively to the partition ids
+	 */
+	if (id == TEGRA_POWERGATE_VDEC)
+		mask = (1 << TEGRA_POWERGATE_PCIE);
+	else if (id == TEGRA_POWERGATE_PCIE)
+		mask = (1 << TEGRA_POWERGATE_VDEC);
+	else
+		mask = (1 << id);
+
+	tegra_pmc_writel(mask, REMOVE_CLAMPING);
+
+out:
 	mutex_unlock(&pmc->powergates_lock);
 
 	return 0;
 }
 
+static void tegra_powergate_disable_clocks(struct tegra_powergate *pg)
+{
+	unsigned int i;
+
+	for (i = 0; i < pg->num_clks; i++)
+		clk_disable_unprepare(pg->clks[i]);
+}
+
+static int tegra_powergate_enable_clocks(struct tegra_powergate *pg)
+{
+	unsigned int i;
+	int err;
+
+	for (i = 0; i < pg->num_clks; i++) {
+		err = clk_prepare_enable(pg->clks[i]);
+		if (err)
+			goto out;
+	}
+
+	return 0;
+
+out:
+	while (i--)
+		clk_disable_unprepare(pg->clks[i]);
+
+	return err;
+}
+
+static int tegra_powergate_reset_assert(struct tegra_powergate *pg)
+{
+	unsigned int i;
+	int err;
+
+	for (i = 0; i < pg->num_resets; i++) {
+		err = reset_control_assert(pg->resets[i]);
+		if (err)
+			return err;
+	}
+
+	return 0;
+}
+
+static int tegra_powergate_reset_deassert(struct tegra_powergate *pg)
+{
+	unsigned int i;
+	int err;
+
+	for (i = 0; i < pg->num_resets; i++) {
+		err = reset_control_deassert(pg->resets[i]);
+		if (err)
+			return err;
+	}
+
+	return 0;
+}
+
+static int tegra_powergate_power_up(struct tegra_powergate *pg,
+				    bool disable_clocks)
+{
+	int err;
+
+	err = tegra_powergate_reset_assert(pg);
+	if (err)
+		return err;
+
+	usleep_range(10, 20);
+
+	err = tegra_powergate_set(pg->id, true);
+	if (err < 0)
+		return err;
+
+	usleep_range(10, 20);
+
+	err = tegra_powergate_enable_clocks(pg);
+	if (err)
+		goto disable_clks;
+
+	usleep_range(10, 20);
+
+	err = __tegra_powergate_remove_clamping(pg->id);
+	if (err)
+		goto disable_clks;
+
+	usleep_range(10, 20);
+
+	err = tegra_powergate_reset_deassert(pg);
+	if (err)
+		goto powergate_off;
+
+	usleep_range(10, 20);
+
+	if (disable_clocks)
+		tegra_powergate_disable_clocks(pg);
+
+	return 0;
+
+disable_clks:
+	tegra_powergate_disable_clocks(pg);
+	usleep_range(10, 20);
+powergate_off:
+	tegra_powergate_set(pg->id, false);
+
+	return err;
+}
+
+static int tegra_powergate_power_down(struct tegra_powergate *pg)
+{
+	int err;
+
+	err = tegra_powergate_enable_clocks(pg);
+	if (err)
+		return err;
+
+	usleep_range(10, 20);
+
+	err = tegra_powergate_reset_assert(pg);
+	if (err)
+		goto disable_clks;
+
+	usleep_range(10, 20);
+
+	tegra_powergate_disable_clocks(pg);
+
+	usleep_range(10, 20);
+
+	err = tegra_powergate_set(pg->id, false);
+	if (err)
+		goto assert_resets;
+
+	return 0;
+
+assert_resets:
+	tegra_powergate_enable_clocks(pg);
+	usleep_range(10, 20);
+	tegra_powergate_reset_deassert(pg);
+	usleep_range(10, 20);
+disable_clks:
+	tegra_powergate_disable_clocks(pg);
+
+	return err;
+}
+
+static int tegra_genpd_power_on(struct generic_pm_domain *domain)
+{
+	struct tegra_powergate *pg = to_powergate(domain);
+	struct tegra_pmc *pmc = pg->pmc;
+	int err;
+
+	err = tegra_powergate_power_up(pg, true);
+	if (err)
+		dev_err(pmc->dev, "failed to turn on PM domain %s: %d\n",
+			pg->genpd.name, err);
+
+	return err;
+}
+
+static int tegra_genpd_power_off(struct generic_pm_domain *domain)
+{
+	struct tegra_powergate *pg = to_powergate(domain);
+	struct tegra_pmc *pmc = pg->pmc;
+	int err;
+
+	err = tegra_powergate_power_down(pg);
+	if (err)
+		dev_err(pmc->dev, "failed to turn off PM domain %s: %d\n",
+			pg->genpd.name, err);
+
+	return err;
+}
+
 /**
  * tegra_powergate_power_on() - power on partition
  * @id: partition ID
  */
-int tegra_powergate_power_on(int id)
+int tegra_powergate_power_on(unsigned int id)
 {
-	if (!pmc->soc || id < 0 || id >= pmc->soc->num_powergates)
+	if (!tegra_powergate_is_available(id))
 		return -EINVAL;
 
 	return tegra_powergate_set(id, true);
@@ -211,9 +486,9 @@ int tegra_powergate_power_on(int id)
  * tegra_powergate_power_off() - power off partition
  * @id: partition ID
  */
-int tegra_powergate_power_off(int id)
+int tegra_powergate_power_off(unsigned int id)
 {
-	if (!pmc->soc || id < 0 || id >= pmc->soc->num_powergates)
+	if (!tegra_powergate_is_available(id))
 		return -EINVAL;
 
 	return tegra_powergate_set(id, false);
@@ -224,53 +499,30 @@ EXPORT_SYMBOL(tegra_powergate_power_off);
  * tegra_powergate_is_powered() - check if partition is powered
  * @id: partition ID
  */
-int tegra_powergate_is_powered(int id)
+int tegra_powergate_is_powered(unsigned int id)
 {
-	u32 status;
+	int status;
 
-	if (!pmc->soc || id < 0 || id >= pmc->soc->num_powergates)
+	if (!tegra_powergate_is_valid(id))
 		return -EINVAL;
 
-	status = tegra_pmc_readl(PWRGATE_STATUS) & (1 << id);
-	return !!status;
+	mutex_lock(&pmc->powergates_lock);
+	status = tegra_powergate_state(id);
+	mutex_unlock(&pmc->powergates_lock);
+
+	return status;
 }
 
 /**
  * tegra_powergate_remove_clamping() - remove power clamps for partition
  * @id: partition ID
  */
-int tegra_powergate_remove_clamping(int id)
+int tegra_powergate_remove_clamping(unsigned int id)
 {
-	u32 mask;
-
-	if (!pmc->soc || id < 0 || id >= pmc->soc->num_powergates)
+	if (!tegra_powergate_is_available(id))
 		return -EINVAL;
 
-	/*
-	 * On Tegra124 and later, the clamps for the GPU are controlled by a
-	 * separate register (with different semantics).
-	 */
-	if (id == TEGRA_POWERGATE_3D) {
-		if (pmc->soc->has_gpu_clamps) {
-			tegra_pmc_writel(0, GPU_RG_CNTRL);
-			return 0;
-		}
-	}
-
-	/*
-	 * Tegra 2 has a bug where PCIE and VDE clamping masks are
-	 * swapped relatively to the partition ids
-	 */
-	if (id == TEGRA_POWERGATE_VDEC)
-		mask = (1 << TEGRA_POWERGATE_PCIE);
-	else if (id == TEGRA_POWERGATE_PCIE)
-		mask = (1 << TEGRA_POWERGATE_VDEC);
-	else
-		mask = (1 << id);
-
-	tegra_pmc_writel(mask, REMOVE_CLAMPING);
-
-	return 0;
+	return __tegra_powergate_remove_clamping(id);
 }
 EXPORT_SYMBOL(tegra_powergate_remove_clamping);
 
@@ -282,38 +534,23 @@ EXPORT_SYMBOL(tegra_powergate_remove_clamping);
  *
  * Must be called with clk disabled, and returns with clk enabled.
  */
-int tegra_powergate_sequence_power_up(int id, struct clk *clk,
+int tegra_powergate_sequence_power_up(unsigned int id, struct clk *clk,
 				      struct reset_control *rst)
 {
-	int ret;
-
-	reset_control_assert(rst);
-
-	ret = tegra_powergate_power_on(id);
-	if (ret)
-		goto err_power;
-
-	ret = clk_prepare_enable(clk);
-	if (ret)
-		goto err_clk;
-
-	usleep_range(10, 20);
-
-	ret = tegra_powergate_remove_clamping(id);
-	if (ret)
-		goto err_clamp;
+	struct tegra_powergate pg;
+	int err;
 
-	usleep_range(10, 20);
-	reset_control_deassert(rst);
+	pg.id = id;
+	pg.clks = &clk;
+	pg.num_clks = 1;
+	pg.resets = &rst;
+	pg.num_resets = 1;
 
-	return 0;
+	err = tegra_powergate_power_up(&pg, false);
+	if (err)
+		pr_err("failed to turn on partition %d: %d\n", id, err);
 
-err_clamp:
-	clk_disable_unprepare(clk);
-err_clk:
-	tegra_powergate_power_off(id);
-err_power:
-	return ret;
+	return err;
 }
 EXPORT_SYMBOL(tegra_powergate_sequence_power_up);
 
@@ -325,9 +562,9 @@ EXPORT_SYMBOL(tegra_powergate_sequence_power_up);
  * Returns the partition ID corresponding to the CPU partition ID or a
  * negative error code on failure.
  */
-static int tegra_get_cpu_powergate_id(int cpuid)
+static int tegra_get_cpu_powergate_id(unsigned int cpuid)
 {
-	if (pmc->soc && cpuid > 0 && cpuid < pmc->soc->num_cpu_powergates)
+	if (pmc->soc && cpuid < pmc->soc->num_cpu_powergates)
 		return pmc->soc->cpu_powergates[cpuid];
 
 	return -EINVAL;
@@ -337,7 +574,7 @@ static int tegra_get_cpu_powergate_id(int cpuid)
  * tegra_pmc_cpu_is_powered() - check if CPU partition is powered
  * @cpuid: CPU partition ID
  */
-bool tegra_pmc_cpu_is_powered(int cpuid)
+bool tegra_pmc_cpu_is_powered(unsigned int cpuid)
 {
 	int id;
 
@@ -352,7 +589,7 @@ bool tegra_pmc_cpu_is_powered(int cpuid)
  * tegra_pmc_cpu_power_on() - power on CPU partition
  * @cpuid: CPU partition ID
  */
-int tegra_pmc_cpu_power_on(int cpuid)
+int tegra_pmc_cpu_power_on(unsigned int cpuid)
 {
 	int id;
 
@@ -367,7 +604,7 @@ int tegra_pmc_cpu_power_on(int cpuid)
  * tegra_pmc_cpu_remove_clamping() - remove power clamps for CPU partition
  * @cpuid: CPU partition ID
  */
-int tegra_pmc_cpu_remove_clamping(int cpuid)
+int tegra_pmc_cpu_remove_clamping(unsigned int cpuid)
 {
 	int id;
 
@@ -416,16 +653,18 @@ static struct notifier_block tegra_pmc_restart_handler = {
 static int powergate_show(struct seq_file *s, void *data)
 {
 	unsigned int i;
+	int status;
 
 	seq_printf(s, " powergate powered\n");
 	seq_printf(s, "------------------\n");
 
 	for (i = 0; i < pmc->soc->num_powergates; i++) {
-		if (!pmc->soc->powergates[i])
+		status = tegra_powergate_is_powered(i);
+		if (status < 0)
 			continue;
 
 		seq_printf(s, " %9s %7s\n", pmc->soc->powergates[i],
-			   tegra_powergate_is_powered(i) ? "yes" : "no");
+			   status ? "yes" : "no");
 	}
 
 	return 0;
@@ -445,17 +684,164 @@ static const struct file_operations powergate_fops = {
 
 static int tegra_powergate_debugfs_init(void)
 {
-	struct dentry *d;
+	pmc->debugfs = debugfs_create_file("powergate", S_IRUGO, NULL, NULL,
+					   &powergate_fops);
+	if (!pmc->debugfs)
+		return -ENOMEM;
+
+	return 0;
+}
+
+static int tegra_powergate_of_get_clks(struct tegra_powergate *pg,
+				       struct device_node *np)
+{
+	struct clk *clk;
+	unsigned int i, count;
+	int err;
+
+	count = of_count_phandle_with_args(np, "clocks", "#clock-cells");
+	if (count == 0)
+		return -ENODEV;
+
+	pg->clks = kcalloc(count, sizeof(clk), GFP_KERNEL);
+	if (!pg->clks)
+		return -ENOMEM;
+
+	for (i = 0; i < count; i++) {
+		pg->clks[i] = of_clk_get(np, i);
+		if (IS_ERR(pg->clks[i])) {
+			err = PTR_ERR(pg->clks[i]);
+			goto err;
+		}
+	}
+
+	pg->num_clks = count;
+
+	return 0;
+
+err:
+	while (i--)
+		clk_put(pg->clks[i]);
+	kfree(pg->clks);
+
+	return err;
+}
+
+static int tegra_powergate_of_get_resets(struct tegra_powergate *pg,
+					 struct device_node *np)
+{
+	struct reset_control *rst;
+	unsigned int i, count;
+	int err;
+
+	count = of_count_phandle_with_args(np, "resets", "#reset-cells");
+	if (count == 0)
+		return -ENODEV;
 
-	d = debugfs_create_file("powergate", S_IRUGO, NULL, NULL,
-				&powergate_fops);
-	if (!d)
+	pg->resets = kcalloc(count, sizeof(rst), GFP_KERNEL);
+	if (!pg->resets)
 		return -ENOMEM;
 
+	for (i = 0; i < count; i++) {
+		pg->resets[i] = of_reset_control_get_by_index(np, i);
+		if (IS_ERR(pg->resets[i])) {
+			err = PTR_ERR(pg->resets[i]);
+			goto error;
+		}
+	}
+
+	pg->num_resets = count;
+
 	return 0;
+
+error:
+	while (i--)
+		reset_control_put(pg->resets[i]);
+	kfree(pg->resets);
+
+	return err;
+}
+
+static void tegra_powergate_add(struct tegra_pmc *pmc, struct device_node *np)
+{
+	struct tegra_powergate *pg;
+	bool off;
+	int id;
+
+	pg = kzalloc(sizeof(*pg), GFP_KERNEL);
+	if (!pg)
+		goto error;
+
+	id = tegra_powergate_lookup(pmc, np->name);
+	if (id < 0)
+		goto free_mem;
+
+	/*
+	 * Clear the bit for this powergate so it cannot be managed
+	 * directly via the legacy APIs for controlling powergates.
+	 */
+	clear_bit(id, pmc->powergates_available);
+
+	pg->id = id;
+	pg->genpd.name = np->name;
+	pg->genpd.power_off = tegra_genpd_power_off;
+	pg->genpd.power_on = tegra_genpd_power_on;
+	pg->pmc = pmc;
+
+	if (tegra_powergate_of_get_clks(pg, np))
+		goto set_available;
+
+	if (tegra_powergate_of_get_resets(pg, np))
+		goto remove_clks;
+
+	off = !tegra_powergate_is_powered(pg->id);
+
+	pm_genpd_init(&pg->genpd, NULL, off);
+
+	if (of_genpd_add_provider_simple(np, &pg->genpd))
+		goto remove_resets;
+
+	dev_dbg(pmc->dev, "added power domain %s\n", pg->genpd.name);
+
+	return;
+
+remove_resets:
+	while (pg->num_resets--)
+		reset_control_put(pg->resets[pg->num_resets]);
+	kfree(pg->resets);
+
+remove_clks:
+	while (pg->num_clks--)
+		clk_put(pg->clks[pg->num_clks]);
+	kfree(pg->clks);
+
+set_available:
+	set_bit(id, pmc->powergates_available);
+
+free_mem:
+	kfree(pg);
+
+error:
+	dev_err(pmc->dev, "failed to create power domain for %s\n", np->name);
+}
+
+static void tegra_powergate_init(struct tegra_pmc *pmc)
+{
+	struct device_node *np, *child;
+
+	np = of_get_child_by_name(pmc->dev->of_node, "powergates");
+	if (!np)
+		return;
+
+	for_each_child_of_node(np, child) {
+		tegra_powergate_add(pmc, child);
+		of_node_put(child);
+	}
+
+	of_node_put(np);
 }
 
-static int tegra_io_rail_prepare(int id, unsigned long *request,
+static int tegra_io_rail_prepare(unsigned int id, unsigned long *request,
 				 unsigned long *status, unsigned int *bit)
 {
 	unsigned long rate, value;
@@ -512,15 +898,17 @@ static void tegra_io_rail_unprepare(void)
 	tegra_pmc_writel(DPD_SAMPLE_DISABLE, DPD_SAMPLE);
 }
 
-int tegra_io_rail_power_on(int id)
+int tegra_io_rail_power_on(unsigned int id)
 {
 	unsigned long request, status, value;
 	unsigned int bit, mask;
 	int err;
 
+	mutex_lock(&pmc->powergates_lock);
+
 	err = tegra_io_rail_prepare(id, &request, &status, &bit);
-	if (err < 0)
-		return err;
+	if (err)
+		goto error;
 
 	mask = 1 << bit;
 
@@ -531,27 +919,32 @@ int tegra_io_rail_power_on(int id)
 	tegra_pmc_writel(value, request);
 
 	err = tegra_io_rail_poll(status, mask, 0, 250);
-	if (err < 0) {
+	if (err) {
 		pr_info("tegra_io_rail_poll() failed: %d\n", err);
-		return err;
+		goto error;
 	}
 
 	tegra_io_rail_unprepare();
 
-	return 0;
+error:
+	mutex_unlock(&pmc->powergates_lock);
+
+	return err;
 }
 EXPORT_SYMBOL(tegra_io_rail_power_on);
 
-int tegra_io_rail_power_off(int id)
+int tegra_io_rail_power_off(unsigned int id)
 {
 	unsigned long request, status, value;
 	unsigned int bit, mask;
 	int err;
 
+	mutex_lock(&pmc->powergates_lock);
+
 	err = tegra_io_rail_prepare(id, &request, &status, &bit);
-	if (err < 0) {
+	if (err) {
 		pr_info("tegra_io_rail_prepare() failed: %d\n", err);
-		return err;
+		goto error;
 	}
 
 	mask = 1 << bit;
@@ -563,12 +956,15 @@ int tegra_io_rail_power_off(int id)
 	tegra_pmc_writel(value, request);
 
 	err = tegra_io_rail_poll(status, mask, mask, 250);
-	if (err < 0)
-		return err;
+	if (err)
+		goto error;
 
 	tegra_io_rail_unprepare();
 
-	return 0;
+error:
+	mutex_unlock(&pmc->powergates_lock);
+
+	return err;
 }
 EXPORT_SYMBOL(tegra_io_rail_power_off);
 
@@ -727,7 +1123,7 @@ static void tegra_pmc_init(struct tegra_pmc *pmc)
 	tegra_pmc_writel(value, PMC_CNTRL);
 }
 
-void tegra_pmc_init_tsense_reset(struct tegra_pmc *pmc)
+static void tegra_pmc_init_tsense_reset(struct tegra_pmc *pmc)
 {
 	static const char disabled[] = "emergency thermal reset disabled";
 	u32 pmu_addr, ctrl_id, reg_addr, reg_data, pinmux;
@@ -805,7 +1201,7 @@ out:
 
 static int tegra_pmc_probe(struct platform_device *pdev)
 {
-	void __iomem *base = pmc->base;
+	void __iomem *base;
 	struct resource *res;
 	int err;
 
@@ -815,11 +1211,9 @@ static int tegra_pmc_probe(struct platform_device *pdev)
 
 	/* take over the memory region from the early initialization */
 	res = platform_get_resource(pdev, IORESOURCE_MEM, 0);
-	pmc->base = devm_ioremap_resource(&pdev->dev, res);
-	if (IS_ERR(pmc->base))
-		return PTR_ERR(pmc->base);
-
-	iounmap(base);
+	base = devm_ioremap_resource(&pdev->dev, res);
+	if (IS_ERR(base))
+		return PTR_ERR(base);
 
 	pmc->clk = devm_clk_get(&pdev->dev, "pclk");
 	if (IS_ERR(pmc->clk)) {
@@ -842,11 +1236,19 @@ static int tegra_pmc_probe(struct platform_device *pdev)
 
 	err = register_restart_handler(&tegra_pmc_restart_handler);
 	if (err) {
+		debugfs_remove(pmc->debugfs);
 		dev_err(&pdev->dev, "unable to register restart handler, %d\n",
 			err);
 		return err;
 	}
 
+	tegra_powergate_init(pmc);
+
+	mutex_lock(&pmc->powergates_lock);
+	iounmap(pmc->base);
+	pmc->base = base;
+	mutex_unlock(&pmc->powergates_lock);
+
 	return 0;
 }
 
@@ -964,7 +1366,6 @@ static const char * const tegra124_powergates[] = {
 	[TEGRA_POWERGATE_VENC] = "venc",
 	[TEGRA_POWERGATE_PCIE] = "pcie",
 	[TEGRA_POWERGATE_VDEC] = "vdec",
-	[TEGRA_POWERGATE_L2] = "l2",
 	[TEGRA_POWERGATE_MPE] = "mpe",
 	[TEGRA_POWERGATE_HEG] = "heg",
 	[TEGRA_POWERGATE_SATA] = "sata",
@@ -1006,17 +1407,13 @@ static const char * const tegra210_powergates[] = {
 	[TEGRA_POWERGATE_3D] = "3d",
 	[TEGRA_POWERGATE_VENC] = "venc",
 	[TEGRA_POWERGATE_PCIE] = "pcie",
-	[TEGRA_POWERGATE_L2] = "l2",
 	[TEGRA_POWERGATE_MPE] = "mpe",
-	[TEGRA_POWERGATE_HEG] = "heg",
 	[TEGRA_POWERGATE_SATA] = "sata",
 	[TEGRA_POWERGATE_CPU1] = "cpu1",
 	[TEGRA_POWERGATE_CPU2] = "cpu2",
 	[TEGRA_POWERGATE_CPU3] = "cpu3",
-	[TEGRA_POWERGATE_CELP] = "celp",
 	[TEGRA_POWERGATE_CPU0] = "cpu0",
 	[TEGRA_POWERGATE_C0NC] = "c0nc",
-	[TEGRA_POWERGATE_C1NC] = "c1nc",
 	[TEGRA_POWERGATE_SOR] = "sor",
 	[TEGRA_POWERGATE_DIS] = "dis",
 	[TEGRA_POWERGATE_DISB] = "disb",
@@ -1080,6 +1477,7 @@ static int __init tegra_pmc_early_init(void)
 	const struct of_device_id *match;
 	struct device_node *np;
 	struct resource regs;
+	unsigned int i;
 	bool invert;
 	u32 value;
 
@@ -1129,6 +1527,11 @@ static int __init tegra_pmc_early_init(void)
 		return -ENXIO;
 	}
 
+	/* Create a bit-map of the available and valid partitions */
+	for (i = 0; i < pmc->soc->num_powergates; i++)
+		if (pmc->soc->powergates[i])
+			set_bit(i, pmc->powergates_available);
+
 	mutex_init(&pmc->powergates_lock);
 
 	/*