#include <linux/pwm.h>
 #include <linux/slab.h>
 
+#include <dt-bindings/mfd/cros_ec.h>
+
 /**
  * struct cros_ec_pwm_device - Driver data for EC PWM
  *
  * @dev: Device node
  * @ec: Pointer to EC device
  * @chip: PWM controller chip
+ * @use_pwm_type: Use PWM types instead of generic channels
  */
 struct cros_ec_pwm_device {
        struct device *dev;
        struct cros_ec_device *ec;
        struct pwm_chip chip;
+       bool use_pwm_type;
 };
 
 /**
        kfree(channel);
 }
 
-static int cros_ec_pwm_set_duty(struct cros_ec_device *ec, u8 index, u16 duty)
+static int cros_ec_dt_type_to_pwm_type(u8 dt_index, u8 *pwm_type)
 {
+       switch (dt_index) {
+       case CROS_EC_PWM_DT_KB_LIGHT:
+               *pwm_type = EC_PWM_TYPE_KB_LIGHT;
+               return 0;
+       case CROS_EC_PWM_DT_DISPLAY_LIGHT:
+               *pwm_type = EC_PWM_TYPE_DISPLAY_LIGHT;
+               return 0;
+       default:
+               return -EINVAL;
+       }
+}
+
+static int cros_ec_pwm_set_duty(struct cros_ec_pwm_device *ec_pwm, u8 index,
+                               u16 duty)
+{
+       struct cros_ec_device *ec = ec_pwm->ec;
        struct {
                struct cros_ec_command msg;
                struct ec_params_pwm_set_duty params;
        } __packed buf;
        struct ec_params_pwm_set_duty *params = &buf.params;
        struct cros_ec_command *msg = &buf.msg;
+       int ret;
 
        memset(&buf, 0, sizeof(buf));
 
        msg->outsize = sizeof(*params);
 
        params->duty = duty;
-       params->pwm_type = EC_PWM_TYPE_GENERIC;
-       params->index = index;
+
+       if (ec_pwm->use_pwm_type) {
+               ret = cros_ec_dt_type_to_pwm_type(index, ¶ms->pwm_type);
+               if (ret) {
+                       dev_err(ec->dev, "Invalid PWM type index: %d\n", index);
+                       return ret;
+               }
+               params->index = 0;
+       } else {
+               params->pwm_type = EC_PWM_TYPE_GENERIC;
+               params->index = index;
+       }
 
        return cros_ec_cmd_xfer_status(ec, msg);
 }
 
-static int cros_ec_pwm_get_duty(struct cros_ec_device *ec, u8 index)
+static int cros_ec_pwm_get_duty(struct cros_ec_pwm_device *ec_pwm, u8 index)
 {
+       struct cros_ec_device *ec = ec_pwm->ec;
        struct {
                struct cros_ec_command msg;
                union {
        msg->insize = sizeof(*resp);
        msg->outsize = sizeof(*params);
 
-       params->pwm_type = EC_PWM_TYPE_GENERIC;
-       params->index = index;
+       if (ec_pwm->use_pwm_type) {
+               ret = cros_ec_dt_type_to_pwm_type(index, ¶ms->pwm_type);
+               if (ret) {
+                       dev_err(ec->dev, "Invalid PWM type index: %d\n", index);
+                       return ret;
+               }
+               params->index = 0;
+       } else {
+               params->pwm_type = EC_PWM_TYPE_GENERIC;
+               params->index = index;
+       }
 
        ret = cros_ec_cmd_xfer_status(ec, msg);
        if (ret < 0)
         */
        duty_cycle = state->enabled ? state->duty_cycle : 0;
 
-       ret = cros_ec_pwm_set_duty(ec_pwm->ec, pwm->hwpwm, duty_cycle);
+       ret = cros_ec_pwm_set_duty(ec_pwm, pwm->hwpwm, duty_cycle);
        if (ret < 0)
                return ret;
 
        struct cros_ec_pwm *channel = pwm_get_chip_data(pwm);
        int ret;
 
-       ret = cros_ec_pwm_get_duty(ec_pwm->ec, pwm->hwpwm);
+       ret = cros_ec_pwm_get_duty(ec_pwm, pwm->hwpwm);
        if (ret < 0) {
                dev_err(chip->dev, "error getting initial duty: %d\n", ret);
                return;
  * of PWMs it supports directly, so we have to read the pwm duty cycle for
  * subsequent channels until we get an error.
  */
-static int cros_ec_num_pwms(struct cros_ec_device *ec)
+static int cros_ec_num_pwms(struct cros_ec_pwm_device *ec_pwm)
 {
        int i, ret;
 
        /* The index field is only 8 bits */
        for (i = 0; i <= U8_MAX; i++) {
-               ret = cros_ec_pwm_get_duty(ec, i);
+               ret = cros_ec_pwm_get_duty(ec_pwm, i);
                /*
                 * We look for SUCCESS, INVALID_COMMAND, or INVALID_PARAM
                 * responses; everything else is treated as an error.
 {
        struct cros_ec_device *ec = dev_get_drvdata(pdev->dev.parent);
        struct device *dev = &pdev->dev;
+       struct device_node *np = pdev->dev.of_node;
        struct cros_ec_pwm_device *ec_pwm;
        struct pwm_chip *chip;
        int ret;
        chip = &ec_pwm->chip;
        ec_pwm->ec = ec;
 
+       if (of_device_is_compatible(np, "google,cros-ec-pwm-type"))
+               ec_pwm->use_pwm_type = true;
+
        /* PWM chip */
        chip->dev = dev;
        chip->ops = &cros_ec_pwm_ops;
        chip->of_xlate = cros_ec_pwm_xlate;
        chip->of_pwm_n_cells = 1;
-       ret = cros_ec_num_pwms(ec);
-       if (ret < 0) {
-               dev_err(dev, "Couldn't find PWMs: %d\n", ret);
-               return ret;
+
+       if (ec_pwm->use_pwm_type) {
+               chip->npwm = CROS_EC_PWM_DT_COUNT;
+       } else {
+               ret = cros_ec_num_pwms(ec_pwm);
+               if (ret < 0) {
+                       dev_err(dev, "Couldn't find PWMs: %d\n", ret);
+                       return ret;
+               }
+               chip->npwm = ret;
        }
-       chip->npwm = ret;
+
        dev_dbg(dev, "Probed %u PWMs\n", chip->npwm);
 
        ret = pwmchip_add(chip);
 #ifdef CONFIG_OF
 static const struct of_device_id cros_ec_pwm_of_match[] = {
        { .compatible = "google,cros-ec-pwm" },
+       { .compatible = "google,cros-ec-pwm-type" },
        {},
 };
 MODULE_DEVICE_TABLE(of, cros_ec_pwm_of_match);