diff --git a/arch/arm/boot/dts/qcom/msmcobalt.dtsi b/arch/arm/boot/dts/qcom/msmcobalt.dtsi
index e2f3d947cc06..2f2824bff4ea 100644
--- a/arch/arm/boot/dts/qcom/msmcobalt.dtsi
+++ b/arch/arm/boot/dts/qcom/msmcobalt.dtsi
@@ -704,6 +704,7 @@
 		reg = <0x100000 0xb0000>;
 		reg-names = "cc_base";
 		vdd_dig-supply = <&pmcobalt_s1_level>;
+		vdd_dig_ao-supply = <&pmcobalt_s1_level_ao>;
 		#clock-cells = <1>;
 	};
 
diff --git a/drivers/clk/msm/clock-gcc-cobalt.c b/drivers/clk/msm/clock-gcc-cobalt.c
index 1756157c996d..cbd8bc686caa 100644
--- a/drivers/clk/msm/clock-gcc-cobalt.c
+++ b/drivers/clk/msm/clock-gcc-cobalt.c
@@ -56,6 +56,7 @@ static void __iomem *virt_dbgbase;
 	}
 
 static DEFINE_VDD_REGULATORS(vdd_dig, VDD_DIG_NUM, 1, vdd_corner, NULL);
+static DEFINE_VDD_REGULATORS(vdd_dig_ao, VDD_DIG_NUM, 1, vdd_corner, NULL);
 
 DEFINE_CLK_RPM_SMD_BRANCH(cxo_clk_src, cxo_clk_src_ao, RPM_MISC_CLK_TYPE,
 			  CXO_CLK_SRC_ID, 19200000);
@@ -211,7 +212,7 @@ static struct rcg_clk hmss_ahb_clk_src = {
 	.c = {
 		.dbg_name = "hmss_ahb_clk_src",
 		.ops = &clk_ops_rcg,
-		VDD_DIG_FMAX_MAP3(LOWER, 19200000, LOW, 50000000,
+		VDD_DIG_FMAX_MAP3_AO(LOWER, 19200000, LOW, 50000000,
 					NOMINAL, 100000000),
 		CLK_INIT(hmss_ahb_clk_src.c),
 	},
@@ -1029,7 +1030,7 @@ static struct rcg_clk hmss_gpll0_clk_src = {
 	.c = {
 		.dbg_name = "hmss_gpll0_clk_src",
 		.ops = &clk_ops_rcg,
-		VDD_DIG_FMAX_MAP1(LOWER, 600000000),
+		VDD_DIG_FMAX_MAP1_AO(LOWER, 600000000),
 		CLK_INIT(hmss_gpll0_clk_src.c),
 	},
 };
@@ -2762,11 +2763,6 @@ static int msm_gcc_cobalt_probe(struct platform_device *pdev)
 		return -ENOMEM;
 	}
 
-	/* Set the HMSS_AHB_CLK_ENA bit to enable the hmss_ahb_clk */
-	regval = readl_relaxed(virt_base + GCC_APCS_CLOCK_BRANCH_ENA_VOTE);
-	regval |= BIT(21);
-	writel_relaxed(regval, virt_base + GCC_APCS_CLOCK_BRANCH_ENA_VOTE);
-
 	/*
 	 * Set the HMSS_AHB_CLK_SLEEP_ENA bit to allow the hmss_ahb_clk to be
 	 * turned off by hardware during certain apps low power modes.
@@ -2783,6 +2779,14 @@ static int msm_gcc_cobalt_probe(struct platform_device *pdev)
 		return PTR_ERR(vdd_dig.regulator[0]);
 	}
 
+	vdd_dig_ao.regulator[0] = devm_regulator_get(&pdev->dev, "vdd_dig_ao");
+	if (IS_ERR(vdd_dig_ao.regulator[0])) {
+		if (!(PTR_ERR(vdd_dig_ao.regulator[0]) == -EPROBE_DEFER))
+			dev_err(&pdev->dev,
+					"Unable to get vdd_dig_ao regulator\n");
+		return PTR_ERR(vdd_dig_ao.regulator[0]);
+	}
+
 	bimc_clk.c.parent = &cxo_clk_src.c;
 	ret = of_msm_clock_register(pdev->dev.of_node, msm_clocks_rpm_cobalt,
 				    ARRAY_SIZE(msm_clocks_rpm_cobalt));
diff --git a/drivers/clk/msm/vdd-level-cobalt.h b/drivers/clk/msm/vdd-level-cobalt.h
index 2cb40afafe3f..f847a4104d4d 100644
--- a/drivers/clk/msm/vdd-level-cobalt.h
+++ b/drivers/clk/msm/vdd-level-cobalt.h
@@ -50,11 +50,19 @@
 	},					\
 	.num_fmax = VDD_DIG_NUM
 
-#define VDD_DIG_FMAX_MAP2_AO(l1, f1, l2, f2) \
+#define VDD_DIG_FMAX_MAP1_AO(l1, f1)		 \
+	.vdd_class = &vdd_dig_ao,		\
+	.fmax = (unsigned long[VDD_DIG_NUM]) {	\
+		[VDD_DIG_##l1] = (f1),		\
+	},					\
+	.num_fmax = VDD_DIG_NUM
+
+#define VDD_DIG_FMAX_MAP3_AO(l1, f1, l2, f2, l3, f3) \
 	.vdd_class = &vdd_dig_ao,			\
 	.fmax = (unsigned long[VDD_DIG_NUM]) {	\
 		[VDD_DIG_##l1] = (f1),		\
 		[VDD_DIG_##l2] = (f2),		\
+		[VDD_DIG_##l3] = (f3),		\
 	},					\
 	.num_fmax = VDD_DIG_NUM