diff --git a/drivers/clk/qcom/krait-cc.c b/drivers/clk/qcom/krait-cc.c
index 410ae8390f1c..2ce38024dc0d 100644
--- a/drivers/clk/qcom/krait-cc.c
+++ b/drivers/clk/qcom/krait-cc.c
@@ -235,7 +235,7 @@ krait_add_pri_mux(struct device *dev, struct clk_hw *hfpll_div, struct clk_hw *s
 		.parent_data = p_data,
 		.num_parents = ARRAY_SIZE(p_data),
 		.ops = &krait_mux_clk_ops,
-		.flags = CLK_SET_RATE_PARENT,
+		.flags = CLK_SET_RATE_PARENT | CLK_IS_CRITICAL,
 	};
 	struct clk_hw *clk;
 	char *hfpll_name;
@@ -324,19 +324,6 @@ static struct clk_hw *krait_add_clks(struct device *dev, int id, bool unique_aux
 	return pri_mux;
 }
 
-static struct clk *krait_of_get(struct of_phandle_args *clkspec, void *data)
-{
-	unsigned int idx = clkspec->args[0];
-	struct clk **clks = data;
-
-	if (idx >= clks_max) {
-		pr_err("%s: invalid clock index %d\n", __func__, idx);
-		return ERR_PTR(-EINVAL);
-	}
-
-	return clks[idx] ? : ERR_PTR(-ENODEV);
-}
-
 static const struct of_device_id krait_cc_match_table[] = {
 	{ .compatible = "qcom,krait-cc-v1", (void *)1UL },
 	{ .compatible = "qcom,krait-cc-v2" },
@@ -344,60 +331,81 @@ static const struct of_device_id krait_cc_match_table[] = {
 };
 MODULE_DEVICE_TABLE(of, krait_cc_match_table);
 
+static int krait_clk_reinit(struct clk_hw *hw, int cpu)
+{
+	struct clk *clk;
+	unsigned long cur_rate, aux_rate;
+	char name[5]; /* CPUn */
+
+	if (cpu == -1)
+		strcpy(name, "L2");
+	else
+		snprintf(name, sizeof(name), "CPU%d", cpu);
+
+	clk = clk_hw_get_clk(hw, clk_hw_get_name(hw));
+	if (IS_ERR(clk))
+		return PTR_ERR(clk);
+
+	aux_rate = 384000000;
+
+	cur_rate = clk_get_rate(clk);
+	if (cur_rate < aux_rate) {
+		pr_info("%s @ Undefined rate %lu. Forcing new rate.\n",
+			name, cur_rate / 1000);
+		cur_rate = aux_rate;
+	}
+
+	clk_set_rate(clk, aux_rate);
+	clk_set_rate(clk, 2);
+	clk_set_rate(clk, cur_rate);
+	pr_info("%s @ %lu KHz\n", name, clk_get_rate(clk) / 1000);
+
+	clk_put(clk);
+
+	return 0;
+}
+
 static int krait_cc_probe(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
 	const struct of_device_id *id;
-	unsigned long cur_rate, aux_rate;
 	int cpu;
-	struct clk_hw *mux, *l2_pri_mux;
-	struct clk *clk, **clks;
+	struct clk_hw *clk;
+	struct clk_hw_onecell_data *clks;
 
 	id = of_match_device(krait_cc_match_table, dev);
 	if (!id)
 		return -ENODEV;
 
 	/* Rate is 1 because 0 causes problems for __clk_mux_determine_rate */
-	clk = clk_register_fixed_rate(dev, "qsb", NULL, 0, 1);
+	clk = clk_hw_register_fixed_rate(dev, "qsb", NULL, 0, 1);
 	if (IS_ERR(clk))
 		return PTR_ERR(clk);
 
 	if (!id->data) {
-		clk = clk_register_fixed_factor(dev, "acpu_aux",
-						"gpll0_vote", 0, 1, 2);
+		clk = clk_hw_register_fixed_factor(dev, "acpu_aux", "gpll0_vote", 0, 1, 2);
 		if (IS_ERR(clk))
 			return PTR_ERR(clk);
 	}
 
 	/* Krait configurations have at most 4 CPUs and one L2 */
-	clks = devm_kcalloc(dev, clks_max, sizeof(*clks), GFP_KERNEL);
+	clks = devm_kzalloc(dev, struct_size(clks, hws, clks_max), GFP_KERNEL);
 	if (!clks)
 		return -ENOMEM;
 
+	clks->num = clks_max;
+
 	for_each_possible_cpu(cpu) {
-		mux = krait_add_clks(dev, cpu, id->data);
-		if (IS_ERR(mux))
-			return PTR_ERR(mux);
-		clks[cpu] = mux->clk;
+		clk = krait_add_clks(dev, cpu, id->data);
+		if (IS_ERR(clk))
+			return PTR_ERR(clk);
+		clks->hws[cpu] = clk;
 	}
 
-	l2_pri_mux = krait_add_clks(dev, -1, id->data);
-	if (IS_ERR(l2_pri_mux))
-		return PTR_ERR(l2_pri_mux);
-	clks[l2_mux] = l2_pri_mux->clk;
-
-	/*
-	 * We don't want the CPU or L2 clocks to be turned off at late init
-	 * if CPUFREQ or HOTPLUG configs are disabled. So, bump up the
-	 * refcount of these clocks. Any cpufreq/hotplug manager can assume
-	 * that the clocks have already been prepared and enabled by the time
-	 * they take over.
-	 */
-	for_each_online_cpu(cpu) {
-		clk_prepare_enable(clks[l2_mux]);
-		WARN(clk_prepare_enable(clks[cpu]),
-		     "Unable to turn on CPU%d clock", cpu);
-	}
+	clk = krait_add_clks(dev, -1, id->data);
+	if (IS_ERR(clk))
+		return PTR_ERR(clk);
+	clks->hws[l2_mux] = clk;
 
 	/*
 	 * Force reinit of HFPLLs and muxes to overwrite any potential
@@ -410,33 +418,11 @@ static int krait_cc_probe(struct platform_device *pdev)
 	 * two different rates to force a HFPLL reinit under all
 	 * circumstances.
 	 */
-	cur_rate = clk_get_rate(clks[l2_mux]);
-	aux_rate = 384000000;
-	if (cur_rate < aux_rate) {
-		pr_info("L2 @ Undefined rate. Forcing new rate.\n");
-		cur_rate = aux_rate;
-	}
-	clk_set_rate(clks[l2_mux], aux_rate);
-	clk_set_rate(clks[l2_mux], 2);
-	clk_set_rate(clks[l2_mux], cur_rate);
-	pr_info("L2 @ %lu KHz\n", clk_get_rate(clks[l2_mux]) / 1000);
-	for_each_possible_cpu(cpu) {
-		clk = clks[cpu];
-		cur_rate = clk_get_rate(clk);
-		if (cur_rate < aux_rate) {
-			pr_info("CPU%d @ Undefined rate. Forcing new rate.\n", cpu);
-			cur_rate = aux_rate;
-		}
-
-		clk_set_rate(clk, aux_rate);
-		clk_set_rate(clk, 2);
-		clk_set_rate(clk, cur_rate);
-		pr_info("CPU%d @ %lu KHz\n", cpu, clk_get_rate(clk) / 1000);
-	}
+	krait_clk_reinit(clks->hws[l2_mux], -1);
+	for_each_possible_cpu(cpu)
+		krait_clk_reinit(clks->hws[cpu], cpu);
 
-	of_clk_add_provider(dev->of_node, krait_of_get, clks);
-
-	return 0;
+	return devm_of_clk_add_hw_provider(dev, of_clk_hw_onecell_get, clks);
 }
 
 static struct platform_driver krait_cc_driver = {
