diff --git a/drivers/pwm/pwm-rockchip.c b/drivers/pwm/pwm-rockchip.c
index 86685fcc802a..7f02221c258b 100644
--- a/drivers/pwm/pwm-rockchip.c
+++ b/drivers/pwm/pwm-rockchip.c
@@ -278,7 +278,8 @@ struct rockchip_pwm_funcs {
 	void (*set_capture)(struct pwm_chip *chip, struct pwm_device *pwm, bool enable);
 	int (*get_capture_result)(struct pwm_chip *chip, struct pwm_device *pwm,
 				  struct pwm_capture *catpure_res);
-	int (*set_counter)(struct pwm_chip *chip, struct pwm_device *pwm, bool enable);
+	int (*set_counter)(struct pwm_chip *chip, struct pwm_device *pwm,
+			   enum rockchip_pwm_counter_input_sel input_sel, bool enable);
 	int (*get_counter_result)(struct pwm_chip *chip, struct pwm_device *pwm,
 				  unsigned long *counter_res, bool is_clear);
 	int (*set_freq_meter)(struct pwm_chip *chip, struct pwm_device *pwm,
@@ -960,12 +961,14 @@ err_disable_pclk:
 }
 
 static int rockchip_pwm_set_counter_v4(struct pwm_chip *chip, struct pwm_device *pwm,
+				       enum rockchip_pwm_counter_input_sel input_sel,
 				       bool enable)
 {
 	struct rockchip_pwm_chip *pc = to_rockchip_pwm_chip(chip);
 	u32 arbiter = 0;
 	u32 channel_sel = 0;
 	u32 val;
+	int ret;
 
 	if (enable) {
 		arbiter = BIT(pc->channel_id) << COUNTER_READ_LOCK_SHIFT |
@@ -980,13 +983,25 @@ static int rockchip_pwm_set_counter_v4(struct pwm_chip *chip, struct pwm_device
 			return -EINVAL;
 	}
 
-	writel_relaxed(COUNTER_EN(enable) | COUNTER_CHANNEL_SEL(channel_sel),
+	if (enable) {
+		ret = clk_enable(pc->clk);
+		if (ret)
+			return ret;
+	}
+
+	writel_relaxed(COUNTER_EN(enable) | COUNTER_CLK_SEL(input_sel) |
+		       COUNTER_CHANNEL_SEL(channel_sel),
 		       pc->base + COUNTER_CTRL);
 
+	if (!enable)
+		clk_disable(pc->clk);
+
 	return 0;
 }
 
-int rockchip_pwm_set_counter(struct pwm_device *pwm, bool enable)
+int rockchip_pwm_set_counter(struct pwm_device *pwm,
+			     enum rockchip_pwm_counter_input_sel input_sel,
+			     bool enable)
 {
 	struct pwm_chip *chip;
 	struct rockchip_pwm_chip *pc;
@@ -1016,7 +1031,13 @@ int rockchip_pwm_set_counter(struct pwm_device *pwm, bool enable)
 	if (ret)
 		return ret;
 
-	ret = pc->data->funcs.set_counter(chip, pwm, enable);
+	ret = pinctrl_select_state(pc->pinctrl, pc->active_state);
+	if (ret) {
+		dev_err(chip->dev, "Failed to select pinctrl state\n");
+		goto err_disable_pclk;
+	}
+
+	ret = pc->data->funcs.set_counter(chip, pwm, input_sel, enable);
 	if (ret) {
 		dev_err(chip->dev, "Failed to abtain counter arbitration for PWM%d\n",
 			pc->channel_id);
diff --git a/include/linux/pwm-rockchip.h b/include/linux/pwm-rockchip.h
index c9462a7d50f4..b70099aa19cb 100644
--- a/include/linux/pwm-rockchip.h
+++ b/include/linux/pwm-rockchip.h
@@ -28,6 +28,16 @@ enum rockchip_pwm_global_ctrl_cmd {
 	PWM_GLOBAL_CTRL_DISABLE,
 };
 
+/**
+ * enum rockchip_pwm_counter_input_sel - select the input src of counter
+ * @PWM_COUNTER_INPUT_FROM_IO: input from PWM IO
+ * @PWM_COUNTER_INPUT_FROM_CRU: input from CRU
+ */
+enum rockchip_pwm_counter_input_sel {
+	PWM_COUNTER_INPUT_FROM_IO,
+	PWM_COUNTER_INPUT_FROM_CRU,
+};
+
 /**
  * enum rockchip_pwm_freq_meter_input_sel - select the input src of frequency meter
  * @PWM_FREQ_METER_INPUT_FROM_IO: input from PWM IO
@@ -122,7 +132,9 @@ struct rockchip_pwm_wave_config {
  *
  * Returns: 0 on success or a negative error code on failure.
  */
-int rockchip_pwm_set_counter(struct pwm_device *pwm, bool enable);
+int rockchip_pwm_set_counter(struct pwm_device *pwm,
+			     enum rockchip_pwm_counter_input_sel input_sel,
+			     bool enable);
 
 /**
  * rockchip_pwm_get_counter_result() - get counter result
@@ -166,7 +178,9 @@ int rockchip_pwm_global_ctrl(struct pwm_device *pwm, enum rockchip_pwm_global_ct
  */
 int rockchip_pwm_set_wave(struct pwm_device *pwm, struct rockchip_pwm_wave_config *config);
 #else
-static inline int rockchip_pwm_set_counter(struct pwm_device *pwm, bool enable)
+static inline int rockchip_pwm_set_counter(struct pwm_device *pwm,
+					   enum rockchip_pwm_counter_input_sel input_sel,
+					   bool enable);
 {
 	return 0;
 }